|
from data.base_dataset import BaseDataset, get_transform |
|
from data.image_folder import make_dataset |
|
from PIL import Image |
|
|
|
|
|
class SingleDataset(BaseDataset): |
|
"""This dataset class can load a set of images specified by the path --dataroot /path/to/data. |
|
|
|
It can be used for generating CycleGAN results only for one side with the model option '-model test'. |
|
""" |
|
|
|
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.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size)) |
|
input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc |
|
self.transform = get_transform(opt, grayscale=(input_nc == 1)) |
|
|
|
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 and A_paths |
|
A(tensor) - - an image in one domain |
|
A_paths(str) - - the path of the image |
|
""" |
|
A_path = self.A_paths[index] |
|
A_img = Image.open(A_path).convert('RGB') |
|
A = self.transform(A_img) |
|
return {'A': A, 'A_paths': A_path} |
|
|
|
def __len__(self): |
|
"""Return the total number of images in the dataset.""" |
|
return len(self.A_paths) |
|
|