概述
最近在跑一篇图像修复论文的代码，配置好环境之后开始运行，发现数据一直加载不进去。
害，还是得看人家代码咋写的，一句一句看逻辑，准能找出问题。通读dataset后，发现了问题所在，终于成功加载了数据集。
项目结构与代码
项目结构
主要的目的就是从数据集中读取到彩色图像和掩码图像。
代码
代码中涉及到torch.transforms、合并路径等知识点，我在代码中都进行了详细的注释，路径要对照着项目结构，如果自己用的话要根据项目结构去将相对路径改过来。
dataset.py ：当前的工作路径：…\OT-GAN-for-Inpainting-master\src\data
import os
import math
import numpy as np
from glob import glob
from random import shuffle
from PIL import Image, ImageFilter
import torch
import torchvision.transforms.functional as F
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
class InpaintingData(Dataset):
def __init__(self, args):
super(Dataset, self).__init__() # 继承Dataset的父类的初始化函数
self.w = self.h = args.image_size # 通过args传入新的属性---图像的w和h
self.mask_type = args.mask_type # 通过args传入新的属性---mask_type
# image and mask
self.image_path = [] #创建image_path的数组
for ext in ['*.jpg', '*.png']: # 获取每一个后缀为.jpg或者.png的图片，为ext
# 将dir_image、data_train和ext拼接作为图片的路径，并将其存入到数组image_path之中,glob()获取一个lsit集合
self.image_path.extend(glob(os.path.join(args.dir_image, args.data_train, ext)))
self.mask_path = glob(os.path.join(args.dir_mask, args.mask_type, '*.png')) #拼接dir_mask、mask_type和路径下所有的.png作为mask_path
# augmentation
self.img_trans = transforms.Compose([ #接收一个 transforms方法的list为参数，将这些操作组合到一起，返回一个新的tranforms
transforms.RandomResizedCrop(args.image_size), #随机随机长宽比裁剪，大小为image_size
transforms.RandomHorizontalFlip(), #随机水平翻转
transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), #改变图像的亮度、对比度、饱和度和色调。
transforms.ToTensor()]) # 转为tensor，并归一化至[0-1]
self.mask_trans = transforms.Compose([
transforms.Resize(args.image_size, interpolation=transforms.InterpolationMode.NEAREST), #将输入图像调整为给定的大小，interpolation是插值方式，此处是默认值NEAREST
transforms.RandomHorizontalFlip(), #随机水平翻转
transforms.RandomRotation( #随机旋转
(0, 45), interpolation=transforms.InterpolationMode.NEAREST), #(0, 45)是角度
])
def __len__(self): # __len__和__getitem__DataSet类必须实现的静态方法
return len(self.image_path)
def __getitem__(self, index):
# load image
image = Image.open(self.image_path[index]).convert('RGB') #获取图像，并将其转化为RGB（3x8位像素）模式
filename = os.path.basename(self.image_path[index]) #获取图片的路径
if self.mask_type == 'pconv': #如果mask_type为pconv
index = np.random.randint(0, len(self.mask_path)) #随机从mask_path中获取一个下标
mask = Image.open(self.mask_path[index]) #根据下标获取mask图片
mask = mask.convert('L') #将mask图片转化为L(8位像素的黑白图片，0表示黑，255表示白)模式
else: # 构造mask，有mask数据集的话就运行不到这里
mask = np.zeros((self.h, self.w)).astype(np.uint8) #构造与h和w一样大的图片，都用0填充，并将其转换为uint8
mask[self.h // 4:self.h // 4 * 3, self.w // 4:self.w // 4 * 3] = 1
mask = Image.fromarray(m).convert('L')
# augment
image = self.img_trans(image) * 2. - 1. # 数据标准化，将输出限定在一定的范围
mask = F.to_tensor(self.mask_trans(mask)) # 将转化后的mask图像转化为tensor
return image, mask, filename #返回
if __name__ == '__main__':
from attrdict import AttrDict
args = {
'dir_image': '../../examples/logos',
'data_train': 'image',
'dir_mask': '../../examples/logos/mask',
'mask_type': 'pconv',
'image_size': 512
}
args = AttrDict(args) # 将上面定义的参数传入AttrDict()作为新参数
data = InpaintingData(args) #创建InpaintingData对象
print(len(data), len(data.mask_path)) #输出data的长度，mask的长度
img, mask, filename = data[0] # 获取第一张图片
print(img.size(), mask.size(), filename) #打印上述信息
输出：
再Debug一下看：
如下图所示，执行玩加载数据的代码之后，已经成功获取到数据
总结
这段代码可以作为读取数据集的一个DataSet类的基础类，可以扩充进行修改，以后有类似需要可以拿过来修改。