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)]##全部图片的绝对路径