recursive dataset stuff

1 files changed, 42 insertions, 0 deletions

diff --git a/data/sequence_dataset.py b/data/sequence_dataset.py
new file mode 100644
index 0000000..3eaa12b
--- /dev/null
+++ b/data/sequence_dataset.py
@@ -0,0 +1,42 @@
+### 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 SequenceDataset(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 = self.A_paths[index]              
+        A = Image.open(A_path)        
+        params = get_params(self.opt, A.size)
+        if self.opt.label_nc == 0:
+            transform_A = get_transform(self.opt, params)
+            A_tensor = transform_A(A.convert('RGB'))
+        else:
+            transform_A = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
+            A_tensor = transform_A(A) * 255.0
+
+        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 'SequenceDataset'
\ No newline at end of file