[Enhance] add plane info in generated pkl file

docs/en/datasets/kitti_det.md

@@ -6,7 +6,7 @@ This page provides specific tutorials about the usage of MMDetection3D for KITTI
 
 ## Prepare dataset
 
-You can download KITTI 3D detection data [HERE](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d) and unzip all zip files.
+You can download KITTI 3D detection data [HERE](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d) and unzip all zip files. Besides, the road planes could be downloaded from [HERE](https://drive.google.com/file/d/1d5mq0RXRnvHPVeKx6Q612z0YRO1t2wAp/view?usp=sharing), which are optional for data augmentation in the training for better performance.
 
 Like the general way to prepare dataset, it is recommended to symlink the dataset root to `$MMDETECTION3D/data`.
 
@@ -29,6 +29,7 @@ mmdetection3d
 │   │   │   ├── image_2
 │   │   │   ├── label_2
 │   │   │   ├── velodyne
+│   │   │   ├── planes (optional)
 ```
 
 ### Create KITTI dataset

docs/zh_cn/datasets/kitti_det.md

@@ -6,7 +6,7 @@
 
 ## 数据准备
 
-您可以在[这里](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d)下载 KITTI 3D 检测数据并解压缩所有 zip 文件。
+您可以在[这里](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=3d)下载 KITTI 3D 检测数据并解压缩所有 zip 文件。此外，您可以在[这里](https://drive.google.com/file/d/1d5mq0RXRnvHPVeKx6Q612z0YRO1t2wAp/view?usp=sharing)下载道路平面信息，其在训练过程中作为一个可选项，用来提高模型的性能.
 
 像准备数据集的一般方法一样，建议将数据集根目录链接到 `$MMDETECTION3D/data`。
 
@@ -29,6 +29,7 @@ mmdetection3d
 │   │   │   ├── image_2
 │   │   │   ├── label_2
 │   │   │   ├── velodyne
+│   │   │   ├── planes (optional)
 ```
 
 ### 创建 KITTI 数据集

setup.cfg

@@ -8,7 +8,7 @@ line_length = 79
 multi_line_output = 0
 known_standard_library = setuptools
 known_first_party = mmdet,mmseg,mmdet3d
-known_third_party = cv2,imageio,indoor3d_util,load_scannet_data,lyft_dataset_sdk,m2r,matplotlib,mmcv,nuimages,numba,numpy,nuscenes,pandas,plyfile,pycocotools,pyquaternion,pytest,pytorch_sphinx_theme,recommonmark,requests,scannet_utils,scipy,seaborn,shapely,skimage,tensorflow,terminaltables,torch,trimesh,ts,waymo_open_dataset
+known_third_party = cv2,imageio,indoor3d_util,load_scannet_data,lyft_dataset_sdk,m2r,matplotlib,mmcv,nuimages,numba,numpy,nuscenes,pandas,plyfile,pycocotools,pyquaternion,pytest,pytorch_sphinx_theme,recommonmark,requests,scannet_utils,scipy,seaborn,shapely,skimage,sphinx,tensorflow,terminaltables,torch,trimesh,ts,waymo_open_dataset
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY
 

tools/data_converter/kitti_converter.py

@@ -94,9 +94,12 @@ def create_kitti_info_file(data_path,
 
     Args:
         data_path (str): Path of the data root.
-        pkl_prefix (str): Prefix of the info file to be generated.
-        save_path (str): Path to save the info file.
-        relative_path (bool): Whether to use relative path.
+        pkl_prefix (str, optional): Prefix of the info file to be generated.
+            Default: 'kitti'.
+        save_path (str, optional): Path to save the info file.
+            Default: None.
+        relative_path (bool, optional): Whether to use relative path.
+            Default: True.
     """
     imageset_folder = Path(data_path) / 'ImageSets'
     train_img_ids = _read_imageset_file(str(imageset_folder / 'train.txt'))
@@ -113,6 +116,7 @@ def create_kitti_info_file(data_path,
         training=True,
         velodyne=True,
         calib=True,
+        plane=True,
         image_ids=train_img_ids,
         relative_path=relative_path)
     _calculate_num_points_in_gt(data_path, kitti_infos_train, relative_path)
@@ -124,6 +128,7 @@ def create_kitti_info_file(data_path,
         training=True,
         velodyne=True,
         calib=True,
+        plane=True,
         image_ids=val_img_ids,
         relative_path=relative_path)
     _calculate_num_points_in_gt(data_path, kitti_infos_val, relative_path)
@@ -140,6 +145,7 @@ def create_kitti_info_file(data_path,
         label_info=False,
         velodyne=True,
         calib=True,
+        plane=False,
         image_ids=test_img_ids,
         relative_path=relative_path)
     filename = save_path / f'{pkl_prefix}_infos_test.pkl'
@@ -158,10 +164,14 @@ def create_waymo_info_file(data_path,
 
     Args:
         data_path (str): Path of the data root.
-        pkl_prefix (str): Prefix of the info file to be generated.
-        save_path (str): Path to save the info file.
-        relative_path (bool): Whether to use relative path.
-        max_sweeps (int): Max sweeps before the detection frame to be used.
+        pkl_prefix (str, optional): Prefix of the info file to be generated.
+            Default: 'waymo'.
+        save_path (str, optional): Path to save the info file.
+            Default: None.
+        relative_path (bool, optional): Whether to use relative path.
+            Default: True.
+        max_sweeps (int, optional): Max sweeps before the detection frame
+            to be used. Default: 5.
     """
     imageset_folder = Path(data_path) / 'ImageSets'
     train_img_ids = _read_imageset_file(str(imageset_folder / 'train.txt'))

tools/data_converter/kitti_data_utils.py

@@ -59,6 +59,17 @@ def get_label_path(idx,
                                relative_path, exist_check, use_prefix_id)
 
 
+def get_plane_path(idx,
+                   prefix,
+                   training=True,
+                   relative_path=True,
+                   exist_check=True,
+                   info_type='planes',
+                   use_prefix_id=False):
+    return get_kitti_info_path(idx, prefix, info_type, '.txt', training,
+                               relative_path, exist_check, use_prefix_id)
+
+
 def get_velodyne_path(idx,
                       prefix,
                       training=True,
@@ -143,6 +154,7 @@ def get_kitti_image_info(path,
                          label_info=True,
                          velodyne=False,
                          calib=False,
+                         plane=False,
                          image_ids=7481,
                          extend_matrix=True,
                          num_worker=8,
@@ -251,6 +263,17 @@ def map_func(idx):
             calib_info['Tr_imu_to_velo'] = Tr_imu_to_velo
             info['calib'] = calib_info
 
+        if plane:
+            plane_path = get_plane_path(idx, path, training, relative_path)
+            if relative_path:
+                plane_path = str(root_path / plane_path)
+            with open(plane_path, 'r') as f:
+                lines = f.readlines()
+            assert lines[0] == '# Matrix\n'
+            assert lines[1] == 'WIDTH 4\n'
+            assert lines[2] == 'HEIGHT 1\n'
+            info['plane'] = np.array([float(i) for i in lines[3].split()])
+
         if annotations is not None:
             info['annos'] = annotations
             add_difficulty_to_annos(info)