Add visualization for lyft

mmdet3d/datasets/lyft_dataset.py

@@ -9,7 +9,8 @@ from pyquaternion import Quaternion
 
 from mmdet3d.core.evaluation.lyft_eval import lyft_eval
 from mmdet.datasets import DATASETS
-from ..core.bbox import LiDARInstance3DBoxes
+from ..core import show_result
+from ..core.bbox import Box3DMode, LiDARInstance3DBoxes
 from .custom_3d import Custom3DDataset
 
 
@@ -354,7 +355,9 @@ class LyftDataset(Custom3DDataset):
                  logger=None,
                  jsonfile_prefix=None,
                  csv_savepath=None,
-                 result_names=['pts_bbox']):
+                 result_names=['pts_bbox'],
+                 show=False,
+                 out_dir=None):
         """Evaluation in Lyft protocol.
 
         Args:
@@ -369,6 +372,10 @@ class LyftDataset(Custom3DDataset):
                 It includes the file path and the csv filename,
                 e.g., "a/b/filename.csv". If not specified,
                 the result will not be converted to csv file.
+            show (bool): Whether to visualize.
+                Default: False.
+            out_dir (str): Path to save the visualization results.
+                Default: None.
 
         Returns:
             dict[str: float]
@@ -387,8 +394,38 @@ class LyftDataset(Custom3DDataset):
 
         if tmp_dir is not None:
             tmp_dir.cleanup()
+
+        if show:
+            self.show(results, out_dir)
         return results_dict
 
+    def show(self, results, out_dir):
+        """Results visualization.
+
+        Args:
+            results (list[dict]): List of bounding boxes results.
+            out_dir (str): Output directory of visualization result.
+        """
+        for i, result in enumerate(results):
+            example = self.prepare_test_data(i)
+            points = example['points'][0]._data.numpy()
+            data_info = self.data_infos[i]
+            pts_path = data_info['lidar_path']
+            file_name = osp.split(pts_path)[-1].split('.')[0]
+            # for now we convert points into depth mode
+            points = points[..., [1, 0, 2]]
+            points[..., 0] *= -1
+            inds = result['pts_bbox']['scores_3d'] > 0.1
+            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor
+            gt_bboxes = Box3DMode.convert(gt_bboxes, Box3DMode.LIDAR,
+                                          Box3DMode.DEPTH)
+            gt_bboxes[..., 2] += gt_bboxes[..., 5] / 2
+            pred_bboxes = result['pts_bbox']['boxes_3d'][inds].tensor.numpy()
+            pred_bboxes = Box3DMode.convert(pred_bboxes, Box3DMode.LIDAR,
+                                            Box3DMode.DEPTH)
+            pred_bboxes[..., 2] += pred_bboxes[..., 5] / 2
+            show_result(points, gt_bboxes, pred_bboxes, out_dir, file_name)
+
     @staticmethod
     def json2csv(json_path, csv_savepath):
         """Convert the json file to csv format for submission.