MLBench/ReaLens

1import os2from data.base_dataset import BaseDataset, get_transform3from data.image_folder import make_dataset4from skimage import color  # require skimage5from PIL import Image6import numpy as np7import torchvision.transforms as transforms8 9 10class ColorizationDataset(BaseDataset):11    """This dataset class can load a set of natural images in RGB, and convert RGB format into (L, ab) pairs in Lab color space.12 13    This dataset is required by pix2pix-based colorization model ('--model colorization')14    """15 16    @staticmethod17    def modify_commandline_options(parser, is_train):18        """Add new dataset-specific options, and rewrite default values for existing options.19 20        Parameters:21            parser          -- original option parser22            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.23 24        Returns:25            the modified parser.26 27        By default, the number of channels for input image  is 1 (L) and28        the number of channels for output image is 2 (ab). The direction is from A to B29        """30        parser.set_defaults(input_nc=1, output_nc=2, direction="AtoB")31        return parser32 33    def __init__(self, opt):34        """Initialize this dataset class.35 36        Parameters:37            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions38        """39        BaseDataset.__init__(self, opt)40        self.dir = os.path.join(opt.dataroot, opt.phase)41        self.AB_paths = sorted(make_dataset(self.dir, opt.max_dataset_size))42        assert opt.input_nc == 1 and opt.output_nc == 2 and opt.direction == "AtoB"43        self.transform = get_transform(self.opt, convert=False)44 45    def __getitem__(self, index):46        """Return a data point and its metadata information.47 48        Parameters:49            index - - a random integer for data indexing50 51        Returns a dictionary that contains A, B, A_paths and B_paths52            A (tensor) - - the L channel of an image53            B (tensor) - - the ab channels of the same image54            A_paths (str) - - image paths55            B_paths (str) - - image paths (same as A_paths)56        """57        path = self.AB_paths[index]58        im = Image.open(path).convert("RGB")59        im = self.transform(im)60        im = np.array(im)61        lab = color.rgb2lab(im).astype(np.float32)62        lab_t = transforms.ToTensor()(lab)63        A = lab_t[[0], ...] / 50.0 - 1.064        B = lab_t[[1, 2], ...] / 110.065        return {"A": A, "B": B, "A_paths": path, "B_paths": path}66 67    def __len__(self):68        """Return the total number of images in the dataset."""69        return len(self.AB_paths)70