Skip to content

Commit

Permalink
Improve lidarseg rendering methods, update tutorial, add unit tests (#…
Browse files Browse the repository at this point in the history 
…429)

* Initial commit to address comments

* Fix bug where passing np.array instead of array into filter_lidarseg_labels throws an error

* Rephrase some parts of tutorial

* Include legend for render_sample_data

* Add arg for show_lidarseg_legend in render_sample_data

* Adjust aesthetics of legend in render_sample_data

* Update colormap

* Make creating legend modular

* Ensure only labels in pc is included in legend

* Explain verbose

* Check that filter_lidarseg_labels is either a list or np.array

* Allow user to list stats by class index

* Update docstrings

* Update tutorial to demo sort_by

* Add unit tests

* Make painting of labels modular

* Add docstrings
  • Loading branch information
@whyekit-motional
whyekit-motional authored Jul 8, 2020
1 parent bea0b34 commit 26252db
Show file tree
Hide file tree
Showing 5 changed files with 262 additions and 70 deletions.
79 changes: 78 additions & 1 deletion python-sdk/nuscenes/lidarseg/lidarseg_utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
# nuScenes dev-kit.
# Code written by Fong Whye Kit, 2020.

import colorsys
from typing import Dict, Iterable, List, Tuple

import cv2
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
Expand Down Expand Up @@ -71,6 +71,9 @@ def plt_to_cv2(points: np.array, coloring: np.array, im, imsize: Tuple[int, int]
mat = cv2.cvtColor(mat, cv2.COLOR_RGB2BGR)
mat = cv2.resize(mat, imsize)

# Clear off the current figure to prevent an accumulation of figures in memory.
plt.close('all')

return mat


Expand Down Expand Up @@ -149,3 +152,77 @@ def _array_in_list(arr: List, list_arrays: List) -> bool:
filter_lidarseg_labels.append(i)

return filter_lidarseg_labels


def create_lidarseg_legend(labels_to_include_in_legend: List[int],
idx2name: Dict[int, str], name2color: Dict[str, List[int]],
loc: str = 'upper center', ncol: int = 3, bbox_to_anchor: Tuple = None):
"""
Given a list of class indices, the mapping from class index to class name, and the mapping from class name
to class color, produce a legend which shows the color and the corresponding class name.
:param labels_to_include_in_legend: Labels to show in the legend.
:param idx2name: The mapping from class index to class name.
:param name2color: The mapping from class name to class color.
:param loc: The location of the legend.
:param ncol: The number of columns that the legend has.
:param bbox_to_anchor: A 2-tuple (x, y) which places the top-left corner of the legend specified by loc
at x, y. The origin is at the bottom-left corner and x and y are normalized between
0 and 1 (i.e. x > 1 and / or y > 1 will place the legend outside the plot.
"""

recs = []
classes_final = []
classes = [name for idx, name in sorted(idx2name.items())]

for i in range(len(classes)):
if labels_to_include_in_legend is None or i in labels_to_include_in_legend:
name = classes[i]
recs.append(mpatches.Rectangle((0, 0), 1, 1, fc=np.array(name2color[name]) / 255))

# Truncate class names to only first 25 chars so that legend is not excessively long.
classes_final.append(classes[i][:25])

plt.legend(recs, classes_final, loc=loc, ncol=ncol, bbox_to_anchor=bbox_to_anchor)


def paint_points_label(lidarseg_labels_filename: str, filter_lidarseg_labels: List[int],
name2idx: Dict[str, int], colormap: Dict[str, List[int]]) -> np.ndarray:
"""
Paint each label in a pointcloud with the corresponding RGB value, and set the opacity of the labels to
be shown to 1 (the opacity of the rest will be set to 0); e.g.:
[30, 5, 12, 34, ...] ------> [[R30, G30, B30, 0], [R5, G5, B5, 1], [R34, G34, B34, 1], ...]
:param lidarseg_labels_filename: Path to the .bin file containing the labels.
:param filter_lidarseg_labels: The labels for which to set opacity to zero; this is to hide those points
thereby preventing them from being displayed.
:param name2idx: A dictionary containing the mapping from class names to class indices.
:param colormap: A dictionary containing the mapping from class names to RGB values.
:return: A numpy array which has length equal to the number of points in the pointcloud, and each value is
a RGBA array.
"""

# Load labels from .bin file.
points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8) # [num_points]

# Given a colormap (class name -> RGB color) and a mapping from class name to class index,
# get an array of RGB values where each color sits at the index in the array corresponding
# to the class index.
colors = colormap_to_colors(colormap, name2idx) # Shape: [num_class, 3]

if filter_lidarseg_labels is not None:
# Ensure that filter_lidarseg_labels is an iterable.
assert isinstance(filter_lidarseg_labels, (list, np.ndarray)), \
'Error: filter_lidarseg_labels should be a list of class indices, eg. [9], [10, 21].'

# Check that class indices in filter_lidarseg_labels are valid.
assert all([0 <= x < len(name2idx) for x in filter_lidarseg_labels]), \
'All class indices in filter_lidarseg_labels should ' \
'be between 0 and {}'.format(len(name2idx) - 1)

# Filter to get only the colors of the desired classes; this is done by setting the
# alpha channel of the classes to be viewed to 1, and the rest to 0.
colors = filter_colors(colors, filter_lidarseg_labels) # Shape: [num_class, 4]

# Paint each label with its respective RGBA value.
coloring = colors[points_label] # Shape: [num_points, 4]

return coloring
109 changes: 60 additions & 49 deletions python-sdk/nuscenes/nuscenes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@

import cv2
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import sklearn.metrics
from PIL import Image
Expand All @@ -22,7 +21,7 @@
from tqdm import tqdm

from nuscenes.lidarseg.lidarseg_utils import filter_colors, colormap_to_colors, plt_to_cv2, get_stats, \
get_key_from_value, get_labels_in_coloring
get_key_from_value, get_labels_in_coloring, create_lidarseg_legend, paint_points_label
from nuscenes.utils.data_classes import LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix
from nuscenes.utils.map_mask import MapMask
Expand Down Expand Up @@ -409,19 +408,23 @@ def box_velocity(self, sample_annotation_token: str, max_time_diff: float = 1.5)
else:
return pos_diff / time_diff

def get_sample_lidarseg_stats(self, sample_token: str, sort_counts: bool = True,
def get_sample_lidarseg_stats(self, sample_token: str, sort_by: str = 'count',
lidarseg_preds_bin_path: str = None) -> None:
"""
Print the number of points for each class in the lidar pointcloud of a sample. Classes with have no
points in the pointcloud will not be printed.
:param sample_token: Sample token.
:param sort_counts: If True, the stats will be printed in ascending order of frequency; if False,
the stats will be printed alphabetically according to class name.
:param sort_by: One of three options: count / name / index. If 'count`, the stats will be printed in
ascending order of frequency; if `name`, the stats will be printed alphabetically
according to class name; if `index`, the stats will be printed in ascending order of
class index.
:param lidarseg_preds_bin_path: A path to the .bin file which contains the user's lidar segmentation
predictions for the sample.
"""
assert hasattr(self, 'lidarseg'), 'Error: You have no lidarseg data; unable to get ' \
'statistics for segmentation of the point cloud.'
assert sort_by in ['count', 'name', 'index'], 'Error: sort_by can only be one of the following: ' \
'count / name / index.'

sample_rec = self.get('sample', sample_token)
ref_sd_token = sample_rec['data']['LIDAR_TOP']
Expand Down Expand Up @@ -457,23 +460,25 @@ def get_sample_lidarseg_stats(self, sample_token: str, sort_counts: bool = True,
for i in range(len(lidarseg_counts)):
lidarseg_counts_dict[self.lidarseg_idx2name_mapping[i]] = lidarseg_counts[i]

if sort_counts:
if sort_by == 'count':
out = sorted(lidarseg_counts_dict.items(), key=lambda item: item[1])
else:
elif sort_by == 'name':
out = sorted(lidarseg_counts_dict.items())
else:
out = lidarseg_counts_dict.items()

for class_name, count in out:
if count > 0:
idx = get_key_from_value(self.lidarseg_idx2name_mapping, class_name)
print('{:3} {:35} n={:12,}'.format(idx, class_name, count))
print('{:3} {:40} n={:12,}'.format(idx, class_name, count))

print('=' * len(header))

def list_categories(self) -> None:
self.explorer.list_categories()

def list_lidarseg_categories(self) -> None:
self.explorer.list_lidarseg_categories()
def list_lidarseg_categories(self, sort_by: str = 'count') -> None:
self.explorer.list_lidarseg_categories(sort_by=sort_by)

def list_attributes(self) -> None:
self.explorer.list_attributes()
Expand Down Expand Up @@ -515,12 +520,14 @@ def render_sample_data(self, sample_data_token: str, with_anns: bool = True,
nsweeps: int = 1, out_path: str = None, underlay_map: bool = True,
use_flat_vehicle_coordinates: bool = True,
show_lidarseg: bool = False,
show_lidarseg_legend: bool = False,
filter_lidarseg_labels: List = None,
lidarseg_preds_bin_path: str = None, verbose: bool = True) -> None:
self.explorer.render_sample_data(sample_data_token, with_anns, box_vis_level, axes_limit, ax, nsweeps=nsweeps,
out_path=out_path, underlay_map=underlay_map,
use_flat_vehicle_coordinates=use_flat_vehicle_coordinates,
show_lidarseg=show_lidarseg,
show_lidarseg_legend=show_lidarseg_legend,
filter_lidarseg_labels=filter_lidarseg_labels,
lidarseg_preds_bin_path=lidarseg_preds_bin_path, verbose=verbose)

Expand Down Expand Up @@ -624,12 +631,18 @@ def list_categories(self) -> None:
np.mean(stats[:, 2]), np.std(stats[:, 2]),
np.mean(stats[:, 3]), np.std(stats[:, 3])))

def list_lidarseg_categories(self) -> None:
def list_lidarseg_categories(self, sort_by: str = 'count') -> None:
"""
Print categories and counts of the lidarseg data. These stats only cover
the split specified in nusc.version.
:param sort_by: One of three options: count / name / index. If 'count`, the stats will be printed in
ascending order of frequency; if `name`, the stats will be printed alphabetically
according to class name; if `index`, the stats will be printed in ascending order of
class index.
"""
assert hasattr(self.nusc, 'lidarseg'), 'Error: nuScenes-lidarseg not installed!'
assert sort_by in ['count', 'name', 'index'], 'Error: sort_by can only be one of the following: ' \
'count / name / index.'

print('Calculating stats for nuScenes-lidarseg...')
start_time = time.time()
Expand All @@ -650,11 +663,17 @@ def list_lidarseg_categories(self) -> None:
for i in range(len(lidarseg_counts)):
lidarseg_counts_dict[self.nusc.lidarseg_idx2name_mapping[i]] = lidarseg_counts[i]

out = sorted(lidarseg_counts_dict.items(), key=lambda item: item[1])
if sort_by == 'count':
out = sorted(lidarseg_counts_dict.items(), key=lambda item: item[1])
elif sort_by == 'name':
out = sorted(lidarseg_counts_dict.items())
else:
out = lidarseg_counts_dict.items()

# Print frequency counts of each class in the lidarseg dataset.
for class_name, count in out:
idx = get_key_from_value(self.nusc.lidarseg_idx2name_mapping, class_name)
print('{:3} {:35} nbr_points={:12,}'.format(idx, class_name, count))
print('{:3} {:40} nbr_points={:12,}'.format(idx, class_name, count))

print('Calculated stats for {} point clouds in {:.1f} seconds.\n====='.format(
len(self.nusc.lidarseg), time.time() - start_time))
Expand Down Expand Up @@ -809,16 +828,9 @@ def map_pointcloud_to_image(self,
lidarseg_labels_filename = None

if lidarseg_labels_filename:
points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)

# A scatter plot is used for displaying the lidarseg points; however, the scatter plot takes in colors
# as an array of RGB values, and thus the colormap needs to be converted to the appropriate format for
# later use.
colors = colormap_to_colors(self.nusc.colormap, self.nusc.lidarseg_name2idx_mapping)

if filter_lidarseg_labels:
colors = filter_colors(colors, filter_lidarseg_labels)
coloring = colors[points_label]
# Paint each label in the pointcloud with a RGBA value.
coloring = paint_points_label(lidarseg_labels_filename, filter_lidarseg_labels,
self.nusc.lidarseg_name2idx_mapping, self.nusc.colormap)
else:
coloring = depths
print('Warning: There are no lidarseg labels in {}. Points will be colored according to distance '
Expand Down Expand Up @@ -901,13 +913,8 @@ def render_pointcloud_in_image(self,

# Produce a legend with the unique colors from the scatter.
if pointsensor_channel == 'LIDAR_TOP' and show_lidarseg and show_lidarseg_legend:
recs = []
classes_final = []
classes = [name for idx, name in sorted(self.nusc.lidarseg_idx2name_mapping.items())]

# A scatter plot is used for displaying the lidarseg points; however, the scatter plot takes in colors
# as an array of RGB values, and thus the colormap needs to be converted to the appropriate format for
# later use.
# Since the labels are stored as class indices, we get the RGB colors from the colormap in an array where
# the position of the RGB color corresponds to the index of the class it represents.
color_legend = colormap_to_colors(self.nusc.colormap, self.nusc.lidarseg_name2idx_mapping)

# If user does not specify a filter, then set the filter to contain the classes present in the pointcloud
Expand All @@ -916,14 +923,8 @@ def render_pointcloud_in_image(self,
if filter_lidarseg_labels is None:
filter_lidarseg_labels = get_labels_in_coloring(color_legend, coloring)

for i in range(len(classes)):
# Create legend only for labels specified in the lidarseg filter.
if filter_lidarseg_labels is None or i in filter_lidarseg_labels:
recs.append(mpatches.Rectangle((0, 0), 1, 1, fc=color_legend[i]))

# Truncate class names to only first 25 chars so that legend is not excessively long.
classes_final.append(classes[i][:25])
plt.legend(recs, classes_final, loc='upper center', ncol=3)
create_lidarseg_legend(filter_lidarseg_labels,
self.nusc.lidarseg_idx2name_mapping, self.nusc.colormap)

if out_path is not None:
plt.savefig(out_path, bbox_inches='tight', pad_inches=0, dpi=200)
Expand Down Expand Up @@ -1092,6 +1093,7 @@ def render_sample_data(self,
underlay_map: bool = True,
use_flat_vehicle_coordinates: bool = True,
show_lidarseg: bool = False,
show_lidarseg_legend: bool = False,
filter_lidarseg_labels: List = None,
lidarseg_preds_bin_path: str = None,
verbose: bool = True) -> None:
Expand All @@ -1111,6 +1113,7 @@ def render_sample_data(self,
setting is more correct and rotates the plot by ~90 degrees.
:param show_lidarseg: When set to True, the lidar data is colored with the segmentation labels. When set
to False, the colors of the lidar data represent the distance from the center of the ego vehicle.
:param show_lidarseg_legend: Whether to display the legend for the lidarseg labels in the frame.
:param filter_lidarseg_labels: Only show lidar points which belong to the given list of classes. If None
or the list is empty, all classes will be displayed.
:param lidarseg_preds_bin_path: A path to the .bin file which contains the user's lidar segmentation
Expand Down Expand Up @@ -1213,16 +1216,25 @@ def render_sample_data(self,
lidarseg_labels_filename = None

if lidarseg_labels_filename:
points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)

# A scatter plot is used for displaying the lidarseg points; however, the scatter plot takes
# in colors as an array of RGB values, and thus the colormap needs to be converted to the
# appropriate format.
coloring = colormap_to_colors(self.nusc.colormap, self.nusc.lidarseg_name2idx_mapping)

if filter_lidarseg_labels:
coloring = filter_colors(coloring, filter_lidarseg_labels)
colors = coloring[points_label]
# Paint each label in the pointcloud with a RGBA value.
colors = paint_points_label(lidarseg_labels_filename, filter_lidarseg_labels,
self.nusc.lidarseg_name2idx_mapping, self.nusc.colormap)

if show_lidarseg_legend:
# Since the labels are stored as class indices, we get the RGB colors from the colormap
# in an array where the position of the RGB color corresponds to the index of the class
# it represents.
color_legend = colormap_to_colors(self.nusc.colormap, self.nusc.lidarseg_name2idx_mapping)

# If user does not specify a filter, then set the filter to contain the classes present in
# the pointcloud after it has been projected onto the image; this will allow displaying the
# legend only for classes which are present in the image (instead of all the classes).
if filter_lidarseg_labels is None:
filter_lidarseg_labels = get_labels_in_coloring(color_legend, colors)

create_lidarseg_legend(filter_lidarseg_labels,
self.nusc.lidarseg_idx2name_mapping, self.nusc.colormap,
loc='upper left', ncol=1, bbox_to_anchor=(1.05, 1.0))
else:
colors = np.minimum(1, dists / axes_limit / np.sqrt(2))
print('Warning: There are no lidarseg labels in {}. Points will be colored according to distance '
Expand Down Expand Up @@ -1260,7 +1272,6 @@ def render_sample_data(self,
# Limit visible range.
ax.set_xlim(-axes_limit, axes_limit)
ax.set_ylim(-axes_limit, axes_limit)

elif sensor_modality == 'camera':
# Load boxes and image.
data_path, boxes, camera_intrinsic = self.nusc.get_sample_data(sample_data_token,
Expand Down
Loading

0 comments on commit 26252db

Please sign in to comment.