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一下看：
如下图所示，执行玩加载数据的代码之后，已经成功获取到数据

来源：https://blog.csdn.net/hshudoudou/article/details/127431107