[Enhance] Speed up evaluation on waymo (#2008)

* support fast eval on waymo

* support waymo evaluatioin more flexible and faster

* support waymo evaluatioin more flexible and faster

* renames

* add docs

* add guides for multi-thread evaluation toolkit

* fix docstring

* add download link for idx2metainfo

* add docstring

* set convert_kitti_format=False in Lidar-based methods

* fix docs

* add docstring

configs/_base_/datasets/waymoD5-3d-3class.py

@@ -151,7 +151,8 @@ val_evaluator = dict(
     ann_file='./data/waymo/kitti_format/waymo_infos_val.pkl',
     waymo_bin_file='./data/waymo/waymo_format/gt.bin',
     data_root='./data/waymo/waymo_format',
-    file_client_args=file_client_args)
+    file_client_args=file_client_args,
+    convert_kitti_format=False)
 test_evaluator = val_evaluator
 
 vis_backends = [dict(type='LocalVisBackend')]

configs/_base_/datasets/waymoD5-3d-car.py

@@ -135,7 +135,8 @@ val_evaluator = dict(
     type='WaymoMetric',
     ann_file='./data/waymo/kitti_format/waymo_infos_val.pkl',
     waymo_bin_file='./data/waymo/waymo_format/gt.bin',
-    data_root='./data/waymo/waymo_format')
+    data_root='./data/waymo/waymo_format',
+    convert_kitti_format=False)
 test_evaluator = val_evaluator
 
 vis_backends = [dict(type='LocalVisBackend')]

docs/en/advanced_guides/datasets/waymo_det.md

@@ -106,6 +106,7 @@ For evaluation on Waymo, please follow the [instruction](https://github.com/waym
 ```shell
 # download the code and enter the base directory
 git clone https://github.com/waymo-research/waymo-open-dataset.git waymo-od
+# git clone https://github.com/Abyssaledge/waymo-open-dataset-master waymo-od # if you want to use faster multi-thread version.
 cd waymo-od
 git checkout remotes/origin/master
 

docs/en/user_guides/dataset_prepare.md

@@ -110,7 +110,25 @@ Download Waymo open dataset V1.2 [HERE](https://waymo.com/open/download/) and it
 python tools/create_data.py waymo --root-path ./data/waymo/ --out-dir ./data/waymo/ --workers 128 --extra-tag waymo
 ```
 
-Note that if your local disk does not have enough space for saving converted data, you can change the `out-dir` to anywhere else. Just remember to create folders and prepare data there in advance and link them back to `data/waymo/kitti_format` after the data conversion.
+Note that:
+
+- If your local disk does not have enough space for saving converted data, you can change the `out-dir` to anywhere else. Just remember to create folders and prepare data there in advance and link them back to `data/waymo/kitti_format` after the data conversion.
+
+- If you want faster evaluation on Waymo, you can download the preprocessed [metainfo](https://download.openmmlab.com/mmdetection3d/data/waymo/idx2metainfo.pkl) containing `contextname` and `timestamp` to the directory `data/waymo/waymo_format/`. Then, the dataset config is modified like the following:
+
+  ```python
+  val_evaluator = dict(
+      type='WaymoMetric',
+      ann_file='./data/waymo/kitti_format/waymo_infos_val.pkl',
+      waymo_bin_file='./data/waymo/waymo_format/gt.bin',
+      data_root='./data/waymo/waymo_format',
+      file_client_args=file_client_args,
+      convert_kitti_format=True,
+      idx2metainfo='data/waymo/waymo_format/idx2metainfo.pkl'
+      )
+  ```
+
+  Now, this trick is only used for LiDAR-based detection methods.
 
 ### NuScenes
 

docs/zh_cn/advanced_guides/datasets/waymo_det.md

@@ -106,6 +106,7 @@ mmdetection3d
 ```shell
 # download the code and enter the base directory
 git clone https://github.com/waymo-research/waymo-open-dataset.git waymo-od
+# git clone https://github.com/Abyssaledge/waymo-open-dataset-master waymo-od # if you want to use faster multi-thread version.
 cd waymo-od
 git checkout remotes/origin/master
 

docs/zh_cn/user_guides/dataset_prepare.md

@@ -104,8 +104,23 @@ python tools/create_data.py kitti --root-path ./data/kitti --out-dir ./data/kitt
 python tools/create_data.py waymo --root-path ./data/waymo/ --out-dir ./data/waymo/ --workers 128 --extra-tag waymo
 ```
 
-注意，如果你的硬盘空间大小不足以存储转换后的数据，你可以将 `out-dir` 参数设定为别的路径。
-你只需要记得在那个路径下创建文件夹并下载数据，然后在数据预处理完成后将其链接回 `data/waymo/kitti_format` 即可。
+注意:
+
+- 如果你的硬盘空间大小不足以存储转换后的数据，你可以将 `out-dir` 参数设定为别的路径。
+  你只需要记得在那个路径下创建文件夹并下载数据，然后在数据预处理完成后将其链接回 `data/waymo/kitti_format` 即可
+- 如果你想在 Waymo 上进行更快的评估，你可以下载已经预处理好的[元信息文件](https://download.openmmlab.com/mmdetection3d/data/waymo/idx2metainfo.pkl) 并将其放置在 `data/waymo/waymo_format/` 目录下. 接着，你可以按照下方来更改数据集的配置：
+  ```python
+  val_evaluator = dict(
+      type='WaymoMetric',
+      ann_file='./data/waymo/kitti_format/waymo_infos_val.pkl',
+      waymo_bin_file='./data/waymo/waymo_format/gt.bin',
+      data_root='./data/waymo/waymo_format',
+      file_client_args=file_client_args,
+      convert_kitti_format=True,
+      idx2metainfo='data/waymo/waymo_format/idx2metainfo.pkl'
+      )
+  ```
+  目前这种方式仅限于纯点云任务。
 
 ### NuScenes
 

mmdet3d/datasets/waymo_dataset.py

@@ -96,7 +96,7 @@ class WaymoDataset(KittiDataset):
         self.max_sweeps = max_sweeps
         # we do not provide file_client_args to custom_3d init
         # because we want disk loading for info
-        # while ceph loading for KITTI2Waymo
+        # while ceph loading for Prediction2Waymo
         super().__init__(
             data_root=data_root,
             ann_file=ann_file,

mmdet3d/evaluation/functional/waymo_utils/__init__.py

 # Copyright (c) OpenMMLab. All rights reserved.
 
-from .prediction_kitti_to_waymo import KITTI2Waymo
+from .prediction_to_waymo import Prediction2Waymo
 
-__all__ = ['KITTI2Waymo']
+__all__ = ['Prediction2Waymo']

mmdet3d/evaluation/functional/waymo_utils/prediction_kitti_to_waymo.py → mmdet3d/evaluation/functional/waymo_utils/prediction_to_waymo.py

@@ -5,29 +5,33 @@ r"""Adapted from `Waymo to KITTI converter
 
 try:
     from waymo_open_dataset import dataset_pb2 as open_dataset
+    from waymo_open_dataset import label_pb2
+    from waymo_open_dataset.protos import metrics_pb2
+    from waymo_open_dataset.protos.metrics_pb2 import Objects
 except ImportError:
+    Objects = None
     raise ImportError(
         'Please run "pip install waymo-open-dataset-tf-2-1-0==1.2.0" '
         'to install the official devkit first.')
 
 from glob import glob
 from os.path import join
+from typing import List, Optional
 
 import mmengine
 import numpy as np
 import tensorflow as tf
-from waymo_open_dataset import label_pb2
-from waymo_open_dataset.protos import metrics_pb2
 
 
-class KITTI2Waymo(object):
-    """KITTI predictions to Waymo converter.
+class Prediction2Waymo(object):
+    """Predictions to Waymo converter. The format of prediction results could
+    be original format or kitti-format.
 
     This class serves as the converter to change predictions from KITTI to
     Waymo format.
 
     Args:
-        kitti_result_files (list[dict]): Predictions in KITTI format.
+        results (list[dict]): Prediction results.
         waymo_tfrecords_dir (str): Directory to load waymo raw data.
         waymo_results_save_dir (str): Directory to save converted predictions
             in waymo format (.bin files).
@@ -35,33 +39,47 @@ class KITTI2Waymo(object):
             predictions in waymo format (.bin file), like 'a/b/c.bin'.
         prefix (str): Prefix of filename. In general, 0 for training, 1 for
             validation and 2 for testing.
-        workers (str): Number of parallel processes.
+        classes (dict): A list of class name.
+        workers (str): Number of parallel processes. Defaults to 2.
+        file_client_args (str): File client for reading gt in waymo format.
+            Defaults to ``dict(backend='disk')``.
+        from_kitti_format (bool, optional): Whether the reuslts are kitti
+            format. Defaults to False.
+        idx2metainfo (Optional[dict], optional): The mapping from sample_idx to
+            metainfo. The metainfo must contain the keys: 'idx2contextname' and
+            'idx2timestamp'. Defaults to None.
     """
 
     def __init__(self,
-                 kitti_result_files,
-                 waymo_tfrecords_dir,
-                 waymo_results_save_dir,
-                 waymo_results_final_path,
-                 prefix,
-                 workers=64,
-                 file_client_args=dict(backend='disk')):
-
-        self.kitti_result_files = kitti_result_files
+                 results: List[dict],
+                 waymo_tfrecords_dir: str,
+                 waymo_results_save_dir: str,
+                 waymo_results_final_path: str,
+                 prefix: str,
+                 classes: dict,
+                 workers: int = 2,
+                 file_client_args: dict = dict(backend='disk'),
+                 from_kitti_format: bool = False,
+                 idx2metainfo: Optional[dict] = None):
+
+        self.results = results
         self.waymo_tfrecords_dir = waymo_tfrecords_dir
         self.waymo_results_save_dir = waymo_results_save_dir
         self.waymo_results_final_path = waymo_results_final_path
         self.prefix = prefix
+        self.classes = classes
         self.workers = int(workers)
         self.file_client_args = file_client_args
-        self.name2idx = {}
-        for idx, result in enumerate(kitti_result_files):
-            if len(result['sample_id']) > 0:
-                self.name2idx[str(result['sample_id'][0])] = idx
+        self.from_kitti_format = from_kitti_format
+        if idx2metainfo is not None:
+            self.idx2metainfo = idx2metainfo
+            # If ``fast_eval``, the metainfo does not need to be read from
+            # original data online. It's preprocessed offline.
+            self.fast_eval = True
+        else:
+            self.fast_eval = False
 
-        # turn on eager execution for older tensorflow versions
-        if int(tf.__version__.split('.')[0]) < 2:
-            tf.enable_eager_execution()
+        self.name2idx = {}
 
         self.k2w_cls_map = {
             'Car': label_pb2.Label.TYPE_VEHICLE,
@@ -70,12 +88,28 @@ class KITTI2Waymo(object):
             'Cyclist': label_pb2.Label.TYPE_CYCLIST,
         }
 
-        self.T_ref_to_front_cam = np.array([[0.0, 0.0, 1.0, 0.0],
-                                            [-1.0, 0.0, 0.0, 0.0],
-                                            [0.0, -1.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 1.0]])
+        if self.from_kitti_format:
+            self.T_ref_to_front_cam = np.array([[0.0, 0.0, 1.0, 0.0],
+                                                [-1.0, 0.0, 0.0, 0.0],
+                                                [0.0, -1.0, 0.0, 0.0],
+                                                [0.0, 0.0, 0.0, 1.0]])
+            # ``sample_idx`` of the sample in kitti-format is an array
+            for idx, result in enumerate(results):
+                if len(result['sample_idx']) > 0:
+                    self.name2idx[str(result['sample_idx'][0])] = idx
+        else:
+            # ``sample_idx`` of the sample in the original prediction
+            # is an int value.
+            for idx, result in enumerate(results):
+                self.name2idx[str(result['sample_idx'])] = idx
+
+        if not self.fast_eval:
+            # need to read original '.tfrecord' file
+            self.get_file_names()
+            # turn on eager execution for older tensorflow versions
+            if int(tf.__version__.split('.')[0]) < 2:
+                tf.enable_eager_execution()
 
-        self.get_file_names()
         self.create_folder()
 
     def get_file_names(self):
@@ -207,22 +241,30 @@ class KITTI2Waymo(object):
 
             filename = f'{self.prefix}{file_idx:03d}{frame_num:03d}'
 
-            for camera in frame.context.camera_calibrations:
-                # FRONT = 1, see dataset.proto for details
-                if camera.name == 1:
-                    T_front_cam_to_vehicle = np.array(
-                        camera.extrinsic.transform).reshape(4, 4)
-
-            T_k2w = T_front_cam_to_vehicle @ self.T_ref_to_front_cam
-
             context_name = frame.context.name
             frame_timestamp_micros = frame.timestamp_micros
 
             if filename in self.name2idx:
-                kitti_result = \
-                    self.kitti_result_files[self.name2idx[filename]]
-                objects = self.parse_objects(kitti_result, T_k2w, context_name,
-                                             frame_timestamp_micros)
+                if self.from_kitti_format:
+                    for camera in frame.context.camera_calibrations:
+                        # FRONT = 1, see dataset.proto for details
+                        if camera.name == 1:
+                            T_front_cam_to_vehicle = np.array(
+                                camera.extrinsic.transform).reshape(4, 4)
+
+                    T_k2w = T_front_cam_to_vehicle @ self.T_ref_to_front_cam
+
+                    kitti_result = \
+                        self.results[self.name2idx[filename]]
+                    objects = self.parse_objects(kitti_result, T_k2w,
+                                                 context_name,
+                                                 frame_timestamp_micros)
+                else:
+                    index = self.name2idx[filename]
+                    objects = self.parse_objects_from_origin(
+                        self.results[index], context_name,
+                        frame_timestamp_micros)
+
             else:
                 print(filename, 'not found.')
                 objects = metrics_pb2.Objects()
@@ -232,11 +274,100 @@ class KITTI2Waymo(object):
                     'wb') as f:
                 f.write(objects.SerializeToString())
 
+    def convert_one_fast(self, res_index: int):
+        """Convert action for single file. It read the metainfo from the
+        preprocessed file offline and will be faster.
+
+        Args:
+            res_index (int): The indices of the results.
+        """
+        sample_idx = self.results[res_index]['sample_idx']
+        if len(self.results[res_index]['pred_instances_3d']) > 0:
+            objects = self.parse_objects_from_origin(
+                self.results[res_index],
+                self.idx2metainfo[str(sample_idx)]['contextname'],
+                self.idx2metainfo[str(sample_idx)]['timestamp'])
+        else:
+            print(sample_idx, 'not found.')
+            objects = metrics_pb2.Objects()
+
+        with open(
+                join(self.waymo_results_save_dir, f'{sample_idx}.bin'),
+                'wb') as f:
+            f.write(objects.SerializeToString())
+
+    def parse_objects_from_origin(self, result: dict, contextname: str,
+                                  timestamp: str) -> Objects:
+        """Parse obejcts from the original prediction results.
+
+        Args:
+            result (dict): The original prediction results.
+            contextname (str): The ``contextname`` of sample in waymo.
+            timestamp (str): The ``timestamp`` of sample in waymo.
+
+        Returns:
+            metrics_pb2.Objects: The parsed object.
+        """
+        lidar_boxes = result['pred_instances_3d']['bboxes_3d'].tensor
+        scores = result['pred_instances_3d']['scores_3d']
+        labels = result['pred_instances_3d']['labels_3d']
+
+        def parse_one_object(index):
+            class_name = self.classes[labels[index].item()]
+
+            box = label_pb2.Label.Box()
+            height = lidar_boxes[index][5].item()
+            heading = lidar_boxes[index][6].item()
+
+            while heading < -np.pi:
+                heading += 2 * np.pi
+            while heading > np.pi:
+                heading -= 2 * np.pi
+
+            box.center_x = lidar_boxes[index][0].item()
+            box.center_y = lidar_boxes[index][1].item()
+            box.center_z = lidar_boxes[index][2].item() + height / 2
+            box.length = lidar_boxes[index][3].item()
+            box.width = lidar_boxes[index][4].item()
+            box.height = height
+            box.heading = heading
+
+            o = metrics_pb2.Object()
+            o.object.box.CopyFrom(box)
+            o.object.type = self.k2w_cls_map[class_name]
+            o.score = scores[index].item()
+            o.context_name = contextname
+            o.frame_timestamp_micros = timestamp
+
+            return o
+
+        objects = metrics_pb2.Objects()
+        for i in range(len(lidar_boxes)):
+            objects.objects.append(parse_one_object(i))
+
+        return objects
+
     def convert(self):
         """Convert action."""
         print('Start converting ...')
-        mmengine.track_parallel_progress(self.convert_one, range(len(self)),
-                                         self.workers)
+        convert_func = self.convert_one_fast if self.fast_eval else \
+            self.convert_one
+
+        # from torch.multiprocessing import set_sharing_strategy
+        # # Force using "file_system" sharing strategy for stability
+        # set_sharing_strategy("file_system")
+
+        # mmengine.track_parallel_progress(convert_func, range(len(self)),
+        #                                  self.workers)
+
+        # TODO: Support multiprocessing. Now, multiprocessing evaluation will
+        # cause shared memory error in torch-1.10 and torch-1.11. Details can
+        # be seen in https://github.com/pytorch/pytorch/issues/67864.
+        prog_bar = mmengine.ProgressBar(len(self))
+        for i in range(len(self)):
+            convert_func(i)
+            prog_bar.update()
+
         print('\nFinished ...')
 
         # combine all files into one .bin
@@ -248,7 +379,8 @@ class KITTI2Waymo(object):
 
     def __len__(self):
         """Length of the filename list."""
-        return len(self.waymo_tfrecord_pathnames)
+        return len(self.results) if self.fast_eval else len(
+            self.waymo_tfrecord_pathnames)
 
     def transform(self, T, x, y, z):
         """Transform the coordinates with matrix T.

mmdet3d/evaluation/metrics/waymo_metric.py

@@ -36,6 +36,9 @@ class WaymoMetric(KittiMetric):
             names to disambiguate homonymous metrics of different evaluators.
             If prefix is not provided in the argument, self.default_prefix
             will be used instead. Defaults to None.
+        convert_kitti_format (bool, optional): Whether convert the reuslts to
+            kitti format. Now, in order to be compatible with camera-based
+            methods, defaults to True.
         pklfile_prefix (str, optional): The prefix of pkl files, including
             the file path and the prefix of filename, e.g., "a/b/prefix".
             If not specified, a temp file will be created. Default: None.
@@ -54,6 +57,11 @@ class WaymoMetric(KittiMetric):
             from different ranks during distributed training. Must be 'cpu' or
             'gpu'. Defaults to 'cpu'.
         file_client_args (dict): file client for reading gt in waymo format.
+            Defaults to ``dict(backend='disk')``.
+        idx2metainfo (Optional[str], optional): The file path of the metainfo
+            in waymmo. It stores the mapping from sample_idx to metainfo.
+            The metainfo must contain the keys: 'idx2contextname' and
+            'idx2timestamp'. Defaults to None.
     """
     num_cams = 5
 
@@ -64,6 +72,7 @@ class WaymoMetric(KittiMetric):
                  split: str = 'training',
                  metric: Union[str, List[str]] = 'mAP',
                  pcd_limit_range: List[float] = [-85, -85, -5, 85, 85, 5],
+                 convert_kitti_format: bool = True,
                  prefix: Optional[str] = None,
                  pklfile_prefix: str = None,
                  submission_prefix: str = None,
@@ -71,12 +80,20 @@ class WaymoMetric(KittiMetric):
                  default_cam_key: str = 'CAM_FRONT',
                  use_pred_sample_idx: bool = False,
                  collect_device: str = 'cpu',
-                 file_client_args: dict = dict(backend='disk')):
+                 file_client_args: dict = dict(backend='disk'),
+                 idx2metainfo: Optional[str] = None):
         self.waymo_bin_file = waymo_bin_file
         self.data_root = data_root
         self.split = split
         self.task = task
         self.use_pred_sample_idx = use_pred_sample_idx
+        self.convert_kitti_format = convert_kitti_format
+
+        if idx2metainfo is not None:
+            self.idx2metainfo = mmengine.load(idx2metainfo)
+        else:
+            self.idx2metainfo = None
+
         super().__init__(
             ann_file=ann_file,
             metric=metric,
@@ -104,6 +121,8 @@ class WaymoMetric(KittiMetric):
 
         # load annotations
         self.data_infos = load(self.ann_file)['data_list']
+        assert len(results) == len(self.data_infos), \
+            'invalid list length of network outputs'
         # different from kitti, waymo do not need to convert the ann file
         # handle the mono3d task
         if self.task == 'mono_det':
@@ -131,7 +150,7 @@ class WaymoMetric(KittiMetric):
 
                     # TODO check if need to modify the sample id
                     # TODO check when will use it except for evaluation.
-                    camera_info['sample_id'] = info['sample_id']
+                    camera_info['sample_idx'] = info['sample_idx']
                     new_data_infos.append(camera_info)
             self.data_infos = new_data_infos
 
@@ -142,8 +161,6 @@ class WaymoMetric(KittiMetric):
             eval_tmp_dir = None
             pklfile_prefix = self.pklfile_prefix
 
-        # load annotations
-
         result_dict, tmp_dir = self.format_results(
             results,
             pklfile_prefix=pklfile_prefix,
@@ -186,11 +203,7 @@ class WaymoMetric(KittiMetric):
                 f'compute_detection_metrics_main {pklfile_prefix}.bin ' + \
                 f'{self.waymo_bin_file}'
             print(eval_str)
-            ret_bytes = subprocess.check_output(
-                'mmdet3d/evaluation/functional/waymo_utils/' +
-                f'compute_detection_metrics_main {pklfile_prefix}.bin ' +
-                f'{self.waymo_bin_file}',
-                shell=True)
+            ret_bytes = subprocess.check_output(eval_str, shell=True)
             ret_texts = ret_bytes.decode('utf-8')
             print_log(ret_texts, logger=logger)
 
@@ -292,7 +305,7 @@ class WaymoMetric(KittiMetric):
                        pklfile_prefix: str = None,
                        submission_prefix: str = None,
                        classes: List[str] = None):
-        """Format the results to pkl file.
+        """Format the results to bin file.
 
         Args:
             results (list[dict]): Testing results of the
@@ -313,9 +326,22 @@ class WaymoMetric(KittiMetric):
                 the formatted result, tmp_dir is the temporal directory created
                 for saving json files when jsonfile_prefix is not specified.
         """
-        result_files, tmp_dir = super().format_results(results, pklfile_prefix,
-                                                       submission_prefix,
-                                                       classes)
+        waymo_save_tmp_dir = tempfile.TemporaryDirectory()
+        waymo_results_save_dir = waymo_save_tmp_dir.name
+        waymo_results_final_path = f'{pklfile_prefix}.bin'
+
+        if self.convert_kitti_format:
+            results_kitti_format, tmp_dir = super().format_results(
+                results, pklfile_prefix, submission_prefix, classes)
+            final_results = results_kitti_format['pred_instances_3d']
+        else:
+            final_results = results
+            for i, res in enumerate(final_results):
+                # Actually, `sample_idx` here is the filename without suffix.
+                # It's for identitying the sample in formating.
+                res['sample_idx'] = self.data_infos[i]['sample_idx']
+                res['pred_instances_3d']['bboxes_3d'].limit_yaw(
+                    offset=0.5, period=np.pi * 2)
 
         waymo_root = self.data_root
         if self.split == 'training':
@@ -326,21 +352,23 @@ class WaymoMetric(KittiMetric):
             prefix = '2'
         else:
             raise ValueError('Not supported split value.')
-        waymo_save_tmp_dir = tempfile.TemporaryDirectory()
-        waymo_results_save_dir = waymo_save_tmp_dir.name
-        waymo_results_final_path = f'{pklfile_prefix}.bin'
-        from ..functional.waymo_utils.prediction_kitti_to_waymo import \
-            KITTI2Waymo
-        converter = KITTI2Waymo(
-            result_files['pred_instances_3d'],
+
+        from ..functional.waymo_utils.prediction_to_waymo import \
+            Prediction2Waymo
+        converter = Prediction2Waymo(
+            final_results,
             waymo_tfrecords_dir,
             waymo_results_save_dir,
             waymo_results_final_path,
             prefix,
-            file_client_args=self.file_client_args)
+            classes,
+            file_client_args=self.file_client_args,
+            from_kitti_format=self.convert_kitti_format,
+            idx2metainfo=self.idx2metainfo)
         converter.convert()
         waymo_save_tmp_dir.cleanup()
-        return result_files, waymo_save_tmp_dir
+
+        return final_results, waymo_save_tmp_dir
 
     def merge_multi_view_boxes(self, box_dict_per_frame: List[dict],
                                cam0_info: dict):
@@ -405,7 +433,7 @@ class WaymoMetric(KittiMetric):
             box3d_lidar=box_preds_lidar.tensor.numpy(),
             scores=scores.numpy(),
             label_preds=labels.numpy(),
-            sample_idx=box_dict['sample_id'],
+            sample_idx=box_dict['sample_idx'],
         )
         return merged_box_dict
 
@@ -431,8 +459,6 @@ class WaymoMetric(KittiMetric):
         Returns:
             list[dict]: A list of dictionaries with the kitti format.
         """
-        assert len(net_outputs) == len(self.data_infos), \
-            'invalid list length of network outputs'
         if submission_prefix is not None:
             mmengine.mkdir_or_exist(submission_prefix)
 
@@ -544,7 +570,7 @@ class WaymoMetric(KittiMetric):
                 # In waymo validation sample_idx in prediction is 000xxx
                 # but in info file it is 1000xxx
                 save_sample_idx = box_dict['sample_idx']
-            annos[-1]['sample_id'] = np.array(
+            annos[-1]['sample_idx'] = np.array(
                 [save_sample_idx] * len(annos[-1]['score']), dtype=np.int64)
 
             det_annos += annos
@@ -566,7 +592,7 @@ class WaymoMetric(KittiMetric):
         Args:
             box_dict (dict): Box dictionaries to be converted.
 
-                - boxes_3d (:obj:`LiDARInstance3DBoxes`): 3D bounding boxes.
+                - bboxes_3d (:obj:`LiDARInstance3DBoxes`): 3D bounding boxes.
                 - scores_3d (torch.Tensor): Scores of boxes.
                 - labels_3d (torch.Tensor): Class labels of boxes.
             info (dict): Data info.
@@ -587,7 +613,7 @@ class WaymoMetric(KittiMetric):
         box_preds = box_dict['bboxes_3d']
         scores = box_dict['scores_3d']
         labels = box_dict['labels_3d']
-        sample_idx = info['sample_id']
+        sample_idx = info['sample_idx']
         box_preds.limit_yaw(offset=0.5, period=np.pi * 2)
 
         if len(box_preds) == 0: