将torch tensor转化为图像

        反向操作transforms.Normalize和transforms.ToTensor函数，将tensor转换成.jpg、.png格式的图片

# -*- coding: utf-8 -*-

"""
Time:     2022.01.17
Author:   Athrunsunny
Version:  V 0.1
File:     tensor_convert_to_img.py
Describe: from torch tensor to image
"""
import torch
from PIL import Image
import torchvision.transforms as T


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, T.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])

    # C x H x W  ---> H x W x C
    img_tensor = img_tensor.transpose(0, 2).transpose(0, 1)

    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


def show_convert(image_path):
    img = Image.open(image_path)
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    imgToTensor = T.Compose([T.ToTensor(),T.Normalize(mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225])])
    img_tensor = imgToTensor(img)
    img = transform_convert(img_tensor, imgToTensor)
    img.save('./new.png')
    img.show()


if __name__ == '__main__':
    path = ''  # 填入图像路径
    show_convert(path)