|
- """
- Copyright (C) 2019 NVIDIA Corporation. All rights reserved.
- Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
- """
-
- from data.pix2pix_dataset import Pix2pixDataset
- from data.image_folder import make_dataset
- import json
-
- class CustomDataset(Pix2pixDataset):
- """ Dataset that loads images from directories
- Use option --label_dir, --image_dir, --instance_dir to specify the directories.
- The images in the directories are sorted in alphabetical order and paired in order.
- """
-
- @staticmethod
- def modify_commandline_options(parser, is_train):
- parser = Pix2pixDataset.modify_commandline_options(parser, is_train)
- # parser.set_defaults(preprocess_mode='resize_and_crop')
- # load_size = 286 if is_train else 256
- # parser.set_defaults(load_size=load_size)
- # parser.set_defaults(crop_size=256)
- # parser.set_defaults(display_winsize=256)
- parser.set_defaults(label_nc=29)
- parser.set_defaults(contain_dontcare_label=False)
-
- parser.add_argument('--label_dir', type=str, required=True,
- help='path to the directory that contains label images')
- parser.add_argument('--image_dir', type=str, default='',
- help='path to the directory that contains photo images')
- parser.add_argument('--instance_dir', type=str, default='',
- help='path to the directory that contains instance maps. Leave black if not exists')
- return parser
-
- def get_paths(self, opt):
- label_dir = opt.label_dir
- label_paths = make_dataset(label_dir, recursive=False, read_cache=True)
- print(f"label_dir: {label_dir}, image_dir: {opt.image_dir}, instance_dir: {opt.instance_dir}")
- # exit(0)
- if len(opt.image_dir) > 0:
- self.image_dir = opt.image_dir
- image_paths = make_dataset(
- self.image_dir, recursive=False, read_cache=True)
- else:
- image_paths = []
-
- if len(opt.instance_dir) > 0:
- instance_dir = opt.instance_dir
- instance_paths = make_dataset(
- instance_dir, recursive=False, read_cache=True)
- else:
- instance_paths = []
-
- if opt.isTrain:
- self.isTrain = True
- assert len(label_paths) == len(
- image_paths), "The #images in %s and %s do not match. Is there something wrong?"
- else:
- self.isTrain = False
- self.ref_dict = {}
- with open('./label_to_img.json', 'r') as f:
- self.ref_dict = json.load(f)
- # for line in f:
- # a = line.strip().split(',')
- # self.ref_dict[a[0]] = a[1]
- self.name = "CustomDataset"
-
- return label_paths, image_paths, instance_paths
|
