Pytorch如何把Tensor转化成图像可视化

2022-12-14 06:01

短信预约 -IT技能 免费直播动态提醒

Pytorch把Tensor转化成图像可视化

在调试程序的时候经常想把tensor可视化成来看看，可以这样操作：

from torchvision import transforms
unloader = transforms.ToPILImage()
image = original_tensor.cpu().clone()  # clone the tensor
image = image.squeeze(0)  # remove the fake batch dimension
image = unloader(image)
image.save('example.jpg')

pytorch标准化的Tensor转图像问题

常常在工作之中遇到将dataloader中出来的tensor成image，numpy格式的数据，然后可以可视化出来

但是这种tensor往往经过了channel变换(RGB2BGR)，以及归一化(减均值除方差)，

然后维度的顺序也发生变化(HWC变成CHW)。为了可视化这种变化比较多的数据，

在tensor转numpy之前需要对tensor做一些处理

如下是一个简单的函数，可以可视化tensor，下次直接拿来用就行

def tensor2im(input_image, imtype=np.uint8):
    """"
    Parameters:
        input_image (tensor) --  输入的tensor，维度为CHW，注意这里没有batch size的维度
        imtype (type)        --  转换后的numpy的数据类型
    """
    mean = [0.485, 0.456, 0.406] # dataLoader中设置的mean参数，需要从dataloader中拷贝过来
    std = [0.229, 0.224, 0.225]  # dataLoader中设置的std参数，需要从dataloader中拷贝过来
    if not isinstance(input_image, np.ndarray):
        if isinstance(input_image, torch.Tensor): # 如果传入的图片类型为torch.Tensor，则读取其数据进行下面的处理
            image_tensor = input_image.data
        else:
            return input_image
        image_numpy = image_tensor.cpu().float().numpy()  # convert it into a numpy array
        if image_numpy.shape[0] == 1:  # grayscale to RGB
            image_numpy = np.tile(image_numpy, (3, 1, 1))
        for i in range(len(mean)): # 反标准化，乘以方差，加上均值
            image_numpy[i] = image_numpy[i] * std[i] + mean[i]
        image_numpy = image_numpy * 255 #反ToTensor(),从[0,1]转为[0,255]
        image_numpy = np.transpose(image_numpy, (1, 2, 0))  # 从(channels, height, width)变为(height, width, channels)
    else:  # 如果传入的是numpy数组,则不做处理
        image_numpy = input_image
    return image_numpy.astype(imtype)