File size: 6,766 Bytes
da2e2ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
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):
# 1. filter box based on ranges
# 2. filter label based on classes
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)
# print('before bda')
# print(features['lidars_warped'][-1][:, 0].max())
# print(features['lidars_warped'][-1][:, 0].min())
# print(features['lidars_warped'][-1][:, 1].max())
# print(features['lidars_warped'][-1][:, 1].min())
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
# print('after bda')
# print(features['lidars_warped'][-1][:, 0].max())
# print(features['lidars_warped'][-1][:, 0].min())
# print(features['lidars_warped'][-1][:, 1].max())
# print(features['lidars_warped'][-1][:, 1].min())
bda_mat[:3, :3] = bda_rot
bda_mat[:3, 3] = torch.from_numpy(tran_bda)
features['bda'] = bda_mat
return features, targets
|