Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Automatically determine which lidarseg classes are present in a pointcloud projected onto an image #410

Merged
merged 4 commits into from
Jun 8, 2020
Merged
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions python-sdk/nuscenes/lidarseg/lidarseg_utils.py
Original file line number Diff line number Diff line change
@@ -15,7 +15,7 @@ def get_stats(points_label: np.array, num_classes: int) -> List[int]:
Get frequency of each label in a point cloud.
:param num_classes: The number of classes.
:param points_label: A numPy array which contains the labels of the point cloud; e.g. np.array([2, 1, 34, ..., 38])
:returns: An array which contains the counts of each label in the point cloud. The index of the point cloud
:return: An array which contains the counts of each label in the point cloud. The index of the point cloud
corresponds to the index of the class label. E.g. [0, 2345, 12, 451] means that there are no points in
class 0, there are 2345 points in class 1, there are 12 points in class 2 etc.
"""
@@ -148,7 +148,7 @@ def filter_colormap(colormap: np.array, classes_to_display: np.array) -> np.ndar
of the labels to be display set to 1.0 and those to be hidden set to 0.0
:param colormap: [n x 3] array where each row consist of the RGB values for the corresponding class index
:param classes_to_display: An array of classes to display (e.g. [1, 8, 32]). The array need not be ordered.
:return (colormap <np.float: n, 4)>).
:return: (colormap <np.float: n, 4)>).

colormap = np.array([[R1, G1, B1], colormap = np.array([[1.0, 1.0, 1.0, 0.0],
[R2, G2, B2], ------> [R2, G2, B2, 1.0],
@@ -165,3 +165,35 @@ def filter_colormap(colormap: np.array, classes_to_display: np.array) -> np.ndar
colormap = np.concatenate((colormap, alpha.T), axis=1)

return colormap


def get_labels_in_coloring(color_legend: np.ndarray, coloring: np.ndarray) -> List[int]:
"""
Find the class labels which are present in a pointcloud which has been projected onto an image.
:param color_legend: A list of arrays in which each array corresponds to the RGB values of a class.
:param coloring: A list of arrays in which each array corresponds to the RGB values of a point in the portion of
the pointcloud projected onto the image.
:return: List of class indices which are present in the image.
"""

def _array_in_list(arr: List, list_arrays: List) -> bool:
"""
Check if an array is in a list of arrays.
:param: arr: An array.
:param: list_arrays: A list of arrays.
:return: Whether the given array is in the list of arrays.
"""
# Credits: https://stackoverflow.com/questions/23979146/check-if-numpy-array-is-in-list-of-numpy-arrays
return next((True for elem in list_arrays if np.array_equal(elem, arr)), False)

filter_lidarseg_labels = []

# Get only the distinct colors present in the pointcloud so that we will not need to compare each color in
# the color legend with every single point in the pointcloud later.
distinct_colors = list(set(tuple(c) for c in coloring))

for i, color in enumerate(color_legend):
if _array_in_list(color, distinct_colors):
filter_lidarseg_labels.append(i)

return filter_lidarseg_labels
21 changes: 16 additions & 5 deletions python-sdk/nuscenes/nuscenes.py
Original file line number Diff line number Diff line change
@@ -12,6 +12,7 @@

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

from nuscenes.lidarseg.lidarseg_utils import filter_colormap, get_arbitrary_colormap, plt_to_cv2, get_stats, \
get_key_from_value
get_key_from_value, get_labels_in_coloring
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
@@ -82,6 +83,11 @@ def __init__(self,
print('Loading nuScenes-lidarseg...')

self.lidarseg = self.__load_table__('lidarseg')
num_lidarseg_recs = len(self.lidarseg)
num_bin_files = len([name for name in os.listdir(os.path.join(self.dataroot, 'lidarseg', self.version))
if name.endswith('.bin')])
assert num_lidarseg_recs == num_bin_files, \
'Error: There are {} .bin files but {} lidarseg records.'.format(num_bin_files, num_lidarseg_recs)
self.table_names.append('lidarseg')

lidarseg_categories = self.__load_table__('category_lidarseg')
@@ -888,15 +894,20 @@ def render_pointcloud_in_image(self,
ax.axis('off')

# Produce a legend with the unique colors from the scatter.
if pointsensor_channel == 'LIDAR_TOP' and show_lidarseg_legend:
import matplotlib.patches as mpatches
if pointsensor_channel == 'LIDAR_TOP' and show_lidarseg_labels and show_lidarseg_legend:
recs = []
classes_final = []
classes = [name for idx, name in sorted(self.nusc.lidarseg_idx2name_mapping.items())]
color_legend = get_arbitrary_colormap(len(classes))

# 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, coloring)

for i in range(len(classes)):
# Create legend only for labels which the user wants to see,
# if the user has specified a lidarseg filter.
# 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]))
classes_final.append(classes[i])