pytorch 自定义dataset

 dataset的定义主要包含三个部分。__init__()初始化，这部分主要是初始化一下数据集的地址和相关变量，比如说patchsize的大小。__len__()这部分主要是用来检测数据集的大小的。__getitem__()方法是是通过所以获取数据集中的数据。在这一部分也经常会出现数据增强的处理。

import os
import numpy as np
from torch.utils.data import Dataset
import torch
from PIL import Image
import torchvision.transforms.functional as TF

import random

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in ['jpeg', 'JPEG', 'jpg', 'png', 'JPG', 'PNG', 'gif'])

class DatasetTrain(Dataset):
    def __init__(self, dir1, patchsize=96):
        super(DatasetTrain, self).__init__()

        inp_files = sorted(os.listdir(os.path.join(dir1, 'input')))  ##获取全部图片的名称
        tar_files = sorted(os.listdir(os.path.join(dir1, 'target')))

        self.inp_filenames = [os.path.join(dir1, 'input', x)  for x in inp_files if is_image_file(x)]##全部图片的绝对路径