MLBench/ReaLens

1import os2from data.base_dataset import BaseDataset, get_params, get_transform3from data.image_folder import make_dataset4from PIL import Image5 6 7class AlignedDataset(BaseDataset):8    """A dataset class for paired image dataset.9 10    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.11    During test time, you need to prepare a directory '/path/to/data/test'.12    """13 14    def __init__(self, opt):15        """Initialize this dataset class.16 17        Parameters:18            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions19        """20        BaseDataset.__init__(self, opt)21        self.dir_AB = os.path.join(opt.dataroot, opt.phase)  # get the image directory22        self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size))  # get image paths23        assert self.opt.load_size >= self.opt.crop_size  # crop_size should be smaller than the size of loaded image24        self.input_nc = self.opt.output_nc if self.opt.direction == "BtoA" else self.opt.input_nc25        self.output_nc = self.opt.input_nc if self.opt.direction == "BtoA" else self.opt.output_nc26 27    def __getitem__(self, index):28        """Return a data point and its metadata information.29 30        Parameters:31            index - - a random integer for data indexing32 33        Returns a dictionary that contains A, B, A_paths and B_paths34            A (tensor) - - an image in the input domain35            B (tensor) - - its corresponding image in the target domain36            A_paths (str) - - image paths37            B_paths (str) - - image paths (same as A_paths)38        """39        # read a image given a random integer index40        AB_path = self.AB_paths[index]41        AB = Image.open(AB_path).convert("RGB")42        # split AB image into A and B43        w, h = AB.size44        w2 = int(w / 2)45        A = AB.crop((0, 0, w2, h))46        B = AB.crop((w2, 0, w, h))47 48        # apply the same transform to both A and B49        transform_params = get_params(self.opt, A.size)50        A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))51        B_transform = get_transform(self.opt, transform_params, grayscale=(self.output_nc == 1))52 53        A = A_transform(A)54        B = B_transform(B)55 56        return {"A": A, "B": B, "A_paths": AB_path, "B_paths": AB_path}57 58    def __len__(self):59        """Return the total number of images in the dataset."""60        return len(self.AB_paths)61