template_dataset.py

"""Dataset class template
This module provides a template for users to implement custom datasets.
The filename should be <dataset>_dataset.py
The class name should be <Dataset>Dataset
You need to implement the following functions:
    <__init__>: Initialize this dataset class.
    <__getitem__>: Return a image and its corresponding lable.
    <__len__>: Return the size of the dataset.
"""

from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image

class TemplateDataset(Dataset):
    def __init__(self, opt):
    '''Initialize this dataset class.
       We need to specific the path of the dataset and the domain label of each image.
    '''
        self.image_list = []
        self.label_list = []
        if opt.is_train:
            trs = [transforms.Resize(opt.load_size, interpolation=Image.ANTIALIAS), transforms.RandomCrop(opt.fine_size)]
        else:
            trs = [transforms.Resize(opt.load_size, interpolation=Image.ANTIALIAS), transforms.CenterCrop(opt.fine_size)]
        if opt.is_flip:
            trs.append(transforms.RandomHorizontalFlip())
        trs.append(transforms.ToTensor())
        trs.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)))
        self.transform = transforms.Compose(trs)
        self.num_data = len(self.image_list)
        
    def __getitem__(self, index):
    '''Return a image and its corresponding lable.'''
        img = Image.open(self.image_list[index]).convert('RGB')
        img = self.transform(img)
        label = torch.FloatTensor(self.label_list[index])
        return img, label

    def __len__(self):
        '''Size of the dataset. '''
        return self.num_data