render_annotation improved (nutonomy#207)

* render_annotation improved

render extra information about the annotation below the CAMERA view

* Minor cleanups

python-sdk/nuscenes/nuscenes.py

@@ -14,6 +14,7 @@
 import numpy as np
 import sklearn.metrics
 from PIL import Image
+from matplotlib import rcParams
 from matplotlib.axes import Axes
 from pyquaternion import Quaternion
 from tqdm import tqdm
@@ -398,12 +399,14 @@ def render_sample_data(self, sample_data_token: str, with_anns: bool = True,
                                          out_path=out_path)
 
     def render_annotation(self, sample_annotation_token: str, margin: float = 10, view: np.ndarray = np.eye(4),
-                          box_vis_level: BoxVisibility = BoxVisibility.ANY, out_path: str = None) -> None:
-        self.explorer.render_annotation(sample_annotation_token, margin, view, box_vis_level, out_path)
+                          box_vis_level: BoxVisibility = BoxVisibility.ANY, out_path: str = None,
+                          extra_info: bool = False) -> None:
+        self.explorer.render_annotation(sample_annotation_token, margin, view, box_vis_level, out_path, extra_info)
 
     def render_instance(self, instance_token: str, margin: float = 10, view: np.ndarray = np.eye(4),
-                        box_vis_level: BoxVisibility = BoxVisibility.ANY, out_path: str = None) -> None:
-        self.explorer.render_instance(instance_token, margin, view, box_vis_level, out_path)
+                        box_vis_level: BoxVisibility = BoxVisibility.ANY, out_path: str = None,
+                        extra_info: bool = False) -> None:
+        self.explorer.render_instance(instance_token, margin, view, box_vis_level, out_path, extra_info)
 
     def render_scene(self, scene_token: str, freq: float = 10, imsize: Tuple[float, float] = (640, 360),
                      out_path: str = None) -> None:
@@ -810,16 +813,17 @@ def render_annotation(self,
                           margin: float = 10,
                           view: np.ndarray = np.eye(4),
                           box_vis_level: BoxVisibility = BoxVisibility.ANY,
-                          out_path: str = None) -> None:
+                          out_path: str = None,
+                          extra_info: bool = False) -> None:
         """
         Render selected annotation.
         :param anntoken: Sample_annotation token.
         :param margin: How many meters in each direction to include in LIDAR view.
         :param view: LIDAR view point.
         :param box_vis_level: If sample_data is an image, this sets required visibility for boxes.
         :param out_path: Optional path to save the rendered figure to disk.
+        :param extra_info: Whether to render extra information below camera view.
         """
-
         ann_record = self.nusc.get('sample_annotation', anntoken)
         sample_record = self.nusc.get('sample', ann_record['sample_token'])
         assert 'LIDAR_TOP' in sample_record['data'].keys(), 'No LIDAR_TOP in data, cant render'
@@ -863,6 +867,32 @@ def render_annotation(self,
             c = np.array(self.get_color(box.name)) / 255.0
             box.render(axes[1], view=camera_intrinsic, normalize=True, colors=(c, c, c))
 
+        # Print extra information about the annotation below the camera view.
+        if extra_info:
+            rcParams['font.family'] = 'monospace'
+
+            w, l, h = ann_record['size']
+            category = ann_record['category_name']
+            lidar_points = ann_record['num_lidar_pts']
+            radar_points = ann_record['num_radar_pts']
+
+            sample_data_record = self.nusc.get('sample_data', sample_record['data']['LIDAR_TOP'])
+            pose_record = self.nusc.get('ego_pose', sample_data_record['ego_pose_token'])
+            dist = np.linalg.norm(np.array(pose_record['translation']) - np.array(ann_record['translation']))
+
+            information = ' \n'.join(['category: {}'.format(category),
+                                      '',
+                                      '# lidar points: {0:>4}'.format(lidar_points),
+                                      '# radar points: {0:>4}'.format(radar_points),
+                                      '',
+                                      'distance: {:>7.3f}m'.format(dist),
+                                      '',
+                                      'width:  {:>7.3f}m'.format(w),
+                                      'length: {:>7.3f}m'.format(l),
+                                      'height: {:>7.3f}m'.format(h)])
+
+            plt.annotate(information, (0, 0), (0, -20), xycoords='axes fraction', textcoords='offset points', va='top')
+
         if out_path is not None:
             plt.savefig(out_path)
 
@@ -871,14 +901,16 @@ def render_instance(self,
                         margin: float = 10,
                         view: np.ndarray = np.eye(4),
                         box_vis_level: BoxVisibility = BoxVisibility.ANY,
-                        out_path: str = None) -> None:
+                        out_path: str = None,
+                        extra_info: bool = False) -> None:
         """
         Finds the annotation of the given instance that is closest to the vehicle, and then renders it.
         :param instance_token: The instance token.
         :param margin: How many meters in each direction to include in LIDAR view.
         :param view: LIDAR view point.
         :param box_vis_level: If sample_data is an image, this sets required visibility for boxes.
         :param out_path: Optional path to save the rendered figure to disk.
+        :param extra_info: Whether to render extra information below camera view.
         """
         ann_tokens = self.nusc.field2token('sample_annotation', 'instance_token', instance_token)
         closest = [np.inf, None]
@@ -891,7 +923,8 @@ def render_instance(self,
             if dist < closest[0]:
                 closest[0] = dist
                 closest[1] = ann_token
-        self.render_annotation(closest[1], margin, view, box_vis_level, out_path)
+
+        self.render_annotation(closest[1], margin, view, box_vis_level, out_path, extra_info)
 
     def render_scene(self,
                      scene_token: str,