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

到此这篇关于pytorch dataset实战----读取数据集的文章就介绍到这了,更多相关pytorch 读取数据集内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多支持我们！