[Fix] add `with_plane` flag while data conversion (#1278)

* add with_plane flag for create_data.py

* update kitti datasets doc

docs/en/datasets/kitti_det.md

@@ -45,10 +45,12 @@ wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/sec
 wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/val.txt --no-check-certificate --content-disposition -O ./data/kitti/ImageSets/val.txt
 wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/trainval.txt --no-check-certificate --content-disposition -O ./data/kitti/ImageSets/trainval.txt
 
-python tools/create_data.py kitti --root-path ./data/kitti --out-dir ./data/kitti --extra-tag kitti
+
+python tools/create_data.py kitti --root-path ./data/kitti --out-dir ./data/kitti --extra-tag kitti --with-plane
+
 ```
 
-Note that if your local disk does not have enough space for saving converted data, you can change the `out-dir` to anywhere else.
+Note that if your local disk does not have enough space for saving converted data, you can change the `out-dir` to anywhere else, and you need to remove the `--with-plane` flag if `planes` are not prepared.
 
 The folder structure after processing should be as below
 
@@ -70,6 +72,7 @@ kitti
 │   ├── label_2
 │   ├── velodyne
 │   ├── velodyne_reduced
+│   ├── planes (optional)
 ├── kitti_gt_database
 │   ├── xxxxx.bin
 ├── kitti_infos_train.pkl

docs/zh_cn/datasets/kitti_det.md

@@ -45,10 +45,10 @@ wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/sec
 wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/val.txt --no-check-certificate --content-disposition -O ./data/kitti/ImageSets/val.txt
 wget -c  https://raw.githubusercontent.com/traveller59/second.pytorch/master/second/data/ImageSets/trainval.txt --no-check-certificate --content-disposition -O ./data/kitti/ImageSets/trainval.txt
 
-python tools/create_data.py kitti --root-path ./data/kitti --out-dir ./data/kitti --extra-tag kitti
+python tools/create_data.py kitti --root-path ./data/kitti --out-dir ./data/kitti --extra-tag kitti --with-plane
 ```
 
-需要注意的是，如果您的本地磁盘没有充足的存储空间来存储转换后的数据，您可以通过改变 `out-dir` 来指定其他任意的存储路径。
+需要注意的是，如果您的本地磁盘没有充足的存储空间来存储转换后的数据，您可以通过改变 `out-dir` 来指定其他任意的存储路径。如果您没有准备 `planes` 数据，您需要移除 `--with-plane` 标志。
 
 处理后的文件夹结构应该如下：
 
@@ -70,6 +70,7 @@ kitti
 │   ├── label_2
 │   ├── velodyne
 │   ├── velodyne_reduced
+│   ├── planes (optional)
 ├── kitti_gt_database
 │   ├── xxxxx.bin
 ├── kitti_infos_train.pkl

tools/create_data.py

@@ -9,7 +9,11 @@ from tools.data_converter import nuscenes_converter as nuscenes_converter
 from tools.data_converter.create_gt_database import create_groundtruth_database
 
 
-def kitti_data_prep(root_path, info_prefix, version, out_dir):
+def kitti_data_prep(root_path,
+                    info_prefix,
+                    version,
+                    out_dir,
+                    with_plane=False):
     """Prepare data related to Kitti dataset.
 
     Related data consists of '.pkl' files recording basic infos,
@@ -20,8 +24,10 @@ def kitti_data_prep(root_path, info_prefix, version, out_dir):
         info_prefix (str): The prefix of info filenames.
         version (str): Dataset version.
         out_dir (str): Output directory of the groundtruth database info.
+        with_plane (bool, optional): Whether to use plane information.
+            Default: False.
     """
-    kitti.create_kitti_info_file(root_path, info_prefix)
+    kitti.create_kitti_info_file(root_path, info_prefix, with_plane)
     kitti.create_reduced_point_cloud(root_path, info_prefix)
 
     info_train_path = osp.join(root_path, f'{info_prefix}_infos_train.pkl')
@@ -204,6 +210,10 @@ parser.add_argument(
     default=10,
     required=False,
     help='specify sweeps of lidar per example')
+parser.add_argument(
+    '--with-plane',
+    action='store_true',
+    help='Whether to use plane information for kitti.')
 parser.add_argument(
     '--out-dir',
     type=str,
@@ -221,7 +231,8 @@ if __name__ == '__main__':
             root_path=args.root_path,
             info_prefix=args.extra_tag,
             version=args.version,
-            out_dir=args.out_dir)
+            out_dir=args.out_dir,
+            with_plane=args.with_plane)
     elif args.dataset == 'nuscenes' and args.version != 'v1.0-mini':
         train_version = f'{args.version}-trainval'
         nuscenes_data_prep(

tools/data_converter/kitti_converter.py

@@ -87,6 +87,7 @@ def _calculate_num_points_in_gt(data_path,
 
 def create_kitti_info_file(data_path,
                            pkl_prefix='kitti',
+                           with_plane=False,
                            save_path=None,
                            relative_path=True):
     """Create info file of KITTI dataset.
@@ -97,6 +98,8 @@ def create_kitti_info_file(data_path,
         data_path (str): Path of the data root.
         pkl_prefix (str, optional): Prefix of the info file to be generated.
             Default: 'kitti'.
+        with_plane (bool, optional): Whether to use plane information.
+            Default: False.
         save_path (str, optional): Path to save the info file.
             Default: None.
         relative_path (bool, optional): Whether to use relative path.
@@ -117,7 +120,7 @@ def create_kitti_info_file(data_path,
         training=True,
         velodyne=True,
         calib=True,
-        with_plane=True,
+        with_plane=with_plane,
         image_ids=train_img_ids,
         relative_path=relative_path)
     _calculate_num_points_in_gt(data_path, kitti_infos_train, relative_path)
@@ -129,7 +132,7 @@ def create_kitti_info_file(data_path,
         training=True,
         velodyne=True,
         calib=True,
-        with_plane=True,
+        with_plane=with_plane,
         image_ids=val_img_ids,
         relative_path=relative_path)
     _calculate_num_points_in_gt(data_path, kitti_infos_val, relative_path)