Fix bug

yezhen17 · Oct 16, 2021 · c01bc22 · c01bc22
1 parent 8af21a4
commit c01bc22
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/sunrgbd/sunrgbd_ssl_dataset.py b/sunrgbd/sunrgbd_ssl_dataset.py
@@ -183,7 +183,7 @@ def __getitem__(self, idx):
 
 class SunrgbdSSLUnlabeledDataset(Dataset):
     def __init__(self, labeled_sample_list=None, num_points=20000, use_color=False, use_height=False, use_v1=False,
-                 aug_num=1, scan_idx_list=None, load_labels=None):
+                 aug_num=1, scan_idx_list=None, load_labels=None, augment=True):
         print('----------------Sunrgbd Unlabeled Dataset Initialization----------------')
         if use_v1:
             self.data_path = os.path.join(ROOT_DIR, 'sunrgbd/sunrgbd_pc_bbox_votes_50k_v1_train')
@@ -211,6 +211,7 @@ def __init__(self, labeled_sample_list=None, num_points=20000, use_color=False,
         self.use_height = use_height
         self.aug_num = aug_num
         self.load_labels = load_labels
+        self.augment = augment
         if load_labels:
             print('Warning! Loading labels for analysis')
 
@@ -308,4 +309,4 @@ def __getitem__(self, idx):
         ret_dict['scale'] = np.array(scale_ratio).astype(np.float32)
         ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
         ret_dict['supervised_mask'] = np.array(0).astype(np.int64)
-        return ret_dict
+        return ret_dict
diff --git a/train.py b/train.py
@@ -122,7 +122,8 @@ def my_worker_init_fn(worker_id):
                                                    use_color=FLAGS.use_color,
                                                    use_height=(not FLAGS.no_height),
                                                    use_v1=(not FLAGS.use_sunrgbd_v2),
-                                                   load_labels=FLAGS.view_stats)
+                                                   load_labels=FLAGS.view_stats,
+                                                   augment=True)
     TEST_DATASET = SunrgbdDetectionVotesDataset('val',
                                                 num_points=NUM_POINT, augment=False,
                                                 use_color=FLAGS.use_color, use_height=(not FLAGS.no_height),
@@ -142,7 +143,8 @@ def my_worker_init_fn(worker_id):
                                                    num_points=NUM_POINT,
                                                    use_color=FLAGS.use_color,
                                                    use_height=(not FLAGS.no_height),
-                                                   load_labels=FLAGS.view_stats)
+                                                   load_labels=FLAGS.view_stats,
+                                                   augment=True)
     TEST_DATASET = ScannetDetectionDataset('val',
                                             num_points=NUM_POINT, augment=False,
                                             use_color=FLAGS.use_color, use_height=(not FLAGS.no_height))