张量保存成图片

1 三维张量转图片

1.1 读取图片

import torchvision
import torchvision.transforms as transforms
import cv2 as cv
import matplotlib.pyplot as plt
import os
import torch

# 读取图片
img = cv.imread('./1.jpg')
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
print(img.shape)        # (405, 405, 3)

1.2 图片转张量

# 图片转张量
img_tensor = transforms.ToTensor()(img)  # tensor数据格式是torch(C,H,W)
print(img_tensor.size())        # torch.Size([3, 405, 405])

1.3 保存图片

方法一：torchvision.utils.save_image（）

# 用torchvision.utils.save_image（）直接把张量保存成图片
input_tensor_1 = img_tensor.clone().detach().to(torch.device('cpu'))
# torchvision.utils.save_image(input_tensor_1, "out_x.jpg")

方法二： 转数组再保存成图片

input_tensor_2 = img_tensor.cpu().numpy().transpose(1, 2, 0)
path = os.path.join('./', '_test_' + '.png')
plt.imsave(path, input_tensor_2)

2 四维张量转图片（卷积神经网络的输出，多了一个batch_size维度）

for (x, _) in zip(data_loader):
    x, z = x.to(device), z.to(device)       # x.shape = torch.Size([1, 3, 256, 256])
    k = x.permute(0, 2, 3, 1)
    k = k.cpu().detach().numpy()

    path = os.path.join('./result/', str(epoch+1) + '_epoch' + '.jpg')
    plt.imsave(path, (k[0] + 1) / 2)