def get_label_data_2D_bounding_boxes(nusc, sample_token, sensor_channels):
"""
Create 2D bounding box labels from the given sample token.
1 bounding box vector contains:
- [0]: box dimensions
- box x_min (normalized to the image size)
- box y_min (normalized to the image size)
- box x_max (normalized to the image size)
- box y_max (normalized to the image size)
- [1]: class_category_name
- [2]: class_index
:param sample: the sample to get the annotation for
:param sensor_channels: list of channels for cropping the labels, e.g. ['CAM_FRONT', 'RADAR_FRONT']
This works only for CAMERA atm
:returns: [(nx4 np.array, list of str, list of int)] Labels (boxes, class_names, class_indices) (used for training)
"""
assert not any([s for s in sensor_channels if 'RADAR' in s]), "RADAR is not supported atm"
sample = nusc.get('sample', sample_token)
box_labels = [] # initialize counter for each category
class_names = []
class_indices = []
category_indices = get_category_index_mapping(nusc)
# Camera parameters
for selected_sensor_channel in sensor_channels:
sd_rec = nusc.get('sample_data', sample['data'][selected_sensor_channel])
# sensor = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token'])
# Create Boxes:
data_path, boxes, camera_intrinsic = nusc.get_sample_data(sd_rec['token'], box_vis_level=BoxVisibility.ANY)
imsize = (sd_rec['width'], sd_rec['height'])
# Add labels to boxes into
# assign_box_labels(nusc, boxes)
# Create labels for all boxes that are visible
for box in boxes:
box.label = category_indices[box.name]
if box_in_image(box=box, intrinsic=camera_intrinsic, imsize=imsize, vis_level=BoxVisibility.ANY):
# Check if box is visible
# If visible, we create the corresponding label
box2d = box.box2d(camera_intrinsic, imsize=imsize, normalize=True)
box_labels.append(box2d)
class_names.append(box.name)
class_indices.append(box.label)
return box_labels, class_names, class_indices
def get_sample_data(self, sample_data_token, box_vis_level=BoxVisibility.IN_FRONT, selected_anntokens=None) -> \
Tuple[str, List[Box], np.array]:
"""
Returns the data path as well as all annotations related to that sample_data.
:param sample_data_token: <str>. Sample_data token.
:param box_vis_level: <BoxVisibility>. If sample_data is an image, this sets required visibility for boxes.
:param selected_anntokens: [<str>]. If provided only return the selected annotation.
:return: (data_path <str>, boxes [<Box>], camera_intrinsic <np.array: 3, 3>)
"""
# Retrieve sensor & pose records
sd_record = self.get('sample_data', sample_data_token)
cs_record = self.get('calibrated_sensor', sd_record['calibrated_sensor_token'])
sensor_record = self.get('sensor', cs_record['sensor_token'])
pose_record = self.get('ego_pose', sd_record['ego_pose_token'])
data_path = self.get_sample_data_path(sample_data_token)
if sensor_record['modality'] == 'camera':
cam_intrinsic = np.array(cs_record['camera_intrinsic'])
imsize = (sd_record['width'], sd_record['height'])
else:
cam_intrinsic = None
imsize = None
# Retrieve all sample annotations and map to sensor coordinate system.
if selected_anntokens is not None:
boxes = list(map(self.get_box, selected_anntokens))
else:
boxes = self.get_boxes(sample_data_token)
# Make list of Box objects including coord system transforms.
box_list = []
for box in boxes:
# Move box to ego vehicle coord system
translation_vector = -np.array(pose_record['translation'])
box.translate(translation_vector)
rotation_matrix_quaternion = Quaternion(pose_record['rotation']).inverse
rotation_matrix = rotation_matrix_quaternion.rotation_matrix
box.rotate(rotation_matrix_quaternion)
# Move box to sensor coord system
translation_vector = -np.array(cs_record['translation'])
box.translate(translation_vector)
rotation_matrix_quaternion = Quaternion(cs_record['rotation']).inverse
rotation_matrix = rotation_matrix_quaternion.rotation_matrix
box.rotate(rotation_matrix_quaternion)
if sensor_record['modality'] == 'camera' and not \
box_in_image(box, cam_intrinsic, imsize, vis_level=box_vis_level):
continue
box_list.append(box)
return data_path, box_list, cam_intrinsic