[Feature] Create SemanticKITTI pkl

tools/create_data.py

@@ -6,6 +6,7 @@
 from tools.dataset_converters import kitti_converter as kitti
 from tools.dataset_converters import lyft_converter as lyft_converter
 from tools.dataset_converters import nuscenes_converter as nuscenes_converter
+from tools.dataset_converters import semantickitti_converter
 from tools.dataset_converters.create_gt_database import (
     GTDatabaseCreater, create_groundtruth_database)
 from tools.dataset_converters.update_infos_to_v2 import update_pkl_infos
@@ -219,6 +220,17 @@ def waymo_data_prep(root_path,
         num_worker=workers).create()
 
 
+def semantickitti_data_prep(info_prefix, out_dir):
+    """Prepare the info file for SemanticKITTI dataset.
+
+    Args:
+        info_prefix (str): The prefix of info filenames.
+        out_dir (str): Output directory of the generated info file.
+    """
+    semantickitti_converter.create_semantickitti_info_file(
+        info_prefix, out_dir)
+
+
 parser = argparse.ArgumentParser(description='Data converter arg parser')
 parser.add_argument('dataset', metavar='kitti', help='name of the dataset')
 parser.add_argument(
@@ -333,5 +345,8 @@ def waymo_data_prep(root_path,
             info_prefix=args.extra_tag,
             out_dir=args.out_dir,
             workers=args.workers)
+    elif args.dataset == 'semantickitti':
+        semantickitti_data_prep(
+            info_prefix=args.extra_tag, out_dir=args.out_dir)
     else:
         raise NotImplementedError(f'Don\'t support {args.dataset} dataset.')

tools/dataset_converters/semantickitti_converter.py

@@ -0,0 +1,87 @@
+from os import path as osp
+from pathlib import Path
+
+import mmengine
+
+total_num = {
+    0: 4541,
+    1: 1101,
+    2: 4661,
+    3: 801,
+    4: 271,
+    5: 2761,
+    6: 1101,
+    7: 1101,
+    8: 4071,
+    9: 1591,
+    10: 1201,
+    11: 921,
+    12: 1061,
+    13: 3281,
+    14: 631,
+    15: 1901,
+    16: 1731,
+    17: 491,
+    18: 1801,
+    19: 4981,
+    20: 831,
+    21: 2721,
+}
+fold_split = {
+    'train': [0, 1, 2, 3, 4, 5, 6, 7, 9, 10],
+    'val': [8],
+    'test': [11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
+}
+split_list = ['train', 'valid', 'test']
+
+
+def get_semantickitti_info(split):
+    data_infos = dict()
+    data_infos['metainfo'] = dict(DATASET='SemanticKITTI')
+    data_list = []
+    for i_folder in fold_split[split]:
+        for j in range(0, total_num[i_folder]):
+            data_list.append({
+                'lidar_points': {
+                    'lidar_path':
+                    osp.join('sequences',
+                             str(i_folder).zfill(2), 'velodyne',
+                             str(j).zfill(6) + '.bin')
+                },
+                'pts_semantic_mask_path':
+                osp.join('sequences',
+                         str(i_folder).zfill(2), 'labels',
+                         str(j).zfill(6) + '.label'),
+                'sample_id':
+                str(i_folder) + str(j)
+            })
+    data_infos.update(dict(data_list=data_list))
+    return data_infos
+
+
+def create_semantickitti_info_file(pkl_prefix='semantickitti', save_path=None):
+    """Create info file of SemanticKITTI dataset.
+
+    Directly generate info file without raw data.
+
+    Args:
+        pkl_prefix (str, optional): Prefix of the info file to be generated.
+            Default: 'semantickitti'.
+        save_path (str, optional): Path to save the info file.
+            Default: None.
+    """
+    print('Generate info.')
+    save_path = Path(save_path)
+
+    semantickitti_infos_train = get_semantickitti_info(split='train')
+    filename = save_path / f'{pkl_prefix}_infos_train.pkl'
+    print(f'SemanticKITTI info train file is saved to {filename}')
+    mmengine.dump(semantickitti_infos_train, filename)
+    semantickitti_infos_val = get_semantickitti_info(split='val')
+    filename = save_path / f'{pkl_prefix}_infos_val.pkl'
+    print(f'SemanticKITTI info val file is saved to {filename}')
+    mmengine.dump(semantickitti_infos_val, filename)
+    semantickitti_infos_test = get_semantickitti_info(split='test')
+    filename = save_path / f'{pkl_prefix}_infos_test.pkl'
+    print(f'SemanticKITTI info test file is saved to {filename}')
+    mmengine.dump(semantickitti_infos_test, filename)