天天看点

Pytorch:保存图片1. 直接保存TensorTensor 转CV2 保存

1. 直接保存Tensor

#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import torch
from torchvision import utils as vutils
 
 
def save_image_tensor(input_tensor: torch.Tensor, filename):
    """
    将tensor保存为图片
    :param input_tensor: 要保存的tensor
    :param filename: 保存的文件名
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 复制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反归一化
    # input_tensor = unnormalize(input_tensor)
    vutils.save_image(input_tensor, filename)
           

Tensor 转CV2 保存

继续阅读