[Enhance] add plane info in generated pkl file (#1162)

* add plane info in pkl file

* add_plane_info

* fix comments

* change to mmcv api

* change road plane link

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://download.openmmlab.com/mmdetection3d/data/train_planes.zip), which are optional for data augmentation during training for better performance. The road planes are generated by [AVOD](https://github.com/kujason/avod), you can see more details [HERE](https://github.com/kujason/avod/issues/19).
 
 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://download.openmmlab.com/mmdetection3d/data/train_planes.zip)下载道路平面信息，其在训练过程中作为一个可选项，用来提高模型的性能。道路平面信息由 [AVOD](https://github.com/kujason/avod) 生成，你可以在[这里](https://github.com/kujason/avod/issues/19)查看更多细节。
 
 像准备数据集的一般方法一样，建议将数据集根目录链接到 `$MMDETECTION3D/data`。
 
@@ -29,6 +29,7 @@ mmdetection3d
 │   │   │   ├── image_2
 │   │   │   ├── label_2
 │   │   │   ├── velodyne
+│   │   │   ├── planes (optional)
 ```
 
 ### 创建 KITTI 数据集

setup.cfg

@@ -13,4 +13,4 @@ no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY
 
 [codespell]
-ignore-words-list = ans,refridgerator,crate,hist,formating,dout,wan,nd,fo
+ignore-words-list = ans,refridgerator,crate,hist,formating,dout,wan,nd,fo,avod,AVOD

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,
+        with_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,
+        with_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,
+        with_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

 # Copyright (c) OpenMMLab. All rights reserved.
+import mmcv
 import numpy as np
 from collections import OrderedDict
 from concurrent import futures as futures
@@ -59,6 +60,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 +155,7 @@ def get_kitti_image_info(path,
                          label_info=True,
                          velodyne=False,
                          calib=False,
+                         with_plane=False,
                          image_ids=7481,
                          extend_matrix=True,
                          num_worker=8,
@@ -251,6 +264,13 @@ def get_kitti_image_info(path,
             calib_info['Tr_imu_to_velo'] = Tr_imu_to_velo
             info['calib'] = calib_info
 
+        if with_plane:
+            plane_path = get_plane_path(idx, path, training, relative_path)
+            if relative_path:
+                plane_path = str(root_path / plane_path)
+            lines = mmcv.list_from_file(plane_path)
+            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)