Test dataset

prismer/experts/depth/generate_dataset.py

@@ -17,7 +17,9 @@ class Dataset(Dataset):
     def __init__(self, config, transform):
         self.data_path = config['data_path']
         self.transform = transform
-        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
+        # self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
+        data_folders = glob.glob(f'{self.data_path}/*/')
+        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.jpg')]
 
     def __len__(self):
         return len(self.data_list)

prismer_model.py

@@ -78,7 +78,6 @@ def run_experts(image_path: str) -> Tuple[str, Tuple[str, ...]]:
 
     config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
     config['im_name'] = im_name
-    print(im_name)
     with open('prismer/configs/experts.yaml', 'w') as yaml_file:
         yaml.dump(config, yaml_file, default_flow_style=False)