def lvx_vis_3d(gt_annos,dt_annos,output_dir):
root_path = dt_annos[0]["metadata"]["image_prefix"]
fig = mlab.figure(figure=None, bgcolor=(0.85,0.85,0.85), fgcolor=None, engine=None, size=(1600, 1000))
if gt_annos is not None:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
vot = cv2.VideoWriter(os.path.join(output_dir,"bev_video_3d.mp4"),fourcc,10,(1600,1000),True)
dataset = lvx_object(root_path)
gt_image_idxes = [str(info["token"]) for info in gt_annos]
for i in range(len(gt_annos)):
print(gt_annos[i]['token'])
points = dataset.get_lidar(gt_annos[i]['token'])
points = points[points[:,2]>0.1,:]
gt_annos[i] = remove_dontcare(gt_annos[i])
fig = draw_lidar(points,color=None,fig=fig, pts_mode = 'sphere')
gt_corners = center_to_corner_box3d(gt_annos[i]['location'],gt_annos[i]['dimensions'], gt_annos[i]['rotation_y'])
fig = draw_gt_boxes3d(gt_corners, gt_annos[i]['id'],color=(0,0,1),fig=fig)
dt_corners = center_to_corner_box3d(dt_annos[i]['location'],dt_annos[i]['dimensions'], dt_annos[i]['rotation_y'])
fig = draw_gt_boxes3d(dt_corners, dt_annos[i]['track_id'],color=(0,1,0),fig=fig)
img_path = os.path.join(output_dir,f"bev_imgs/lvx_{gt_annos[i]['token']}_3d.png")
mlab.savefig(img_path)
mlab.clf(figure=fig)
img = cv2.imread(img_path)
vot.write(img)
vot.release()
else:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
vot = cv2.VideoWriter(os.path.join(output_dir,"beicao_video_3d.mp4"),fourcc,10,(1600,1000),True)
dataset = lvx_object(root_path,split="testing")
for i in range(len(dt_annos)):
idx = dt_annos[i]['metadata']['token']
print(idx)
points = dataset.get_lidar(idx)
plane = points[points[:,2]<0.1,:]
points = points[points[:,2]>0.1,:]
fig = draw_lidar(points,color=None,fig=fig, pts_mode = 'sphere')
fig = draw_lidar(plane,color=None,fig=fig)
dt_corners = center_to_corner_box3d(dt_annos[i]['location'],dt_annos[i]['dimensions'], dt_annos[i]['rotation_y'])
color_list = np.zeros((dt_annos[i]['track_id'].shape[0],3))
for j in range(dt_annos[i]['track_id'].shape[0]):
np.random.seed(int(dt_annos[i]['track_id'][j]))
color_list[j] = np.random.rand(3)
view_setting = [180+30*np.sin(np.pi/100.*i),50+20*np.sin(np.pi/200.*i),60-20*np.sin(np.pi/100*i)]
print(view_setting)
fig = draw_gt_boxes3d(dt_corners, dt_annos[i]['track_id'],line_width=3,fig=fig,color_list=color_list,view_setting=view_setting)
img_path = os.path.join(output_dir,f"beicao_imgs/lvx_{idx}_3d.png")
mlab.savefig(img_path)
mlab.clf(figure=fig)
img = cv2.imread(img_path)
vot.write(img)
vot.release()
def get_track_scores(boxes, boxes_1):
''' DIoU'''
predict_pre = box3d_overlap(boxes, boxes_1, z_axis=2, z_center=0.5)
corners = center_to_corner_box3d(boxes[:, :3], boxes[:, 3:6], boxes[:, 6])
corners_1 = center_to_corner_box3d(boxes_1[:, :3], boxes_1[:, 3:6],
boxes_1[:, 6])
for i in range(boxes.shape[0]):
for j in range(boxes_1.shape[0]):
c = ((boxes[i, :3] - boxes_1[j, :3])**2).sum()
d = get_farest(corners[i], corners_1[j])
predict_pre[i, j] -= c / d
return predict_pre
def noise_per_object_v4_(gt_boxes, points=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0,\
global_random_rot_range=np.pi / 4, num_try=5, group_ids=None, data_aug_with_context=-1.0,\
data_aug_random_drop=-1.0):
"""
perform random rotation and translation on each groundtruth independently.
Args:
gt_boxes: [N, 7], gt box in lidar, [x, y, z, w, l, h, ry]
points: [M, 4], point cloud in lidar, all points in the scene.
rotation_perturb: [-0.785, 0.785],
center_noise_std: [1.0, 1.0, 0.5],
global_random_rot_range: [0, 0]
num_try: 100
"""
num_boxes = gt_boxes.shape[0]
center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype)
loc_noises = np.random.normal(scale=center_noise_std,
size=[num_boxes, num_try, 3])
rot_noises = np.random.uniform(rotation_perturb[0],
rotation_perturb[1],
size=[num_boxes, num_try])
origin = [0.5, 0.5, 0.5]
offset = [0.0, 0.0, 0.0, 0.0, 0.0]
if data_aug_with_context > 0: # to enlarge boxes and keep context
offset = [0.0, 0.0, data_aug_with_context, data_aug_with_context, 0.0]
# gt_boxes: x,y,z(lidar), w, l, h, ry(cam).
gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3],
gt_boxes[:, 3:6] +
offset[2:5],
gt_boxes[:, 6],
origin=origin,
axis=2)
# noise on bev_box (x, y, w, l, ry). # todo: seems without noise in z-axis velo, because only cal bev iou, but z-axis noise added later.
selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]] + offset,
valid_mask, loc_noises, rot_noises) #
loc_transforms = _select_transform(loc_noises, selected_noise)
rot_transforms = _select_transform(rot_noises, selected_noise)
surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners)
if points is not None:
point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces)
# todo: perform random drop here.
if data_aug_random_drop > 0.0:
point_masks = point_masks
points_transform_(
points,
gt_boxes[:, :3],
point_masks,
loc_transforms,
rot_transforms,
valid_mask,
)
box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def track_vis_all(dt_annos,output_dir):
fig = mlab.figure(figure=None, bgcolor=(0.85,0.85,0.85), fgcolor=None, engine=None, size=(3200, 2000))
# pcd = o3d.io.read_point_cloud("/Extra/zhangmeng/3d_detection/BBOX_x4_track/testing/Lidar/PC_300.pcd")
pcd = o3d.io.read_point_cloud("/home/zhangmeng/PC_2.pcd")
points = np.asarray(pcd.points)
plane = points[points[:,2]<0.1,:]
points = points[points[:,2]>0.1,:]
fig = draw_lidar(points,color=None,fig=fig)
fig = draw_lidar(plane,color=None,fig=fig)
print(dt_annos[0])
dt_corners = center_to_corner_box3d(dt_annos[0]['location'],dt_annos[0]['dimensions'], dt_annos[0]['rotation_y'])
color_list = np.zeros((dt_annos[0]['id'].shape[0],3))
pos = dt_annos[0]['location']
track = dt_annos[0]['id']
for j in range(dt_annos[0]['id'].shape[0]):
np.random.seed(int(dt_annos[0]['id'][j]))
color_list[j] = np.random.rand(3)
fig = draw_gt_boxes3d(dt_corners, dt_annos[0]['id'],draw_text=False,line_width=2,fig=fig,color_list=color_list,view_setting=(235,70,60))
for i in range(3,len(dt_annos),3):
for j in range(dt_annos[i]['id'].shape[0]):
ind = np.where(track == dt_annos[i]['id'][j])[0]
print(ind)
if ind.size:
mlab.plot3d([pos[ind[0]][0], dt_annos[i]['location'][j,0]],[pos[ind[0]][1], dt_annos[i]['location'][j,1]],[pos[ind[0]][2], dt_annos[i]['location'][j,2]], color=tuple(color_list[ind[0]]), line_width=7, tube_radius=None, figure=fig)
pos[ind[0]] = dt_annos[i]['location'][j]
mlab.view(azimuth=230, elevation=45, focalpoint=[0,0,0], distance=65.0)
img_path = os.path.join(output_dir,f"track_all_3d.png")
mlab.savefig(img_path)
def draw_track_3d(dt_objects,fig,idx):
for obj in dt_objects:
if obj.track == idx:
dt_corners = center_to_corner_box3d(np.array(obj.t)[np.newaxis,:],np.array([obj.l,obj.w,obj.h])[np.newaxis,:], np.array([obj.rz]))
print(dt_corners)
fig = draw_gt_boxes3d(dt_corners,None,draw_text=False,color=(0,0,1),fig=fig, view_setting=(235,70,60))
return fig
def plot_boxes(boxes, score_thresh):
visuals = []
num_det = boxes['scores'].shape[0]
for i in range(num_det):
score = boxes['scores'][i]
if score < score_thresh:
continue
box = boxes['boxes'][i:i + 1]
label = boxes['classes'][i]
corner = center_to_corner_box3d(box[:, :3], box[:, 3:6],
box[:, -1])[0].tolist()
color = label2color(label)
visuals.append(corners_to_lines(corner, color))
return visuals
def get_pyramids(gt_boxes):
'''
corners (in velo coord): [N, 8, 3]
6 -------- 5
/| /|
2 -------- 1 .
| | | |
. 7 -------- 4
|/ |/
3 -------- 0
gt_boxes: [N, 7].
return [N, 6, 15=5*3] pyramids.
surface order:
3
----------------
| | ^ x
4| 5(6) |2 |
| | y |
---------------- <------
1
'''
pyramid_orders = np.array([[0, 1, 5, 4], [4, 5, 6, 7], [7, 6, 2, 3],
[3, 2, 1, 0], [1, 2, 6, 5], [0, 4, 7, 3]])
boxes_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, 0:3], gt_boxes[:, 3:6], gt_boxes[:, 6], \
origin=[0.5, 0.5, 0.5], axis=2).reshape(-1, 24)
pyramid_list = []
for order in pyramid_orders:
# frustum polygon: 5 corners, 5 surfaces
pyramid = np.concatenate(
(gt_boxes[:, 0:3], boxes_corners[:, 3 * order[0]:3 * order[0] + 3],
boxes_corners[:, 3 * order[1]:3 * order[1] + 3],
boxes_corners[:, 3 * order[2]:3 * order[2] + 3],
boxes_corners[:, 3 * order[3]:3 * order[3] + 3]),
axis=1)
pyramid_list.append(pyramid[:, None, :])
pyramids = np.concatenate(pyramid_list, axis=1) # [N, 6, 15], 15=5*3
return pyramids
def noise_per_object_v2_(
gt_boxes,
points=None,
valid_mask=None,
rotation_perturb=np.pi / 4,
center_noise_std=1.0,
global_random_rot_range=np.pi / 4,
num_try=100,
):
"""random rotate or remove each groundtrutn independently.
use kitti viewer to test this function points_transform_
Args:
gt_boxes: [N, 7], gt box in lidar.points_transform_
points: [M, 4], point cloud in lidar.
"""
num_boxes = gt_boxes.shape[0]
if not isinstance(rotation_perturb, (list, tuple, np.ndarray)):
rotation_perturb = [-rotation_perturb, rotation_perturb]
if not isinstance(global_random_rot_range, (list, tuple, np.ndarray)):
global_random_rot_range = [
-global_random_rot_range, global_random_rot_range
]
if not isinstance(center_noise_std, (list, tuple, np.ndarray)):
center_noise_std = [
center_noise_std, center_noise_std, center_noise_std
]
if valid_mask is None:
valid_mask = np.ones((num_boxes, ), dtype=np.bool_)
center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype)
loc_noises = np.random.normal(scale=center_noise_std,
size=[num_boxes, num_try, 3])
# loc_noises = np.random.uniform(
# -center_noise_std, center_noise_std, size=[num_boxes, num_try, 3])
rot_noises = np.random.uniform(rotation_perturb[0],
rotation_perturb[1],
size=[num_boxes, num_try])
gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1])
grot_lowers = global_random_rot_range[0] - gt_grots
grot_uppers = global_random_rot_range[1] - gt_grots
global_rot_noises = np.random.uniform(
grot_lowers[..., np.newaxis],
grot_uppers[..., np.newaxis],
size=[num_boxes, num_try],
)
origin = [0.5, 0.5, 0]
gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3],
gt_boxes[:, 3:6],
gt_boxes[:, 6],
origin=origin,
axis=2)
if np.abs(global_random_rot_range[0] - global_random_rot_range[1]) < 1e-3:
selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask, loc_noises, rot_noises)
else:
selected_noise = noise_per_box_v2_(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
global_rot_noises,
)
loc_transforms = _select_transform(loc_noises, selected_noise)
rot_transforms = _select_transform(rot_noises, selected_noise)
if points is not None:
surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners)
point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces)
points_transform_(
points,
gt_boxes[:, :3],
point_masks,
loc_transforms,
rot_transforms,
valid_mask,
)
box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def convert_detection_to_kitti_annos(self, detection, partial=False):
class_names = self._class_names
det_image_idxes = [k for k in detection.keys()]
gt_image_idxes = [str(info["image"]["image_idx"]) for info in self._kitti_infos]
image_idxes = [gt_image_idxes, det_image_idxes]
# print(f"det_image_idxes: {det_image_idxes[:10]}")
# print(f"gt_image_idxes: {gt_image_idxes[:10]}")
annos = []
# for i in range(len(detection)):
for det_idx in image_idxes[int(partial==True)]:
det = detection[det_idx]
info = self._kitti_infos[gt_image_idxes.index(det_idx)]
# info = self._kitti_infos[i]
calib = info["calib"]
rect = calib["R0_rect"]
Trv2c = calib["Tr_velo_to_cam"]
P2 = calib["P2"]
final_box_preds = det["box3d_lidar"].detach().cpu().numpy()
label_preds = det["label_preds"].detach().cpu().numpy()
scores = det["scores"].detach().cpu().numpy()
anno = get_start_result_anno()
num_example = 0
if final_box_preds.shape[0] != 0:
final_box_preds[:, -1] = box_np_ops.limit_period(final_box_preds[:, -1], offset=0.5, period=np.pi * 2,)
final_box_preds[:, 2] -= final_box_preds[:, 5] / 2 # center_z -> bottom_z
# aim: x, y, z, w, l, h, r -> -y, -z, x, h, w, l, r
# (x, y, z, w, l, h r) in lidar -> (x', y', z', l, h, w, r) in camera
box3d_camera = box_np_ops.box_lidar_to_camera(final_box_preds, rect, Trv2c)
camera_box_origin = [0.5, 1.0, 0.5]
box_corners = box_np_ops.center_to_corner_box3d(box3d_camera[:, :3], box3d_camera[:, 3:6], box3d_camera[:, 6], camera_box_origin, axis=1,)
box_corners_in_image = box_np_ops.project_to_image(box_corners, P2)
# box_corners_in_image: [N, 8, 2]
minxy = np.min(box_corners_in_image, axis=1)
maxxy = np.max(box_corners_in_image, axis=1)
bbox = np.concatenate([minxy, maxxy], axis=1)
for j in range(box3d_camera.shape[0]):
image_shape = info["image"]["image_shape"]
if bbox[j, 0] > image_shape[1] or bbox[j, 1] > image_shape[0]:
continue
if bbox[j, 2] < 0 or bbox[j, 3] < 0:
continue
bbox[j, 2:] = np.minimum(bbox[j, 2:], image_shape[::-1])
bbox[j, :2] = np.maximum(bbox[j, :2], [0, 0])
anno["bbox"].append(bbox[j])
anno["alpha"].append(-np.arctan2(-final_box_preds[j, 1], final_box_preds[j, 0]) + box3d_camera[j, 6])
# anno["dimensions"].append(box3d_camera[j, [4, 5, 3]])
anno["dimensions"].append(box3d_camera[j, 3:6])
anno["location"].append(box3d_camera[j, :3])
anno["rotation_y"].append(box3d_camera[j, 6])
anno["name"].append(class_names[int(label_preds[j])])
anno["truncated"].append(0.0)
anno["occluded"].append(0)
anno["score"].append(scores[j])
num_example += 1
if num_example != 0:
anno = {n: np.stack(v) for n, v in anno.items()}
annos.append(anno)
else:
annos.append(empty_result_anno())
num_example = annos[-1]["name"].shape[0]
annos[-1]["metadata"] = det["metadata"]
return annos
def noise_per_object_v3_( gt_boxes, points=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0,\
global_random_rot_range=np.pi / 4, num_try=5, group_ids=None,):
"""random rotate or remove each groundtruth independently.
use kitti viewer to test this function points_transform_
Args:
gt_boxes: [N, 7], gt box in lidar.points_transform_
points: [M, 4], point cloud in lidar.
rotation_perturb: [-0.785, 0.785],
center_noise_std: [1.0, 1.0, 0.5],
global_random_rot_range: [0, 0]
num_try: 100
"""
num_boxes = gt_boxes.shape[0]
if not isinstance(rotation_perturb, (list, tuple, np.ndarray)): # False
rotation_perturb = [-rotation_perturb, rotation_perturb]
if not isinstance(global_random_rot_range,
(list, tuple, np.ndarray)): # False
global_random_rot_range = [
-global_random_rot_range, global_random_rot_range
]
enable_grot = (
np.abs(global_random_rot_range[0] - global_random_rot_range[1]) >= 1e-3
) # False
if not isinstance(center_noise_std, (list, tuple, np.ndarray)): # False
center_noise_std = [
center_noise_std, center_noise_std, center_noise_std
]
if valid_mask is None: # False
valid_mask = np.ones((num_boxes, ), dtype=np.bool_)
center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype)
loc_noises = np.random.normal(scale=center_noise_std,
size=[num_boxes, num_try, 3])
rot_noises = np.random.uniform(rotation_perturb[0],
rotation_perturb[1],
size=[num_boxes, num_try])
gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1])
grot_lowers = global_random_rot_range[0] - gt_grots
grot_uppers = global_random_rot_range[1] - gt_grots
global_rot_noises = np.random.uniform(
grot_lowers[..., np.newaxis],
grot_uppers[..., np.newaxis],
size=[num_boxes, num_try],
)
if group_ids is not None: # False
if enable_grot:
set_group_noise_same_v2_(loc_noises, rot_noises, global_rot_noises,
group_ids)
else:
set_group_noise_same_(loc_noises, rot_noises, group_ids)
group_centers, group_id_num_dict = get_group_center(
gt_boxes[:, :3], group_ids)
if enable_grot:
group_transform_v2_(
loc_noises,
rot_noises,
gt_boxes[:, :3],
gt_boxes[:, 6],
group_centers,
global_rot_noises,
valid_mask,
)
else:
group_transform_(
loc_noises,
rot_noises,
gt_boxes[:, :3],
gt_boxes[:, 6],
group_centers,
valid_mask,
)
group_nums = np.array(list(group_id_num_dict.values()), dtype=np.int64)
origin = [0.5, 0.5, 0.5]
gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3],
gt_boxes[:, 3:6],
gt_boxes[:, 6],
origin=origin,
axis=2)
if group_ids is not None:
if not enable_grot:
selected_noise = noise_per_box_group(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
group_nums,
)
else:
selected_noise = noise_per_box_group_v2_(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
group_nums,
global_rot_noises,
)
else:
if not enable_grot: # True
selected_noise = noise_per_box(
gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises,
rot_noises) # todo: seems without noise in z-axis velo
else:
selected_noise = noise_per_box_v2_(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
global_rot_noises,
)
loc_transforms = _select_transform(loc_noises, selected_noise)
rot_transforms = _select_transform(rot_noises, selected_noise)
surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners)
if points is not None:
point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces)
points_transform_(
points,
gt_boxes[:, :3],
point_masks,
loc_transforms,
rot_transforms,
valid_mask,
)
box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def noise_per_object_v3_(
gt_boxes,
gt_boxes_1,
gt_boxes_2,
gt_boxes_3,
points=None,
points_1=None,
points_2=None,
valid_mask=None,
rotation_perturb=np.pi / 4,
center_noise_std=1.0,
global_random_rot_range=np.pi / 4,
num_try=5,
group_ids=None,
):
"""random rotate or remove each groundtruth independently.
use kitti viewer to test this function points_transform_
Args:
gt_boxes: [N, 7], gt box in lidar.points_transform_
points: [M, 4], point cloud in lidar.
"""
num_boxes = gt_boxes.shape[0]
if not isinstance(rotation_perturb, (list, tuple, np.ndarray)):
rotation_perturb = [-rotation_perturb, rotation_perturb]
if not isinstance(global_random_rot_range, (list, tuple, np.ndarray)):
global_random_rot_range = [
-global_random_rot_range, global_random_rot_range
]
enable_grot = (np.abs(global_random_rot_range[0] -
global_random_rot_range[1]) >= 1e-3)
if not isinstance(center_noise_std, (list, tuple, np.ndarray)):
center_noise_std = [
center_noise_std, center_noise_std, center_noise_std
]
if valid_mask is None:
valid_mask = np.ones((num_boxes, ), dtype=np.bool_)
center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype)
loc_noises = np.random.normal(scale=center_noise_std,
size=[num_boxes, num_try, 3])
# loc_noises = np.random.uniform(
# -center_noise_std, center_noise_std, size=[num_boxes, num_try, 3])
rot_noises = np.random.uniform(rotation_perturb[0],
rotation_perturb[1],
size=[num_boxes, num_try])
gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1])
grot_lowers = global_random_rot_range[0] - gt_grots
grot_uppers = global_random_rot_range[1] - gt_grots
global_rot_noises = np.random.uniform(
grot_lowers[..., np.newaxis],
grot_uppers[..., np.newaxis],
size=[num_boxes, num_try],
)
if group_ids is not None:
if enable_grot:
set_group_noise_same_v2_(loc_noises, rot_noises, global_rot_noises,
group_ids)
else:
set_group_noise_same_(loc_noises, rot_noises, group_ids)
group_centers, group_id_num_dict = get_group_center(
gt_boxes[:, :3], group_ids)
if enable_grot:
group_transform_v2_(
loc_noises,
rot_noises,
gt_boxes[:, :3],
gt_boxes[:, 6],
group_centers,
global_rot_noises,
valid_mask,
)
else:
group_transform_(
loc_noises,
rot_noises,
gt_boxes[:, :3],
gt_boxes[:, 6],
group_centers,
valid_mask,
)
group_nums = np.array(list(group_id_num_dict.values()), dtype=np.int64)
origin = [0.5, 0.5, 0.5]
# TODO: 3个gt_box
gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3],
gt_boxes[:, 3:6],
gt_boxes[:, 6],
origin=origin,
axis=2)
# T-1
gt_box_corners_1 = box_np_ops.center_to_corner_box3d(gt_boxes_1[:, :3],
gt_boxes_1[:, 3:6],
gt_boxes_1[:, 6],
origin=origin,
axis=2)
# T+1
gt_box_corners_2 = box_np_ops.center_to_corner_box3d(gt_boxes_2[:, :3],
gt_boxes_2[:, 3:6],
gt_boxes_2[:, 6],
origin=origin,
axis=2)
# T-2
gt_box_corners_3 = box_np_ops.center_to_corner_box3d(gt_boxes_3[:, :3],
gt_boxes_3[:, 3:6],
gt_boxes_3[:, 6],
origin=origin,
axis=2)
if group_ids is not None:
if not enable_grot:
selected_noise = noise_per_box_group(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
group_nums,
)
else:
selected_noise = noise_per_box_group_v2_(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
group_nums,
global_rot_noises,
)
else:
if not enable_grot:
selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]],
gt_boxes_1[:, [0, 1, 3, 4, 6]],
gt_boxes_2[:, [0, 1, 3, 4, 6]],
gt_boxes_3[:, [0, 1, 3, 4, 6]],
valid_mask, loc_noises, rot_noises)
else:
selected_noise = noise_per_box_v2_(
gt_boxes[:, [0, 1, 3, 4, 6]],
valid_mask,
loc_noises,
rot_noises,
global_rot_noises,
)
loc_transforms = _select_transform(loc_noises, selected_noise)
rot_transforms = _select_transform(rot_noises, selected_noise)
surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners)
if points is not None:
point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces)
points_transform_(
points,
gt_boxes[:, :3],
point_masks,
loc_transforms,
rot_transforms,
valid_mask,
)
# T-1 帧的增广
loc_transforms_1 = get_loc_transform(loc_transforms, rot_transforms,
gt_boxes, gt_boxes_1)
surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners_1)
if points_1 is not None:
point_masks = points_in_convex_polygon_3d_jit(points_1[:, :3],
surfaces)
points_transform_(
points_1,
gt_boxes_1[:, :3],
point_masks,
loc_transforms_1,
rot_transforms,
valid_mask,
)
box3d_transform_(gt_boxes_1, loc_transforms_1, rot_transforms, valid_mask)
# T+1帧的增广
loc_transforms_2 = get_loc_transform(loc_transforms, rot_transforms,
gt_boxes, gt_boxes_2)
box3d_transform_(gt_boxes_2, loc_transforms_2, rot_transforms, valid_mask)
# T-2帧的增广
loc_transforms_3 = get_loc_transform(loc_transforms, rot_transforms,
gt_boxes, gt_boxes_3)
if points_2 is not None:
point_masks = points_in_convex_polygon_3d_jit(points_2[:, :3],
surfaces)
points_transform_(
points_2,
gt_boxes_3[:, :3],
point_masks,
loc_transforms_3,
rot_transforms,
valid_mask,
)
# T帧的box需要做为基准，所以最后变换
box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
