|
import numpy as np
|
|
import torch
|
|
from nuplan.common.actor_state.tracked_objects_types import (
|
|
TrackedObjectType,
|
|
)
|
|
|
|
OBJECT_TYPE_DICT = {
|
|
"vehicle": TrackedObjectType.VEHICLE,
|
|
"pedestrian": TrackedObjectType.PEDESTRIAN,
|
|
"bicycle": TrackedObjectType.BICYCLE,
|
|
"traffic_cone": TrackedObjectType.TRAFFIC_CONE,
|
|
"barrier": TrackedObjectType.BARRIER,
|
|
"czone_sign": TrackedObjectType.CZONE_SIGN,
|
|
"generic_object": TrackedObjectType.GENERIC_OBJECT,
|
|
}
|
|
|
|
|
|
def limit_period(val, offset=0.5, period=2 * np.pi):
|
|
"""Limit the value into a period for periodic function.
|
|
|
|
Args:
|
|
val (torch.Tensor | np.ndarray): The value to be converted.
|
|
offset (float, optional): Offset to set the value range.
|
|
Defaults to 0.5.
|
|
period ([type], optional): Period of the value. Defaults to np.pi.
|
|
|
|
Returns:
|
|
(torch.Tensor | np.ndarray): Value in the range of
|
|
[-offset * period, (1-offset) * period]
|
|
"""
|
|
limited_val = val - torch.floor(val / period + offset) * period
|
|
return limited_val
|
|
|
|
|
|
class LiDARAug(object):
|
|
def __init__(self,
|
|
bda_aug_conf, is_train,
|
|
x_range='(-50.0, 50.0)',
|
|
y_range='(-50.0, 50.0)',
|
|
z_range='(-5, 20)',
|
|
):
|
|
for k in ['rot_lim', 'scale_lim', 'tran_lim']:
|
|
bda_aug_conf[k] = eval(bda_aug_conf[k])
|
|
self.bda_aug_conf = bda_aug_conf
|
|
self.is_train = False
|
|
self.x_range = eval(x_range)
|
|
self.y_range = eval(y_range)
|
|
self.z_range = eval(z_range)
|
|
|
|
def sample_bda_augmentation(self):
|
|
"""Generate bda augmentation values based on bda_config."""
|
|
if self.is_train:
|
|
rotate_bda = np.random.uniform(*self.bda_aug_conf['rot_lim'])
|
|
scale_bda = np.random.uniform(*self.bda_aug_conf['scale_lim'])
|
|
flip_dx = np.random.uniform() < self.bda_aug_conf['flip_dx_ratio']
|
|
flip_dy = np.random.uniform() < self.bda_aug_conf['flip_dy_ratio']
|
|
translation_std = self.bda_aug_conf.get('tran_lim', [0.0, 0.0, 0.0])
|
|
tran_bda = np.random.normal(scale=translation_std, size=3).T
|
|
else:
|
|
rotate_bda = 0
|
|
scale_bda = 1.0
|
|
flip_dx = False
|
|
flip_dy = False
|
|
tran_bda = np.zeros((1, 3), dtype=np.float32)
|
|
return rotate_bda, scale_bda, flip_dx, flip_dy, tran_bda
|
|
|
|
def bev_transform(self, gt_boxes, rotate_angle, scale_ratio, flip_dx,
|
|
flip_dy, tran_bda, rot_mat):
|
|
if gt_boxes.shape[0] > 0:
|
|
gt_boxes[:, :3] = (
|
|
rot_mat @ gt_boxes[:, :3].unsqueeze(-1)).squeeze(-1)
|
|
gt_boxes[:, 3:6] *= scale_ratio
|
|
gt_boxes[:, 6] += rotate_angle
|
|
if flip_dx:
|
|
gt_boxes[:,
|
|
6] = 2 * torch.asin(torch.tensor(1.0)) - gt_boxes[:,
|
|
6]
|
|
if flip_dy:
|
|
gt_boxes[:, 6] = -gt_boxes[:, 6]
|
|
gt_boxes[:, 7:] = (
|
|
rot_mat[:2, :2] @ gt_boxes[:, 7:].unsqueeze(-1)).squeeze(-1)
|
|
gt_boxes[:, :3] = gt_boxes[:, :3] + tran_bda
|
|
return gt_boxes
|
|
|
|
def __call__(self, features, targets):
|
|
|
|
|
|
if 'dets' in targets and 'labels' in targets:
|
|
boxes = targets['dets']
|
|
labels = targets['labels']
|
|
|
|
for t, (box, label) in enumerate(zip(boxes, labels)):
|
|
label_mask = np.array([n in OBJECT_TYPE_DICT for n in label], dtype=np.bool_)
|
|
label_mask = torch.from_numpy(label_mask)
|
|
range_mask = ((box[:, 0] > self.x_range[0]) &
|
|
(box[:, 0] < self.x_range[1]) &
|
|
(box[:, 1] > self.y_range[0]) &
|
|
(box[:, 1] < self.y_range[1]))
|
|
mask = range_mask & label_mask
|
|
box_of_interest = box[mask]
|
|
box_of_interest[:, 6] = limit_period(box_of_interest[:, 6])
|
|
boxes[t] = box_of_interest.float()
|
|
|
|
labels[t] = torch.from_numpy(np.array([OBJECT_TYPE_DICT[x].value for
|
|
x in label], dtype=np.int64))[mask]
|
|
targets['dets'] = boxes
|
|
targets['labels'] = labels
|
|
|
|
rotate_bda, scale_bda, flip_dx, flip_dy, tran_bda = \
|
|
self.sample_bda_augmentation()
|
|
bda_mat = torch.zeros(4, 4)
|
|
bda_mat[3, 3] = 1
|
|
rotate_angle = torch.tensor(rotate_bda / 180 * np.pi)
|
|
rot_sin = torch.sin(rotate_angle)
|
|
rot_cos = torch.cos(rotate_angle)
|
|
rot_mat = torch.Tensor([[rot_cos, -rot_sin, 0], [rot_sin, rot_cos, 0],
|
|
[0, 0, 1]])
|
|
scale_mat = torch.Tensor([[scale_bda, 0, 0], [0, scale_bda, 0],
|
|
[0, 0, scale_bda]])
|
|
flip_mat = torch.Tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
|
|
if flip_dx:
|
|
flip_mat = flip_mat @ torch.Tensor([[-1, 0, 0], [0, 1, 0],
|
|
[0, 0, 1]])
|
|
if flip_dy:
|
|
flip_mat = flip_mat @ torch.Tensor([[1, 0, 0], [0, -1, 0],
|
|
[0, 0, 1]])
|
|
bda_rot = flip_mat @ (scale_mat @ rot_mat)
|
|
|
|
if 'dets' in targets:
|
|
for idx, boxes in enumerate(targets['dets']):
|
|
targets['dets'][idx] = self.bev_transform(boxes, rotate_bda, scale_bda,
|
|
flip_dx, flip_dy, tran_bda, bda_rot)
|
|
|
|
|
|
|
|
|
|
|
|
for idx, points in enumerate(features['lidars_warped']):
|
|
points_aug = (bda_rot @ points[:, :3].unsqueeze(-1)).squeeze(-1)
|
|
points[:, :3] = points_aug + tran_bda
|
|
features['lidars_warped'][idx] = points
|
|
|
|
|
|
|
|
|
|
|
|
|
|
bda_mat[:3, :3] = bda_rot
|
|
bda_mat[:3, 3] = torch.from_numpy(tran_bda)
|
|
features['bda'] = bda_mat
|
|
return features, targets
|
|
|
