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)]##全部图片的绝对路径
self.tar_filenames = [os.path.join(dir1, 'target', x) for x in tar_files if is_image_file(x)]
self.ps = patchsize
def __len__(self):
return len(self.tar_filenames)
def __getitem__(self, index):
index_ = index % len(self.tar_filenames) ##这里是为了防止索引超出范围
ps = self.ps
inp_path = self.inp_filenames[index_]
tar_path = self.tar_filenames[index_]
inp_img = Image.open(inp_path)
tar_img = Image.open(tar_path)
w,h = tar_img.size
padw = ps-w if w
class DataLoaderVal(Dataset):
def __init__(self, rgb_dir, img_options=None, rgb_dir2=None):
super(DataLoaderVal, self).__init__()
inp_files = sorted(os.listdir(os.path.join(rgb_dir, 'input')))
tar_files = sorted(os.listdir(os.path.join(rgb_dir, 'target')))
self.inp_filenames = [os.path.join(rgb_dir, 'input', x) for x in inp_files if is_image_file(x)]
self.tar_filenames = [os.path.join(rgb_dir, 'target', x) for x in tar_files if is_image_file(x)]
self.img_options = img_options
self.sizex = len(self.tar_filenames) # get the size of target
self.ps = self.img_options['patch_size']
def __len__(self):
return self.sizex
def __getitem__(self, index):
index_ = index % self.sizex
ps = self.ps
inp_path = self.inp_filenames[index_]
tar_path = self.tar_filenames[index_]
inp_img = Image.open(inp_path)
tar_img = Image.open(tar_path)
# Validate on center crop
if self.ps is not None:
inp_img = TF.center_crop(inp_img, (ps,ps))
tar_img = TF.center_crop(tar_img, (ps,ps))
inp_img = TF.to_tensor(inp_img)
tar_img = TF.to_tensor(tar_img)
filename = os.path.splitext(os.path.split(tar_path)[-1])[0]
return tar_img, inp_img, filename
class DataLoaderTest(Dataset):
def __init__(self, inp_dir, img_options):
super(DataLoaderTest, self).__init__()
inp_files = sorted(os.listdir(inp_dir))
self.inp_filenames = [os.path.join(inp_dir, x) for x in inp_files if is_image_file(x)]
self.inp_size = len(self.inp_filenames)
self.img_options = img_options
def __len__(self):
return self.inp_size
def __getitem__(self, index):
path_inp = self.inp_filenames[index]
filename = os.path.splitext(os.path.split(path_inp)[-1])[0]
inp = Image.open(path_inp)
inp = TF.to_tensor(inp)
return inp, filename
