69 lines
2.7 KiB
Python
69 lines
2.7 KiB
Python
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')
|
|
"""
|
|
@staticmethod
|
|
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)
|