def aug_roi_by_noise_batch(self, roi_boxes3d, gt_boxes3d, aug_times=10):
"""
:param roi_boxes3d: (N, 7)
:param gt_boxes3d: (N, 7)
:return:
"""
iou_of_rois = np.zeros(roi_boxes3d.shape[0], dtype=np.float32)
for k in range(roi_boxes3d.__len__()):
temp_iou = cnt = 0
roi_box3d = roi_boxes3d[k]
gt_box3d = gt_boxes3d[k]
pos_thresh = min(cfg.RCNN.REG_FG_THRESH, cfg.RCNN.CLS_FG_THRESH)
gt_corners = kitti_utils.boxes3d_to_corners3d(gt_box3d.reshape(1, 7))
aug_box3d = roi_box3d
while temp_iou < pos_thresh and cnt < aug_times:
if np.random.rand() < 0.2:
aug_box3d = roi_box3d # p=0.2 to keep the original roi box
else:
aug_box3d = self.random_aug_box3d(roi_box3d)
aug_corners = kitti_utils.boxes3d_to_corners3d(aug_box3d.reshape(1, 7))
iou3d = kitti_utils.get_iou3d(aug_corners, gt_corners)
temp_iou = iou3d[0][0]
cnt += 1
roi_boxes3d[k] = aug_box3d
iou_of_rois[k] = temp_iou
return roi_boxes3d, iou_of_rois
def aug_roi_by_noise(self, roi_info):
"""
add noise to original roi to get aug_box3d
:param roi_info:
:return:
"""
roi_box3d, gt_box3d = roi_info['roi_box3d'], roi_info['gt_box3d']
original_iou = roi_info['iou3d']
temp_iou = cnt = 0
pos_thresh = min(cfg.RCNN.REG_FG_THRESH, cfg.RCNN.CLS_FG_THRESH)
gt_corners = kitti_utils.boxes3d_to_corners3d(gt_box3d.reshape(-1, 7))
aug_box3d = roi_box3d
while temp_iou < pos_thresh and cnt < 10:
if roi_info['type'] == 'gt':
aug_box3d = self.random_aug_box3d(roi_box3d) # GT, must random
else:
if np.random.rand() < 0.2:
aug_box3d = roi_box3d # p=0.2 to keep the original roi box
else:
aug_box3d = self.random_aug_box3d(roi_box3d)
aug_corners = kitti_utils.boxes3d_to_corners3d(aug_box3d.reshape(-1, 7))
iou3d = kitti_utils.get_iou3d(aug_corners, gt_corners)
temp_iou = iou3d[0][0]
cnt += 1
if original_iou < pos_thresh: # original bg, break
break
return aug_box3d
def get_rcnn_training_sample_batch(self, index):
sample_id = int(self.sample_id_list[index])
rpn_xyz, rpn_features, rpn_intensity, seg_mask = \
self.get_rpn_features(self.rcnn_training_feature_dir, sample_id)
# load rois and gt_boxes3d for this sample
roi_file = os.path.join(self.rcnn_training_roi_dir, '%06d.txt' % sample_id)
roi_obj_list = kitti_utils.get_objects_from_label(roi_file)
roi_boxes3d = kitti_utils.objs_to_boxes3d(roi_obj_list)
# roi_scores = kitti_utils.objs_to_scores(roi_obj_list)
gt_obj_list = self.filtrate_objects(self.get_label(sample_id))
gt_boxes3d = kitti_utils.objs_to_boxes3d(gt_obj_list)
# calculate original iou
iou3d = kitti_utils.get_iou3d(kitti_utils.boxes3d_to_corners3d(roi_boxes3d),
kitti_utils.boxes3d_to_corners3d(gt_boxes3d))
max_overlaps, gt_assignment = iou3d.max(axis=1), iou3d.argmax(axis=1)
max_iou_of_gt, roi_assignment = iou3d.max(axis=0), iou3d.argmax(axis=0)
roi_assignment = roi_assignment[max_iou_of_gt > 0].reshape(-1)
# sample fg, easy_bg, hard_bg
fg_rois_per_lidar = int(np.round(cfg.RCNN.FG_RATIO * cfg.RCNN.ROI_PER_lidar))
fg_thresh = min(cfg.RCNN.REG_FG_THRESH, cfg.RCNN.CLS_FG_THRESH)
fg_inds = np.nonzero(max_overlaps >= fg_thresh)[0]
fg_inds = np.concatenate((fg_inds, roi_assignment), axis=0) # consider the roi which has max_overlaps with gt as fg
easy_bg_inds = np.nonzero((max_overlaps < cfg.RCNN.CLS_BG_THRESH_LO))[0]
hard_bg_inds = np.nonzero((max_overlaps < cfg.RCNN.CLS_BG_THRESH) &
(max_overlaps >= cfg.RCNN.CLS_BG_THRESH_LO))[0]
fg_num_rois = fg_inds.size
bg_num_rois = hard_bg_inds.size + easy_bg_inds.size
if fg_num_rois > 0 and bg_num_rois > 0:
# sampling fg
fg_rois_per_this_lidar = min(fg_rois_per_lidar, fg_num_rois)
rand_num = np.random.permutation(fg_num_rois)
fg_inds = fg_inds[rand_num[:fg_rois_per_this_lidar]]
# sampling bg
bg_rois_per_this_lidar = cfg.RCNN.ROI_PER_lidar - fg_rois_per_this_lidar
bg_inds = self.sample_bg_inds(hard_bg_inds, easy_bg_inds, bg_rois_per_this_lidar)
elif fg_num_rois > 0 and bg_num_rois == 0:
# sampling fg
rand_num = np.floor(np.random.rand(cfg.RCNN.ROI_PER_lidar ) * fg_num_rois)
rand_num = torch.from_numpy(rand_num).type_as(gt_boxes3d).long()
fg_inds = fg_inds[rand_num]
fg_rois_per_this_lidar = cfg.RCNN.ROI_PER_lidar
bg_rois_per_this_lidar = 0
elif bg_num_rois > 0 and fg_num_rois == 0:
# sampling bg
bg_rois_per_this_lidar = cfg.RCNN.ROI_PER_lidar
bg_inds = self.sample_bg_inds(hard_bg_inds, easy_bg_inds, bg_rois_per_this_lidar)
fg_rois_per_this_lidar = 0
else:
import pdb
pdb.set_trace()
raise NotImplementedError
# augment the rois by noise
roi_list, roi_iou_list, roi_gt_list = [], [], []
if fg_rois_per_this_lidar > 0:
fg_rois_src = roi_boxes3d[fg_inds].copy()
gt_of_fg_rois = gt_boxes3d[gt_assignment[fg_inds]]
fg_rois, fg_iou3d = self.aug_roi_by_noise_batch(fg_rois_src, gt_of_fg_rois, aug_times=10)
roi_list.append(fg_rois)
roi_iou_list.append(fg_iou3d)
roi_gt_list.append(gt_of_fg_rois)
if bg_rois_per_this_lidar > 0:
bg_rois_src = roi_boxes3d[bg_inds].copy()
gt_of_bg_rois = gt_boxes3d[gt_assignment[bg_inds]]
bg_rois, bg_iou3d = self.aug_roi_by_noise_batch(bg_rois_src, gt_of_bg_rois, aug_times=1)
roi_list.append(bg_rois)
roi_iou_list.append(bg_iou3d)
roi_gt_list.append(gt_of_bg_rois)
rois = np.concatenate(roi_list, axis=0)
iou_of_rois = np.concatenate(roi_iou_list, axis=0)
gt_of_rois = np.concatenate(roi_gt_list, axis=0)
# collect extra features for point cloud pooling
if cfg.RCNN.USE_INTENSITY:
pts_extra_input_list = [rpn_intensity.reshape(-1, 1), seg_mask.reshape(-1, 1)]
else:
pts_extra_input_list = [seg_mask.reshape(-1, 1)]
if cfg.RCNN.USE_DEPTH:
pts_depth = (np.linalg.norm(rpn_xyz, ord=2, axis=1) / 70.0) - 0.5
pts_extra_input_list.append(pts_depth.reshape(-1, 1))
pts_extra_input = np.concatenate(pts_extra_input_list, axis=1)
pts_input, pts_features, pts_empty_flag = roipool3d_utils.roipool3d_cpu(rois, rpn_xyz, rpn_features,
pts_extra_input,
cfg.RCNN.POOL_EXTRA_WIDTH,
sampled_pt_num=cfg.RCNN.NUM_POINTS,
canonical_transform=False)
valid_mask = (pts_empty_flag == 0).astype(np.int32)
# regression valid mask
reg_valid_mask = (iou_of_rois > cfg.RCNN.REG_FG_THRESH).astype(np.int32) & valid_mask
# classification label
cls_label = (iou_of_rois > cfg.RCNN.CLS_FG_THRESH).astype(np.int32)
invalid_mask = (iou_of_rois > cfg.RCNN.CLS_BG_THRESH) & (iou_of_rois < cfg.RCNN.CLS_FG_THRESH)
cls_label[invalid_mask] = -1
cls_label[valid_mask == 0] = -1
# canonical transform and sampling
pts_input_ct, gt_boxes3d_ct = self.canonical_transform_batch(pts_input, rois, gt_of_rois)
sample_info = {'sample_id': sample_id,
'pts_input': pts_input_ct,
'pts_features': pts_features,
'cls_label': cls_label,
'reg_valid_mask': reg_valid_mask,
'gt_boxes3d_ct': gt_boxes3d_ct,
'roi_boxes3d': rois,
'roi_size': rois[:, 3:6],
'gt_boxes3d': gt_of_rois}
return sample_info
def get_proposal_from_file(self, index):
sample_id = int(self.lidar_idx_list[index])
proposal_file = os.path.join(self.rcnn_eval_roi_dir, '%06d.txt' % sample_id)
roi_obj_list = kitti_utils.get_objects_from_label(proposal_file)
rpn_xyz, rpn_features, rpn_intensity, seg_mask = self.get_rpn_features(self.rcnn_eval_feature_dir, sample_id)
pts_rect, pts_rpn_features, pts_intensity = rpn_xyz, rpn_features, rpn_intensity
roi_box3d_list, roi_scores = [], []
for obj in roi_obj_list:
box3d = np.array([obj.pos[0], obj.pos[1], obj.pos[2], obj.h, obj.w, obj.l, obj.ry], dtype=np.float32)
roi_box3d_list.append(box3d.reshape(1, 7))
roi_scores.append(obj.score)
roi_boxes3d = np.concatenate(roi_box3d_list, axis=0) # (N, 7)
roi_scores = np.array(roi_scores, dtype=np.float32) # (N)
if cfg.RCNN.ROI_SAMPLE_JIT:
sample_dict = {'sample_id': sample_id,
'rpn_xyz': rpn_xyz,
'rpn_features': rpn_features,
'seg_mask': seg_mask,
'roi_boxes3d': roi_boxes3d,
'roi_scores': roi_scores,
'pts_depth': np.linalg.norm(rpn_xyz, ord=2, axis=1)}
if self.mode != 'TEST':
gt_obj_list = self.filtrate_objects(self.get_label(sample_id))
gt_boxes3d = kitti_utils.objs_to_boxes3d(gt_obj_list)
roi_corners = kitti_utils.boxes3d_to_corners3d(roi_boxes3d)
gt_corners = kitti_utils.boxes3d_to_corners3d(gt_boxes3d)
iou3d = kitti_utils.get_iou3d(roi_corners, gt_corners)
if gt_boxes3d.shape[0] > 0:
gt_iou = iou3d.max(axis=1)
else:
gt_iou = np.zeros(roi_boxes3d.shape[0]).astype(np.float32)
sample_dict['gt_boxes3d'] = gt_boxes3d
sample_dict['gt_iou'] = gt_iou
return sample_dict
if cfg.RCNN.USE_INTENSITY:
pts_extra_input_list = [pts_intensity.reshape(-1, 1), seg_mask.reshape(-1, 1)]
else:
pts_extra_input_list = [seg_mask.reshape(-1, 1)]
if cfg.RCNN.USE_DEPTH:
cur_depth = np.linalg.norm(pts_rect, axis=1, ord=2)
cur_depth_norm = (cur_depth / 70.0) - 0.5
pts_extra_input_list.append(cur_depth_norm.reshape(-1, 1))
pts_extra_input = np.concatenate(pts_extra_input_list, axis=1)
pts_input, pts_features = roipool3d_utils.roipool3d_cpu(roi_boxes3d, pts_rect, pts_rpn_features,
pts_extra_input, cfg.RCNN.POOL_EXTRA_WIDTH,
sampled_pt_num=cfg.RCNN.NUM_POINTS)
sample_dict = {'sample_id': sample_id,
'pts_input': pts_input,
'pts_features': pts_features,
'roi_boxes3d': roi_boxes3d,
'roi_scores': roi_scores,
'roi_size': roi_boxes3d[:, 3:6]}
if self.mode == 'TEST':
return sample_dict
gt_obj_list = self.filtrate_objects(self.get_label(sample_id))
gt_boxes3d = np.zeros((gt_obj_list.__len__(), 7), dtype=np.float32)
for k, obj in enumerate(gt_obj_list):
gt_boxes3d[k, 0:3], gt_boxes3d[k, 3], gt_boxes3d[k, 4], gt_boxes3d[k, 5], gt_boxes3d[k, 6] \
= obj.pos, obj.h, obj.w, obj.l, obj.ry
if gt_boxes3d.__len__() == 0:
gt_iou = np.zeros((roi_boxes3d.shape[0]), dtype=np.float32)
else:
roi_corners = kitti_utils.boxes3d_to_corners3d(roi_boxes3d)
gt_corners = kitti_utils.boxes3d_to_corners3d(gt_boxes3d)
iou3d = kitti_utils.get_iou3d(roi_corners, gt_corners)
gt_iou = iou3d.max(axis=1)
sample_dict['gt_boxes3d'] = gt_boxes3d
sample_dict['gt_iou'] = gt_iou
return sample_dict
def get_rcnn_sample_info(self, roi_info):
sample_id, gt_box3d = roi_info['sample_id'], roi_info['gt_box3d']
rpn_xyz, rpn_features, rpn_intensity, seg_mask = self.rpn_feature_list[sample_id]
# augmentation original roi by adding noise
roi_box3d = self.aug_roi_by_noise(roi_info)
# point cloud pooling based on roi_box3d
pooled_boxes3d = kitti_utils.enlarge_box3d(roi_box3d.reshape(1, 7), cfg.RCNN.POOL_EXTRA_WIDTH)
boxes_pts_mask_list = roipool3d_utils.pts_in_boxes3d_cpu(torch.from_numpy(rpn_xyz),
torch.from_numpy(pooled_boxes3d))
pt_mask_flag = (boxes_pts_mask_list[0].numpy() == 1)
cur_pts = rpn_xyz[pt_mask_flag].astype(np.float32)
# data augmentation
aug_pts = cur_pts.copy()
aug_gt_box3d = gt_box3d.copy().astype(np.float32)
aug_roi_box3d = roi_box3d.copy()
if cfg.AUG_DATA and self.mode == 'TRAIN':
# calculate alpha by ry
temp_boxes3d = np.concatenate([aug_roi_box3d.reshape(1, 7), aug_gt_box3d.reshape(1, 7)], axis=0)
temp_x, temp_z, temp_ry = temp_boxes3d[:, 0], temp_boxes3d[:, 2], temp_boxes3d[:, 6]
temp_beta = np.arctan2(temp_z, temp_x).astype(np.float64)
temp_alpha = -np.sign(temp_beta) * np.pi / 2 + temp_beta + temp_ry
# data augmentation
aug_pts, aug_boxes3d, aug_method = self.data_augmentation(aug_pts, temp_boxes3d, temp_alpha, mustaug=True, stage=2)
aug_roi_box3d, aug_gt_box3d = aug_boxes3d[0], aug_boxes3d[1]
aug_gt_box3d = aug_gt_box3d.astype(gt_box3d.dtype)
# Pool input points
valid_mask = 1 # whether the input is valid
if aug_pts.shape[0] == 0:
pts_features = np.zeros((1, 128), dtype=np.float32)
input_channel = 3 + int(cfg.RCNN.USE_INTENSITY) + int(cfg.RCNN.USE_MASK) + int(cfg.RCNN.USE_DEPTH)
pts_input = np.zeros((1, input_channel), dtype=np.float32)
valid_mask = 0
else:
pts_features = rpn_features[pt_mask_flag].astype(np.float32)
pts_intensity = rpn_intensity[pt_mask_flag].astype(np.float32)
pts_input_list = [aug_pts, pts_intensity.reshape(-1, 1)]
if cfg.RCNN.USE_INTENSITY:
pts_input_list = [aug_pts, pts_intensity.reshape(-1, 1)]
else:
pts_input_list = [aug_pts]
if cfg.RCNN.USE_MASK:
if cfg.RCNN.MASK_TYPE == 'seg':
pts_mask = seg_mask[pt_mask_flag].astype(np.float32)
elif cfg.RCNN.MASK_TYPE == 'roi':
pts_mask = roipool3d_utils.pts_in_boxes3d_cpu(torch.from_numpy(aug_pts),
torch.from_numpy(aug_roi_box3d.reshape(1, 7)))
pts_mask = (pts_mask[0].numpy() == 1).astype(np.float32)
else:
raise NotImplementedError
pts_input_list.append(pts_mask.reshape(-1, 1))
if cfg.RCNN.USE_DEPTH:
pts_depth = np.linalg.norm(aug_pts, axis=1, ord=2)
pts_depth_norm = (pts_depth / 70.0) - 0.5
pts_input_list.append(pts_depth_norm.reshape(-1, 1))
pts_input = np.concatenate(pts_input_list, axis=1) # (N, C)
aug_gt_corners = kitti_utils.boxes3d_to_corners3d(aug_gt_box3d.reshape(-1, 7))
aug_roi_corners = kitti_utils.boxes3d_to_corners3d(aug_roi_box3d.reshape(-1, 7))
iou3d = kitti_utils.get_iou3d(aug_roi_corners, aug_gt_corners)
cur_iou = iou3d[0][0]
# regression valid mask
reg_valid_mask = 1 if cur_iou >= cfg.RCNN.REG_FG_THRESH and valid_mask == 1 else 0
# classification label
cls_label = 1 if cur_iou > cfg.RCNN.CLS_FG_THRESH else 0
if cfg.RCNN.CLS_BG_THRESH < cur_iou < cfg.RCNN.CLS_FG_THRESH or valid_mask == 0:
cls_label = -1
# canonical transform and sampling
pts_input_ct, gt_box3d_ct = self.canonical_transform(pts_input, aug_roi_box3d, aug_gt_box3d)
pts_input_ct, pts_features = self.rcnn_input_sample(pts_input_ct, pts_features)
sample_info = {'sample_id': sample_id,
'pts_input': pts_input_ct,
'pts_features': pts_features,
'cls_label': cls_label,
'reg_valid_mask': reg_valid_mask,
'gt_boxes3d_ct': gt_box3d_ct,
'roi_boxes3d': aug_roi_box3d,
'roi_size': aug_roi_box3d[3:6],
'gt_boxes3d': aug_gt_box3d}
return sample_info
def apply_gt_aug_to_one_scene(self, sample_id, pts_rect, pts_intensity, all_gt_boxes3d):
"""
:param pts_rect: (N, 3)
:param all_gt_boxex3d: (M2, 7)
:return:
"""
assert self.gt_database is not None
# extra_gt_num = np.random.randint(10, 15)
# try_times = 50
if cfg.GT_AUG_RAND_NUM:
extra_gt_num = np.random.randint(10, cfg.GT_EXTRA_NUM)
else:
extra_gt_num = cfg.GT_EXTRA_NUM
try_times = 100
cnt = 0
cur_gt_boxes3d = all_gt_boxes3d.copy()
cur_gt_boxes3d[:, 4] += 0.5 # TODO: consider different objects
cur_gt_boxes3d[:, 5] += 0.5 # enlarge new added box to avoid too nearby boxes
cur_gt_corners = kitti_utils.boxes3d_to_corners3d(cur_gt_boxes3d)
extra_gt_obj_list = []
extra_gt_boxes3d_list = []
new_pts_list, new_pts_intensity_list = [], []
src_pts_flag = np.ones(pts_rect.shape[0], dtype=np.int32)
road_plane = self.get_road_plane(sample_id)
a, b, c, d = road_plane
while try_times > 0:
if cnt > extra_gt_num:
break
try_times -= 1
if cfg.GT_AUG_HARD_RATIO > 0:
p = np.random.rand()
if p > cfg.GT_AUG_HARD_RATIO:
# use easy sample
rand_idx = np.random.randint(0, len(self.gt_database[0]))
new_gt_dict = self.gt_database[0][rand_idx]
else:
# use hard sample
rand_idx = np.random.randint(0, len(self.gt_database[1]))
new_gt_dict = self.gt_database[1][rand_idx]
else:
rand_idx = np.random.randint(0, self.gt_database.__len__())
new_gt_dict = self.gt_database[rand_idx]
new_gt_box3d = new_gt_dict['gt_box3d'].copy()
new_gt_points = new_gt_dict['points'].copy()
new_gt_intensity = new_gt_dict['intensity'].copy()
new_gt_obj = new_gt_dict['obj']
center = new_gt_box3d[0:3]
if cfg.PC_REDUCE_BY_RANGE and (self.check_pc_range(center) is False):
continue
if new_gt_points.__len__() < 5: # too few points
continue
# put it on the road plane
cur_height = (-d - a * center[0] - c * center[2]) / b
move_height = new_gt_box3d[1] - cur_height
new_gt_box3d[1] -= move_height
new_gt_points[:, 1] -= move_height
new_gt_obj.pos[1] -= move_height
new_enlarged_box3d = new_gt_box3d.copy()
new_enlarged_box3d[4] += 0.5
new_enlarged_box3d[5] += 0.5 # enlarge new added box to avoid too nearby boxes
cnt += 1
new_corners = kitti_utils.boxes3d_to_corners3d(new_enlarged_box3d.reshape(1, 7))
iou3d = kitti_utils.get_iou3d(new_corners, cur_gt_corners)
valid_flag = iou3d.max() < 1e-8
if not valid_flag:
continue
enlarged_box3d = new_gt_box3d.copy()
enlarged_box3d[3] += 2 # remove the points above and below the object
boxes_pts_mask_list = roipool3d_utils.pts_in_boxes3d_cpu(
torch.from_numpy(pts_rect), torch.from_numpy(enlarged_box3d.reshape(1, 7)))
pt_mask_flag = (boxes_pts_mask_list[0].numpy() == 1)
src_pts_flag[pt_mask_flag] = 0 # remove the original points which are inside the new box
new_pts_list.append(new_gt_points)
new_pts_intensity_list.append(new_gt_intensity)
cur_gt_boxes3d = np.concatenate((cur_gt_boxes3d, new_enlarged_box3d.reshape(1, 7)), axis=0)
cur_gt_corners = np.concatenate((cur_gt_corners, new_corners), axis=0)
extra_gt_boxes3d_list.append(new_gt_box3d.reshape(1, 7))
extra_gt_obj_list.append(new_gt_obj)
if new_pts_list.__len__() == 0:
return False, pts_rect, pts_intensity, None, None
extra_gt_boxes3d = np.concatenate(extra_gt_boxes3d_list, axis=0)
# remove original points and add new points
pts_rect = pts_rect[src_pts_flag == 1]
pts_intensity = pts_intensity[src_pts_flag == 1]
new_pts_rect = np.concatenate(new_pts_list, axis=0)
new_pts_intensity = np.concatenate(new_pts_intensity_list, axis=0)
pts_rect = np.concatenate((pts_rect, new_pts_rect), axis=0)
pts_intensity = np.concatenate((pts_intensity, new_pts_intensity), axis=0)
return True, pts_rect, pts_intensity, extra_gt_boxes3d, extra_gt_obj_list
