这是一个工具包,功能:反向操作transforms.Normalize和transforms.ToTensor函数,将tensor格式的图片转换成.jpg、.png格式的图片

 (注:这里是我原始的写法,但是是存在着一些改进空间的【如评论区所言,会改变img_tensor的内容,但我这边就暂时保持原样了,详细的改进方法可见评论区】【感谢评论区大佬!】)

import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms

def transform_convert(img_tensor, transform):
    """
    param img_tensor: tensor
    param transforms: torchvision.transforms
    """
    if 'Normalize' in str(transform):
        normal_transform = list(filter(lambda x: isinstance(x, transforms.Normalize), transform.transforms))
        mean = torch.tensor(normal_transform[0].mean, dtype=img_tensor.dtype, device=img_tensor.device)
        std = torch.tensor(normal_transform[0].std, dtype=img_tensor.dtype, device=img_tensor.device)
        img_tensor.mul_(std[:,None,None]).add_(mean[:,None,None])
        
    img_tensor = img_tensor.transpose(0,2).transpose(0,1)  # C x H x W  ---> H x W x C
    
    if 'ToTensor' in str(transform) or img_tensor.max() < 1:
        img_tensor = img_tensor.detach().numpy()*255
    
    if isinstance(img_tensor, torch.Tensor):
    	img_tensor = img_tensor.numpy()
    
    if img_tensor.shape[2] == 3:
        img = Image.fromarray(img_tensor.astype('uint8')).convert('RGB')
    elif img_tensor.shape[2] == 1:
        img = Image.fromarray(img_tensor.astype('uint8')).squeeze()
    else:
        raise Exception("Invalid img shape, expected 1 or 3 in axis 2, but got {}!".format(img_tensor.shape[2]))
        
    return img

代码示例

import matplotlib.pyplot as plt
img = Image.open('./lena.jpg')
ToTensor_transform = transforms.Compose([transforms.ToTensor()])
img_tensor = ToTensor_transform(img)
img = transform_convert(img_tensor, ToTensor_transform)
plt.imshow(img)
plt.savefig('./lenaa.png')

在这里插入图片描述

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐