1.数据处理代码utils.py:
1）
# coding:utf-8 import os import torch.nn as nn import numpy as np import scipy.misc import imageio import matplotlib.pyplot as plt import torch def tensor2im(input_image, imtype=np.uint8): """"将tensor的数据类型转成numpy类型，并反归一化. Parameters: input_image (tensor) -- 输入的图像tensor数组 imtype (type) -- 转换后的numpy的数据类型 """ mean = [0.485,0.456,0.406] #dataLoader中设置的mean参数 std = [0.229,0.224,0.225] #dataLoader中设置的std参数 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) def save_img(im, path, size): """im可是没经过任何处理的tensor类型的数据,将数据存储到path中 Parameters: im (tensor) -- 输入的图像tensor数组 path (str) -- 图像寻出的路径 size (list/tuple) -- 图像合并的高宽(heigth, width) """ scipy.misc.imsave(path, merge(im, size)) #将合并后的图保存到相应path中 def merge(images, size): """ 将batch size张图像合成一张大图，一行有size张图 :param images: 输入的图像tensor数组,shape = (batch_size, channels, height, width) :param size: 合并的高宽(heigth, width) :return: 合并后的图 """ h, w = images[0].shape[1], images[0].shape[1] if (images[0].shape[0] in (3,4)): # 彩色图像 c = images[0].shape[0] img = np.zeros((h * size[0], w * size[1], c)) for idx, image in enumerate(images): i = idx % size[1] j = idx // size[1] image = tensor2im(image) img[j * h:j * h + h, i * w:i * w + w, :] = image return img elif images.shape[3]==1: # 灰度图像 img = np.zeros((h * size[0], w * size[1])) for idx, image in enumerate(images): i = idx % size[1] j = idx // size[1] image = tensor2im(image) img[j * h:j * h + h, i * w:i * w + w] = image[:,:,0] return img else: raise ValueError('in merge(images,size) images parameter ''must have dimensions: HxW or HxWx3 or HxWx4')
2）
后面发现torchvision.utils有一个make_grid()函数能够直接实现将（batchsize，channels，height，width）格式的tensor图像数据合并成一张图。
同时其也有一个save_img(tensor, file_path)的方法，如果你的归一化的均值和方差都设置为0.5，那么你可以很简单地使用这个方法保存图片
但是因为我这里的均值和方差是自定义的，所以要自己写一个。所以上面的代码的merge()函数就可以不用了，可以简化为：
# coding:utf-8 import os, torchvision import torch.nn as nn import numpy as np import imageio import matplotlib.pyplot as plt from PIL import Image import torch def tensor2im(input_image, imtype=np.uint8): """"将tensor的数据类型转成numpy类型，并反归一化. Parameters: input_image (tensor) -- 输入的图像tensor数组 imtype (type) -- 转换后的numpy的数据类型 """ mean = [0.485,0.456,0.406] #自己设置的 std = [0.229,0.224,0.225] #自己设置的 if not isinstance(input_image, np.ndarray): if isinstance(input_image, torch.Tensor): # get the data from a variable 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 image_numpy = np.transpose(image_numpy, (1, 2, 0)) # post-processing: tranpose and scaling else: # if it is a numpy array, do nothing image_numpy = input_image return image_numpy.astype(imtype) def save_img(im, path, size): """im可是没经过任何处理的tensor类型的数据,将数据存储到path中 Parameters: im (tensor) -- 输入的图像tensor数组 path (str) -- 图像保存的路径 size (int) -- 一行有size张图,最好是2的倍数 """ im_grid = torchvision.utils.make_grid(im, size) #将batchsize的图合成一张图 im_numpy = tensor2im(im_grid) #转成numpy类型并反归一化 im_array = Image.fromarray(im_numpy) im_array.save(path)
2.数据读取代码dataLoader.py为：
# coding:utf-8 from torch.utils.data import DataLoader import utils import torch.utils.data as data from PIL import Image import os import torchvision.transforms as transforms import torch class ListDataset(data.Dataset): """处理数据，返回图片数据和数据类型""" def __init__(self, root, transform, type): self.type_list = [] self.imgsList = [] self.transform = transform self.imgs = os.listdir(root) for img in self.imgs: #得到所有数据的路径 self.imgsList.append(os.path.join(root, img)) self.type_list.append(int(type)) def __getitem__(self, idx): img_path = self.imgsList[idx] img = Image.open(img_path) img = self.transform(img) type_pred = self.type_list[idx] return img, type_pred def __len__(self): return len(self.imgs) def getTransform(input_size): transform = transforms.Compose([ transforms.Resize((input_size, input_size)),#重置大小 transforms.ToTensor(), #转为[0,1]值 transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225)) #标准化处理（mean, std） ]) return transform def dataloader0(input_size, batch_size, type): transform = getTransform(input_size) dataset = ListDataset(root='./GAN/data/0', transform=transform, type=type) loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8) return loader if __name__ == '__main__': batch_size = 4 dataloader0 = dataloader0(input_size=224, batch_size=batch_size, type=1) fix_images, _ = next(iter(dataloader0)) utils.save_img(fix_images, './real.png', (1, batch_size))
运行该代码，保存图像为：
使用简化后的utils.py代码,dataloader.py也要相应更改为：
if __name__ == '__main__': batch_size = 4 dataloader0 = dataloader0(input_size=256, batch_size=batch_size, type=1) fix_images, _ = next(iter(dataloader0)) utils.save_img(fix_images, './real.png', batch_size)
保存的图片为，效果相同：
