pytorch dataset实战----读取数据集

---

概述

最近在跑一篇图像修复论文的代码，配置好环境之后开始运行，发现数据一直加载不进去。
害，还是得看人家代码咋写的，一句一句看逻辑，准能找出问题。通读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类的基础类，可以扩充进行修改，以后有类似需要可以拿过来修改。

参考资料

[1] https://github.com/researchmm/AOT-GAN-for-Inpainting