blob: 14ce9062726162a5b203f27f2df1d9f13e780d75 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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'
|
