change based on the fourth comment

tools/create_data.py

 import argparse
 import os.path as osp
 
+import tools.data_converter.indoor_converter as indoor
 import tools.data_converter.kitti_converter as kitti
 import tools.data_converter.nuscenes_converter as nuscenes_converter
-import tools.data_converter.scannet_converter as scannet
-import tools.data_converter.sunrgbd_converter as sunrgbd
 from tools.data_converter.create_gt_database import create_groundtruth_database
 
 
@@ -46,11 +45,11 @@ def nuscenes_data_prep(root_path,
 
 
 def scannet_data_prep(root_path, info_prefix, out_dir):
-    scannet.create_scannet_info_file(root_path, info_prefix, out_dir)
+    indoor.create_indoor_info_file(root_path, info_prefix, out_dir)
 
 
 def sunrgbd_data_prep(root_path, info_prefix, out_dir):
-    sunrgbd.create_sunrgbd_info_file(root_path, info_prefix, out_dir)
+    indoor.create_indoor_info_file(root_path, info_prefix, out_dir)
 
 
 parser = argparse.ArgumentParser(description='Data converter arg parser')

tools/data_converter/sunrgbd_converter.py → tools/data_converter/indoor_converter.py

 import os
 
 import mmcv
+from tools.data_converter.scannet_data_utils import ScanNetData
 from tools.data_converter.sunrgbd_data_utils import SUNRGBDData
 
 
-def create_sunrgbd_info_file(data_path,
-                             pkl_prefix='sunrgbd',
-                             save_path=None,
-                             use_v1=False):
-    '''
-    Create sunrgbd information file.
+def create_indoor_info_file(data_path,
+                            pkl_prefix='sunrgbd',
+                            save_path=None,
+                            use_v1=False):
+    """
+    Create indoor  information file.
 
     Get information of the raw data and save it to the pkl file.
 
     Args:
         data_path (str): Path of the data.
-        pkl_prefix (str): Prefix ofr the pkl to be saved. Default: 'sunrgbd'.
+        pkl_prefix (str): Prefix of the pkl to be saved. Default: 'sunrgbd'.
         save_path (str): Path of the pkl to be saved. Default: None.
         use_v1 (bool): Whether to use v1. Default: False.
 
     Returns:
         None
 
-    '''
+    """
     assert os.path.exists(data_path)
+    assert pkl_prefix in ['sunrgbd', 'scannet']
     if save_path is None:
         save_path = data_path
-    else:
-        save_path = save_path
     assert os.path.exists(save_path)
     train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl')
     val_filename = os.path.join(save_path, f'{pkl_prefix}_infos_val.pkl')
-    train_dataset = SUNRGBDData(
-        root_path=data_path, split='train', use_v1=use_v1)
-    val_dataset = SUNRGBDData(root_path=data_path, split='val', use_v1=use_v1)
-    sunrgbd_infos_train = train_dataset.get_sunrgbd_infos(has_label=True)
+    if pkl_prefix == 'sunrgbd':
+        train_dataset = SUNRGBDData(
+            root_path=data_path, split='train', use_v1=use_v1)
+        val_dataset = SUNRGBDData(
+            root_path=data_path, split='val', use_v1=use_v1)
+    else:
+        train_dataset = ScanNetData(root_path=data_path, split='train')
+        val_dataset = ScanNetData(root_path=data_path, split='val')
+    infos_train = train_dataset.get_infos(has_label=True)
     with open(train_filename, 'wb') as f:
-        mmcv.dump(sunrgbd_infos_train, f, 'pkl')
-    print(f'Sunrgbd info train file is saved to {train_filename}')
-    sunrgbd_infos_val = val_dataset.get_sunrgbd_infos(has_label=True)
+        mmcv.dump(infos_train, f, 'pkl')
+    print(f'{pkl_prefix} info train file is saved to {train_filename}')
+    infos_val = val_dataset.get_infos(has_label=True)
     with open(val_filename, 'wb') as f:
-        mmcv.dump(sunrgbd_infos_val, f, 'pkl')
-    print(f'Sunrgbd info val file is saved to {val_filename}')
+        mmcv.dump(infos_val, f, 'pkl')
+    print(f'{pkl_prefix} info val file is saved to {val_filename}')

tools/data_converter/scannet_converter.py

-import os
-
-import mmcv
-from tools.data_converter.scannet_data_utils import ScanNetData
-
-
-def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None):
-    '''
-        Create scannet information file.
-
-        Get information of the raw data and save it to the pkl file.
-
-        Args:
-            data_path (str): Path of the data.
-            pkl_prefix (str): Prefix ofr the pkl to be saved. Default: 'scannet'. # noqa: E501
-            save_path (str): Path of the pkl to be saved. Default: None.
-
-        Returns:
-            None
-
-    '''
-    assert os.path.exists(data_path)
-    if save_path is None:
-        save_path = data_path
-    else:
-        save_path = save_path
-    assert os.path.exists(save_path)
-    train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl')
-    val_filename = os.path.join(save_path, f'{pkl_prefix}_infos_val.pkl')
-    train_dataset = ScanNetData(root_path=data_path, split='train')
-    val_dataset = ScanNetData(root_path=data_path, split='val')
-    scannet_infos_train = train_dataset.get_scannet_infos(has_label=True)
-    with open(train_filename, 'wb') as f:
-        mmcv.dump(scannet_infos_train, f, 'pkl')
-    print(f'Scannet info train file is saved to {train_filename}')
-    scannet_infos_val = val_dataset.get_scannet_infos(has_label=True)
-    with open(val_filename, 'wb') as f:
-        mmcv.dump(scannet_infos_val, f, 'pkl')
-    print(f'Scannet info val file is saved to {val_filename}')

tools/data_converter/scannet_data_utils.py

@@ -5,7 +5,7 @@ import numpy as np
 
 
 class ScanNetData(object):
-    '''
+    """
     ScanNet Data
 
     Generate scannet infos for scannet_converter
@@ -13,7 +13,7 @@ class ScanNetData(object):
     Args:
         root_path (str): Root path of the raw data
         split (str): Set split type of the data. Default: 'train'.
-    '''
+    """
 
     def __init__(self, root_path, split='train'):
         self.root_dir = root_path
@@ -48,23 +48,21 @@ class ScanNetData(object):
         assert os.path.exists(box_file)
         return np.load(box_file)
 
-    def get_scannet_infos(self,
-                          num_workers=4,
-                          has_label=True,
-                          sample_id_list=None):
-        '''
-        Get scannet infos.
+    def get_infos(self, num_workers=4, has_label=True, sample_id_list=None):
+        """
+        Get data infos.
 
         This method gets information from the raw data.
 
         Args:
             num_workers (int): Number of threads to be used. Default: 4.
             has_label (bool): Whether the data has label. Default: True.
-            sample_id_list (List[int]): Index list of the sample. Default: None. # noqa: E501
+            sample_id_list (List[int]): Index list of the sample.
+                Default: None.
 
         Returns:
             infos (List[dict]): Information of the raw data.
-        '''
+        """
 
         def process_single_scene(sample_idx):
             print(f'{self.split} sample_idx: {sample_idx}')

tools/data_converter/sunrgbd_data_utils.py

 import concurrent.futures as futures
 import os
 
-import cv2
+import mmcv
 import numpy as np
 import scipy.io as sio
 
 
 def random_sampling(pc, num_samples, replace=None, return_choices=False):
-    '''
+    """
     Random Sampling.
 
     Sampling point cloud to num_samples points.
@@ -20,7 +20,8 @@ def random_sampling(pc, num_samples, replace=None, return_choices=False):
 
     Returns:
         pc (ndarray): Point cloud after sampling.
-    '''
+    """
+
     if replace is None:
         replace = (pc.shape[0] < num_samples)
     choices = np.random.choice(pc.shape[0], num_samples, replace=replace)
@@ -57,7 +58,7 @@ class SUNRGBDInstance(object):
 
 
 class SUNRGBDData(object):
-    '''
+    """
     SUNRGBD Data
 
     Generate scannet infos for sunrgbd_converter
@@ -66,7 +67,7 @@ class SUNRGBDData(object):
         root_path (str): Root path of the raw data.
         split (str): Set split type of the data. Default: 'train'.
         use_v1 (bool): Whether to use v1. Default: False.
-    '''
+    """
 
     def __init__(self, root_path, split='train', use_v1=False):
         self.root_dir = root_path
@@ -100,7 +101,7 @@ class SUNRGBDData(object):
 
     def get_image(self, idx):
         img_filename = os.path.join(self.image_dir, f'{idx:06d}.jpg')
-        return cv2.imread(img_filename)
+        return mmcv.imread(img_filename)
 
     def get_image_shape(self, idx):
         image = self.get_image(idx)
@@ -125,23 +126,21 @@ class SUNRGBDData(object):
         objects = [SUNRGBDInstance(line) for line in lines]
         return objects
 
-    def get_sunrgbd_infos(self,
-                          num_workers=4,
-                          has_label=True,
-                          sample_id_list=None):
-        '''
-        Get sunrgbd infos.
+    def get_infos(self, num_workers=4, has_label=True, sample_id_list=None):
+        """
+        Get data infos.
 
         This method gets information from the raw data.
 
         Args:
             num_workers (int): Number of threads to be used. Default: 4.
             has_label (bool): Whether the data has label. Default: True.
-            sample_id_list (List[int]): Index list of the sample. Default: None. # noqa: E501
+            sample_id_list (List[int]): Index list of the sample.
+                Default: None.
 
         Returns:
             infos (List[dict]): Information of the raw data.
-        '''
+        """
 
         def process_single_scene(sample_idx):
             print(f'{self.split} sample_idx: {sample_idx}')