Tracking eval BEV visualization (nutonomy#254)

* Replace ego_diff by ego_translation

* Implement basic visualization

* Render each object with unique color

* Encapsulate and expose renderer

* Pass output directory to renderer

* Visualize switches by changing front-back color

* Cleanup

* Remove warnings

* Fixed all but one unittest

* Fixed bug where ego_dist was not set

* Change output folder name

python-sdk/nuscenes/eval/common/data_classes.py

@@ -17,7 +17,6 @@ def __init__(self,
                  size: Tuple[float, float, float] = (0, 0, 0),
                  rotation: Tuple[float, float, float, float] = (0, 0, 0, 0),
                  velocity: Tuple[float, float] = (0, 0),
-                 ego_dist: float = 0.0,  # Distance to ego vehicle in meters.
                  num_pts: int = -1):  # Nbr. LIDAR or RADAR inside the box. Only for gt boxes.
 
         # Assert data for shape and NaNs.
@@ -35,9 +34,6 @@ def __init__(self,
         # Velocity can be NaN from our database for certain annotations.
         assert len(velocity) == 2, 'Error: Velocity must have 2 elements!'
 
-        assert type(ego_dist) == float, 'Error: ego_dist must be a float!'
-        assert not np.any(np.isnan(ego_dist)), 'Error: ego_dist may not be NaN!'
-
         assert type(num_pts) == int, 'Error: num_pts must be int!'
         assert not np.any(np.isnan(num_pts)), 'Error: num_pts may not be NaN!'
 
@@ -47,7 +43,6 @@ def __init__(self,
         self.size = size
         self.rotation = rotation
         self.velocity = velocity
-        self.ego_dist = ego_dist
         self.num_pts = num_pts
 
     def __repr__(self):

python-sdk/nuscenes/eval/common/loaders.py

@@ -191,8 +191,16 @@ def add_center_dist(nusc: NuScenes,
 
         for box in eval_boxes[sample_token]:
             # Both boxes and ego pose are given in global coord system, so distance can be calculated directly.
-            diff = np.array(pose_record['translation'][:2]) - np.array(box.translation[:2])
-            box.ego_dist = np.sqrt(np.sum(diff ** 2))
+            # Note that the z component of the ego pose is 0.
+            ego_translation = (box.translation[0] - pose_record['translation'][0],
+                               box.translation[1] - pose_record['translation'][1],
+                               box.translation[2] - pose_record['translation'][2])
+            if isinstance(box, DetectionBox):
+                box.ego_dist = np.sqrt(np.sum(np.array(ego_translation[:2]) ** 2))
+            elif isinstance(box, TrackingBox):
+                box.ego_translation = ego_translation
+            else:
+                raise NotImplementedError
 
     return eval_boxes
 

python-sdk/nuscenes/eval/detection/data_classes.py

@@ -327,7 +327,7 @@ def __init__(self,
                  detection_score: float = -1.0,  # GT samples do not have a score.
                  attribute_name: str = ''):  # Box attribute. Each box can have at most 1 attribute.
 
-        super().__init__(sample_token, translation, size, rotation, velocity, ego_dist, num_pts)
+        super().__init__(sample_token, translation, size, rotation, velocity, num_pts)
 
         assert detection_name is not None, 'Error: detection_name cannot be empty!'
         assert detection_name in DETECTION_NAMES, 'Error: Unknown detection_name %s' % detection_name
@@ -339,21 +339,22 @@ def __init__(self,
         assert not np.any(np.isnan(detection_score)), 'Error: detection_score may not be NaN!'
 
         # Assign.
+        self.ego_dist = ego_dist
         self.detection_name = detection_name
-        self.attribute_name = attribute_name
         self.detection_score = detection_score
+        self.attribute_name = attribute_name
 
     def __eq__(self, other):
         return (self.sample_token == other.sample_token and
                 self.translation == other.translation and
                 self.size == other.size and
                 self.rotation == other.rotation and
                 self.velocity == other.velocity and
-                self.detection_name == other.detection_name and
-                self.attribute_name == other.attribute_name and
                 self.ego_dist == other.ego_dist and
+                self.num_pts == other.num_pts and
+                self.detection_name == other.detection_name and
                 self.detection_score == other.detection_score and
-                self.num_pts == other.num_pts)
+                self.attribute_name == other.attribute_name)
 
     def serialize(self) -> dict:
         """ Serialize instance into json-friendly format. """
@@ -363,11 +364,11 @@ def serialize(self) -> dict:
             'size': self.size,
             'rotation': self.rotation,
             'velocity': self.velocity,
-            'detection_name': self.detection_name,
-            'attribute_name': self.attribute_name,
             'ego_dist': self.ego_dist,
+            'num_pts': self.num_pts,
+            'detection_name': self.detection_name,
             'detection_score': self.detection_score,
-            'num_pts': self.num_pts
+            'attribute_name': self.attribute_name
         }
 
     @classmethod
@@ -378,11 +379,11 @@ def deserialize(cls, content: dict):
                    size=tuple(content['size']),
                    rotation=tuple(content['rotation']),
                    velocity=tuple(content['velocity']),
-                   detection_name=content['detection_name'],
-                   attribute_name=content['attribute_name'],
                    ego_dist=0.0 if 'ego_dist' not in content else float(content['ego_dist']),
+                   num_pts=-1 if 'num_pts' not in content else int(content['num_pts']),
+                   detection_name=content['detection_name'],
                    detection_score=-1.0 if 'detection_score' not in content else float(content['detection_score']),
-                   num_pts=-1 if 'num_pts' not in content else int(content['num_pts']))
+                   attribute_name=content['attribute_name'])
 
 
 class DetectionMetricDataList:

python-sdk/nuscenes/eval/tracking/algo.py

@@ -10,17 +10,19 @@
 py-motmetrics at:
 https://github.com/cheind/py-motmetrics
 """
+import os
 from typing import List, Dict, Callable, Tuple
 
 import numpy as np
 import pandas
 import sklearn
 import tqdm
 
+from nuscenes.eval.tracking.constants import MOT_METRIC_MAP, TRACKING_METRICS
 from nuscenes.eval.tracking.data_classes import TrackingBox, TrackingMetricData
-from nuscenes.eval.tracking.utils import print_threshold_metrics, create_motmetrics
 from nuscenes.eval.tracking.mot import MOTAccumulatorCustom
-from nuscenes.eval.tracking.constants import MOT_METRIC_MAP, TRACKING_METRICS
+from nuscenes.eval.tracking.render import TrackingRenderer
+from nuscenes.eval.tracking.utils import print_threshold_metrics, create_motmetrics
 
 
 class TrackingEvaluation(object):
@@ -33,7 +35,9 @@ def __init__(self,
                  min_recall: float,
                  num_thresholds: int,
                  metric_worst: Dict[str, float],
-                 verbose: bool = True):
+                 verbose: bool = True,
+                 output_dir: str = None,
+                 render_classes: List[str] = None):
         """
         Create a TrackingEvaluation object which computes all metrics for a given class.
         :param tracks_gt: The ground-truth tracks.
@@ -46,6 +50,8 @@ def __init__(self,
         :param metric_worst: Mapping from metric name to the fallback value assigned if a recall threshold
             is not achieved.
         :param verbose: Whether to print to stdout.
+        :param output_dir: Output directory to save renders.
+        :param render_classes: Classes to render to disk or None.
 
         Computes the metrics defined in:
         - Stiefelhagen 2008: Evaluating Multiple Object Tracking Performance: The CLEAR MOT Metrics.
@@ -64,6 +70,8 @@ def __init__(self,
         self.num_thresholds = num_thresholds
         self.metric_worst = metric_worst
         self.verbose = verbose
+        self.output_dir = output_dir
+        self.render_classes = [] if render_classes is None else render_classes
 
         self.n_scenes = len(self.tracks_gt)
 
@@ -212,6 +220,14 @@ def accumulate_threshold(self, threshold: float = None) -> Tuple[pandas.DataFram
             scene_tracks_gt = self.tracks_gt[scene_id]
             scene_tracks_pred = self.tracks_pred[scene_id]
 
+            # Visualize the boxes in this frame.
+            if self.class_name in self.render_classes and threshold is None:
+                save_path = os.path.join(self.output_dir, 'render', str(scene_id), self.class_name)
+                os.makedirs(save_path, exist_ok=True)
+                renderer = TrackingRenderer(save_path)
+            else:
+                renderer = None
+
             for timestamp in scene_tracks_gt.keys():
                 # Select only the current class.
                 frame_gt = scene_tracks_gt[timestamp]
@@ -254,9 +270,16 @@ def accumulate_threshold(self, threshold: float = None) -> Tuple[pandas.DataFram
                     match_ids = matches.HId.values
                     match_scores = [tt.tracking_score for tt in frame_pred if tt.tracking_id in match_ids]
                     scores.extend(match_scores)
+                else:
+                    events = None
 
-                # Increment the frame_id, unless there were no boxes (equivalent to what motmetrics does).
+                # Render the boxes in this frame.
+                if self.class_name in self.render_classes and threshold is None:
+                    renderer.render(events, timestamp, frame_gt, frame_pred)
+
+                # Increment the frame_id, unless there are no boxes (equivalent to what motmetrics does).
                 frame_id += 1
+
             accs.append(acc)
 
         # Merge accumulators

python-sdk/nuscenes/eval/tracking/data_classes.py

@@ -1,7 +1,7 @@
 # nuScenes dev-kit.
 # Code written by Holger Caesar, Caglayan Dicle and Oscar Beijbom, 2019.
 
-from typing import List, Dict, Tuple, Any
+from typing import Dict, Tuple, Any
 
 import numpy as np
 
@@ -254,13 +254,16 @@ def __init__(self,
                  size: Tuple[float, float, float] = (0, 0, 0),
                  rotation: Tuple[float, float, float, float] = (0, 0, 0, 0),
                  velocity: Tuple[float, float] = (0, 0),
-                 ego_dist: float = 0.0,  # Distance to ego vehicle in meters.
+                 ego_translation: [float, float, float] = (0, 0, 0),  # Translation to ego vehicle in meters.
                  num_pts: int = -1,  # Nbr. LIDAR or RADAR inside the box. Only for gt boxes.
                  tracking_id: str = '',  # Instance id of this object.
                  tracking_name: str = '',  # The class name used in the tracking challenge.
                  tracking_score: float = -1.0):  # Does not apply to GT.
 
-        super().__init__(sample_token, translation, size, rotation, velocity, ego_dist, num_pts)
+        super().__init__(sample_token, translation, size, rotation, velocity, num_pts)
+
+        assert len(ego_translation) == 3, 'Error: Translation must have 3 elements!'
+        assert not np.any(np.isnan(ego_translation)), 'Error: Translation may not be NaN!'
 
         assert tracking_name is not None, 'Error: tracking_name cannot be empty!'
         assert tracking_name in TRACKING_NAMES, 'Error: Unknown tracking_name %s' % tracking_name
@@ -269,17 +272,23 @@ def __init__(self,
         assert not np.any(np.isnan(tracking_score)), 'Error: tracking_score may not be NaN!'
 
         # Assign.
+        self.ego_translation = ego_translation
         self.tracking_id = tracking_id
         self.tracking_name = tracking_name
         self.tracking_score = tracking_score
 
+    @ property
+    def ego_dist(self) -> float:
+        """ Compute the distance from this box to the ego vehicle in 2D. """
+        return np.sqrt(np.sum(np.array(self.ego_translation[:2]) ** 2))
+
     def __eq__(self, other):
         return (self.sample_token == other.sample_token and
                 self.translation == other.translation and
                 self.size == other.size and
                 self.rotation == other.rotation and
                 self.velocity == other.velocity and
-                self.ego_dist == other.ego_dist and
+                self.ego_translation == other.ego_translation and
                 self.num_pts == other.num_pts and
                 self.tracking_id == other.tracking_id and
                 self.tracking_name == other.tracking_name and
@@ -293,7 +302,7 @@ def serialize(self) -> dict:
             'size': self.size,
             'rotation': self.rotation,
             'velocity': self.velocity,
-            'ego_dist': self.ego_dist,
+            'ego_translation': self.ego_translation,
             'num_pts': self.num_pts,
             'tracking_id': self.tracking_id,
             'tracking_name': self.tracking_name,
@@ -308,7 +317,8 @@ def deserialize(cls, content: dict):
                    size=tuple(content['size']),
                    rotation=tuple(content['rotation']),
                    velocity=tuple(content['velocity']),
-                   ego_dist=0.0 if 'ego_dist' not in content else float(content['ego_dist']),
+                   ego_translation=(0.0, 0.0, 0.0) if 'ego_translation' not in content
+                   else tuple(content['ego_translation']),
                    num_pts=-1 if 'num_pts' not in content else int(content['num_pts']),
                    tracking_id=content['tracking_id'],
                    tracking_name=content['tracking_name'],

python-sdk/nuscenes/eval/tracking/evaluate.py

@@ -5,7 +5,7 @@
 import json
 import os
 import time
-from typing import Tuple
+from typing import Tuple, List
 
 import numpy as np
 
@@ -43,7 +43,8 @@ def __init__(self,
                  output_dir: str,
                  nusc_version: str,
                  nusc_dataroot: str,
-                 verbose: bool = True):
+                 verbose: bool = True,
+                 render_classes: List[str] = None):
         """
         Initialize a TrackingEval object.
         :param config: A TrackingConfig object.
@@ -53,12 +54,14 @@ def __init__(self,
         :param nusc_version: The version of the NuScenes dataset.
         :param nusc_dataroot: Path of the nuScenes dataset on disk.
         :param verbose: Whether to print to stdout.
+        :param render_classes: Classes to render to disk or None.
         """
         self.cfg = config
         self.result_path = result_path
         self.eval_set = eval_set
         self.output_dir = output_dir
         self.verbose = verbose
+        self.render_classes = render_classes
 
         # Check result file exists.
         assert os.path.exists(result_path), 'Error: The result file does not exist!'
@@ -122,7 +125,9 @@ def accumulate_class(curr_class_name):
                                          self.cfg.dist_th_tp, self.cfg.min_recall,
                                          num_thresholds=TrackingMetricData.nelem,
                                          metric_worst=self.cfg.metric_worst,
-                                         verbose=self.verbose)
+                                         verbose=self.verbose,
+                                         output_dir=self.output_dir,
+                                         render_classes=self.render_classes)
             curr_md = curr_ev.accumulate()
             metric_data_list.set(curr_class_name, curr_md)
 
@@ -241,6 +246,8 @@ def main(self, render_curves: bool = True) -> TrackingMetrics:
                         help='Whether to render statistic curves to disk.')
     parser.add_argument('--verbose', type=int, default=1,
                         help='Whether to print to stdout.')
+    parser.add_argument('--render_classes', type=str, default='', nargs='+',
+                        help='For which classes we render tracking results to disk.')
     args = parser.parse_args()
 
     result_path_ = os.path.expanduser(args.result_path)
@@ -251,6 +258,7 @@ def main(self, render_curves: bool = True) -> TrackingMetrics:
     config_path = args.config_path
     render_curves_ = bool(args.render_curves)
     verbose_ = bool(args.verbose)
+    render_classes_ = args.render_classes
 
     if config_path == '':
         cfg_ = config_factory('tracking_nips_2019')
@@ -259,5 +267,6 @@ def main(self, render_curves: bool = True) -> TrackingMetrics:
             cfg_ = TrackingConfig.deserialize(json.load(_f))
 
     nusc_eval = TrackingEval(config=cfg_, result_path=result_path_, eval_set=eval_set_, output_dir=output_dir_,
-                             nusc_version=version_, nusc_dataroot=dataroot_, verbose=verbose_)
+                             nusc_version=version_, nusc_dataroot=dataroot_, verbose=verbose_,
+                             render_classes=render_classes_)
     nusc_eval.main(render_curves=render_curves_)

python-sdk/nuscenes/eval/tracking/loaders.py

@@ -46,7 +46,8 @@ def interp_float(left, right, rratio):
                        size=interp_list(left_box.size, right_box.size, right_ratio),
                        rotation=rotation,
                        velocity=interp_list(left_box.velocity, right_box.velocity, right_ratio),
-                       ego_dist=interp_float(left_box.ego_dist, right_box.ego_dist, right_ratio),  # May be inaccurate.
+                       ego_translation=interp_list(left_box.ego_translation, right_box.ego_translation,
+                                                   right_ratio),  # May be inaccurate.
                        tracking_id=right_box.tracking_id,
                        tracking_name=right_box.tracking_name,
                        tracking_score=tracking_score)

python-sdk/nuscenes/eval/tracking/mot.py

@@ -19,7 +19,6 @@ class MOTAccumulatorCustom(motmetrics.mot.MOTAccumulator):
     def __init__(self):
         super().__init__()
 
-
     @staticmethod
     def new_event_dataframe_with_data(indices, events):
         """