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)
1.2 图片转张量
img_tensor = transforms.ToTensor()(img)
print(img_tensor.size())
1.3 保存图片
方法一:torchvision.utils.save_image()
input_tensor_1 = img_tensor.clone().detach().to(torch.device('cpu'))
方法二: 转数组再保存成图片
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)
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)