diff options
Diffstat (limited to 'data/recursive_dataset.py')
| -rw-r--r-- | data/recursive_dataset.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py new file mode 100644 index 0000000..14ce906 --- /dev/null +++ b/data/recursive_dataset.py @@ -0,0 +1,44 @@ +### Copyright (C) 2017 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). +import os.path +from data.base_dataset import BaseDataset, get_params, get_transform, normalize +from data.image_folder import make_dataset +from PIL import Image + +class RecursiveDataset(BaseDataset): + def initialize(self, opt): + self.opt = opt + self.root = opt.dataroot + + ### input A (label maps) + self.dir_A = opt.dataroot + self.A_paths = sorted(make_dataset(self.dir_A)) + + self.dataset_size = len(self.A_paths) + + def __getitem__(self, index): + ### input A (label maps) + A_path = os.path.join(self.opt.dataroot, "frame_{:05d}.png".format(index)) + if not os.path.exists(A_path): + # print() + while not os.path.exists(A_path): + # print('sleeping for {}'.format(self.opt.poll_delay)) + time.sleep(self.opt.poll_delay) + # print("got {}".format(A_path)) + A = Image.open(A_path) + params = get_params(self.opt, A.size) + transform_A = get_transform(self.opt, params) + A_tensor = transform_A(A.convert('RGB')) + + B_tensor = inst_tensor = feat_tensor = 0 + + input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, + 'feat': feat_tensor, 'path': A_path} + + return input_dict + + def __len__(self): + return len(self.A_paths) + + def name(self): + return 'RecursiveDataset'
\ No newline at end of file |