[Refactor] update `metainfo` in pkl file and add `categories` into `metainfo` (#1934)

* update metainfo in pkl file

* update

* add version for waymo

* update kitti metric

mmdet3d/evaluation/metrics/kitti_metric.py

@@ -74,25 +74,19 @@ class KittiMetric(BaseMetric):
                 raise KeyError("metric should be one of 'bbox', 'img_bbox', "
                                'but got {metric}.')
 
-    def convert_annos_to_kitti_annos(
-        self,
-        data_annos: list,
-        classes: list = [
-            'Pedestrian', 'Cyclist', 'Car', 'Van', 'Truck', 'Person_sitting',
-            'Tram', 'Misc'
-        ]
-    ) -> list:
+    def convert_annos_to_kitti_annos(self, data_infos: dict) -> list:
         """Convert loading annotations to Kitti annotations.
 
         Args:
-            data_annos (list[dict]): Annotations loaded from ann_file.
-            classes (list[str]): Classes used in the dataset. Default used
-                ['Pedestrian', 'Cyclist', 'Car', 'Van', 'Truck',
-                'Person_sitting', 'Tram', 'Misc'].
+            data_infos (dict): Data infos including metainfo and annotations
+                loaded from ann_file.
 
         Returns:
             List[dict]: List of Kitti annotations.
         """
+        cat2label = data_infos['metainfo']['categories']
+        data_annos = data_infos['data_list']
+        label2cat = dict((v, k) for (k, v) in cat2label.items())
         assert 'instances' in data_annos[0]
         for i, annos in enumerate(data_annos):
             if len(annos['instances']) == 0:
@@ -120,11 +114,8 @@ class KittiMetric(BaseMetric):
                     'score': []
                 }
                 for instance in annos['instances']:
-                    labels = instance['bbox_label']
-                    if labels == -1:
-                        kitti_annos['name'].append('DontCare')
-                    else:
-                        kitti_annos['name'].append(classes[labels])
+                    label = instance['bbox_label']
+                    kitti_annos['name'].append(label2cat[label])
                     kitti_annos['truncated'].append(instance['truncated'])
                     kitti_annos['occluded'].append(instance['occluded'])
                     kitti_annos['alpha'].append(instance['alpha'])
@@ -179,9 +170,8 @@ class KittiMetric(BaseMetric):
         self.classes = self.dataset_meta['CLASSES']
 
         # load annotations
-        pkl_annos = load(
-            self.ann_file, file_client_args=self.file_client_args)['data_list']
-        self.data_infos = self.convert_annos_to_kitti_annos(pkl_annos)
+        pkl_infos = load(self.ann_file, file_client_args=self.file_client_args)
+        self.data_infos = self.convert_annos_to_kitti_annos(pkl_infos)
         result_dict, tmp_dir = self.format_results(
             results,
             pklfile_prefix=self.pklfile_prefix,

tools/dataset_converters/update_infos_to_v2.py

@@ -264,10 +264,6 @@ def update_nuscenes_infos(pkl_path, out_dir):
         'CLASSES':
         ('car', 'truck', 'trailer', 'bus', 'construction_vehicle', 'bicycle',
          'motorcycle', 'pedestrian', 'traffic_cone', 'barrier'),
-        'DATASET':
-        'Nuscenes',
-        'version':
-        data_list['metadata']['version']
     }
     nusc = NuScenes(
         version=data_list['metadata']['version'],
@@ -367,7 +363,16 @@ def update_nuscenes_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(metainfo=METAINFO, data_list=converted_list)
+
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'nuscenes'
+    metainfo['version'] = data_list['metadata']['version']
+    metainfo['info_version'] = '1.1'
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -497,8 +502,16 @@ def update_kitti_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(
-        metainfo={'DATASET': 'KITTI'}, data_list=converted_list)
+
+    # dataset metainfo
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'kitti'
+    metainfo['info_version'] = '1.1'
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -558,8 +571,17 @@ def update_s3dis_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(
-        metainfo={'DATASET': 'S3DIS'}, data_list=converted_list)
+
+    # dataset metainfo
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 's3dis'
+    metainfo['info_version'] = '1.1'
+
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -623,8 +645,17 @@ def update_scannet_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(
-        metainfo={'DATASET': 'SCANNET'}, data_list=converted_list)
+
+    # dataset metainfo
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'scannet'
+    metainfo['info_version'] = '1.1'
+
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -692,8 +723,17 @@ def update_sunrgbd_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(
-        metainfo={'DATASET': 'SUNRGBD'}, data_list=converted_list)
+
+    # dataset metainfo
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'sunrgbd'
+    metainfo['info_version'] = '1.1'
+
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -709,10 +749,6 @@ def update_lyft_infos(pkl_path, out_dir):
         'CLASSES':
         ('car', 'truck', 'bus', 'emergency_vehicle', 'other_vehicle',
          'motorcycle', 'bicycle', 'pedestrian', 'animal'),
-        'DATASET':
-        'Nuscenes',
-        'version':
-        data_list['metadata']['version']
     }
     print('Start updating:')
     converted_list = []
@@ -797,7 +833,16 @@ def update_lyft_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(metainfo=METAINFO, data_list=converted_list)
+
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'lyft'
+    metainfo['version'] = data_list['metadata']['version']
+    metainfo['info_version'] = '1.1'
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')
 
@@ -995,8 +1040,18 @@ def update_waymo_infos(pkl_path, out_dir):
     out_path = osp.join(out_dir, pkl_name)
     print(f'Writing to output file: {out_path}.')
     print(f'ignore classes: {ignore_class_name}')
-    converted_data_info = dict(
-        metainfo={'DATASET': 'Waymo'}, data_list=converted_list)
+
+    # dataset metainfo
+    metainfo = dict()
+    metainfo['categories'] = {k: i for i, k in enumerate(METAINFO['CLASSES'])}
+    if ignore_class_name:
+        for ignore_class in ignore_class_name:
+            metainfo['categories'][ignore_class] = -1
+    metainfo['dataset'] = 'waymo'
+    metainfo['version'] = '1.2'
+    metainfo['info_version'] = '1.1'
+
+    converted_data_info = dict(metainfo=metainfo, data_list=converted_list)
 
     mmengine.dump(converted_data_info, out_path, 'pkl')