Spaces:
Running
Running
| import os | |
| from data.base_dataset import BaseDataset, get_transform | |
| from data.image_folder import make_dataset | |
| from skimage import color # require skimage | |
| from PIL import Image | |
| import numpy as np | |
| import torchvision.transforms as transforms | |
| class ColorizationDataset(BaseDataset): | |
| """This dataset class can load a set of natural images in RGB, and convert RGB format into (L, ab) pairs in Lab color space. | |
| This dataset is required by pix2pix-based colorization model ('--model colorization') | |
| """ | |
| def modify_commandline_options(parser, is_train): | |
| """Add new dataset-specific options, and rewrite default values for existing options. | |
| Parameters: | |
| parser -- original option parser | |
| is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options. | |
| Returns: | |
| the modified parser. | |
| By default, the number of channels for input image is 1 (L) and | |
| the number of channels for output image is 2 (ab). The direction is from A to B | |
| """ | |
| parser.set_defaults(input_nc=1, output_nc=2, direction='AtoB') | |
| return parser | |
| def __init__(self, opt): | |
| """Initialize this dataset class. | |
| Parameters: | |
| opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions | |
| """ | |
| BaseDataset.__init__(self, opt) | |
| self.dir = os.path.join(opt.dataroot, opt.phase) | |
| self.AB_paths = sorted(make_dataset(self.dir, opt.max_dataset_size)) | |
| assert(opt.input_nc == 1 and opt.output_nc == 2 and opt.direction == 'AtoB') | |
| self.transform = get_transform(self.opt, convert=False) | |
| def __getitem__(self, index): | |
| """Return a data point and its metadata information. | |
| Parameters: | |
| index - - a random integer for data indexing | |
| Returns a dictionary that contains A, B, A_paths and B_paths | |
| A (tensor) - - the L channel of an image | |
| B (tensor) - - the ab channels of the same image | |
| A_paths (str) - - image paths | |
| B_paths (str) - - image paths (same as A_paths) | |
| """ | |
| path = self.AB_paths[index] | |
| im = Image.open(path).convert('RGB') | |
| im = self.transform(im) | |
| im = np.array(im) | |
| lab = color.rgb2lab(im).astype(np.float32) | |
| lab_t = transforms.ToTensor()(lab) | |
| A = lab_t[[0], ...] / 50.0 - 1.0 | |
| B = lab_t[[1, 2], ...] / 110.0 | |
| return {'A': A, 'B': B, 'A_paths': path, 'B_paths': path} | |
| def __len__(self): | |
| """Return the total number of images in the dataset.""" | |
| return len(self.AB_paths) | |