Spaces:
Runtime error
Runtime error
# Copyright (c) Microsoft Corporation. | |
# Licensed under the MIT License. | |
from data.base_dataset import BaseDataset, get_params, get_transform | |
from PIL import Image | |
import util.util as util | |
import os | |
class Pix2pixDataset(BaseDataset): | |
def modify_commandline_options(parser, is_train): | |
parser.add_argument( | |
"--no_pairing_check", | |
action="store_true", | |
help="If specified, skip sanity check of correct label-image file pairing", | |
) | |
return parser | |
def initialize(self, opt): | |
self.opt = opt | |
label_paths, image_paths, instance_paths = self.get_paths(opt) | |
util.natural_sort(label_paths) | |
util.natural_sort(image_paths) | |
if not opt.no_instance: | |
util.natural_sort(instance_paths) | |
label_paths = label_paths[: opt.max_dataset_size] | |
image_paths = image_paths[: opt.max_dataset_size] | |
instance_paths = instance_paths[: opt.max_dataset_size] | |
if not opt.no_pairing_check: | |
for path1, path2 in zip(label_paths, image_paths): | |
assert self.paths_match(path1, path2), ( | |
"The label-image pair (%s, %s) do not look like the right pair because the filenames are quite different. Are you sure about the pairing? Please see data/pix2pix_dataset.py to see what is going on, and use --no_pairing_check to bypass this." | |
% (path1, path2) | |
) | |
self.label_paths = label_paths | |
self.image_paths = image_paths | |
self.instance_paths = instance_paths | |
size = len(self.label_paths) | |
self.dataset_size = size | |
def get_paths(self, opt): | |
label_paths = [] | |
image_paths = [] | |
instance_paths = [] | |
assert False, "A subclass of Pix2pixDataset must override self.get_paths(self, opt)" | |
return label_paths, image_paths, instance_paths | |
def paths_match(self, path1, path2): | |
filename1_without_ext = os.path.splitext(os.path.basename(path1))[0] | |
filename2_without_ext = os.path.splitext(os.path.basename(path2))[0] | |
return filename1_without_ext == filename2_without_ext | |
def __getitem__(self, index): | |
# Label Image | |
label_path = self.label_paths[index] | |
label = Image.open(label_path) | |
params = get_params(self.opt, label.size) | |
transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) | |
label_tensor = transform_label(label) * 255.0 | |
label_tensor[label_tensor == 255] = self.opt.label_nc # 'unknown' is opt.label_nc | |
# input image (real images) | |
image_path = self.image_paths[index] | |
assert self.paths_match( | |
label_path, image_path | |
), "The label_path %s and image_path %s don't match." % (label_path, image_path) | |
image = Image.open(image_path) | |
image = image.convert("RGB") | |
transform_image = get_transform(self.opt, params) | |
image_tensor = transform_image(image) | |
# if using instance maps | |
if self.opt.no_instance: | |
instance_tensor = 0 | |
else: | |
instance_path = self.instance_paths[index] | |
instance = Image.open(instance_path) | |
if instance.mode == "L": | |
instance_tensor = transform_label(instance) * 255 | |
instance_tensor = instance_tensor.long() | |
else: | |
instance_tensor = transform_label(instance) | |
input_dict = { | |
"label": label_tensor, | |
"instance": instance_tensor, | |
"image": image_tensor, | |
"path": image_path, | |
} | |
# Give subclasses a chance to modify the final output | |
self.postprocess(input_dict) | |
return input_dict | |
def postprocess(self, input_dict): | |
return input_dict | |
def __len__(self): | |
return self.dataset_size | |