def test_points_in_boxes():
if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda')
boxes = torch.tensor(
# x, y, z, l, h, w
[[1.0, 3.0, 2.0, 5.0, 6.0, 4.0, 0.3],
[-10.0, 23.0, 16.0, 10, 20, 20, 0.5]
],
dtype=torch.float32
).numpy() # boxes (m, 7) with bottom center in lidar coordinate
pts = torch.tensor(
[[1, 2, 3.3], [1.2, 2.5, 3.0], [0.8, 2.1, 3.5], [1.6, 2.6, 3.6],
[0.8, 1.2, 3.9], [-9.2, 21.0, 18.2], [3.8, 7.9, 6.3],
[4.7, 3.5, -12.2], [3.8, 7.6, -2], [-10.6, -12.9, -20], [-16, -18, 9],
[-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4]],
dtype=torch.float32).numpy() # points (n, 3) in lidar coordinate
expected_point_indices = torch.tensor(
[[1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
dtype=torch.bool).numpy().T
point_indices = check_inside_points(points=pts, cur_boxes=boxes)
assert point_indices.shape == torch.Size([15, 2])
assert (point_indices == expected_point_indices).all()
def generate_mixup_sample(self, sample_dict):
label_boxes_3d = sample_dict[maps_dict.KEY_LABEL_BOXES_3D]
label_classes = sample_dict[maps_dict.KEY_LABEL_CLASSES]
points = sample_dict[maps_dict.KEY_POINT_CLOUD]
label_class_names = np.array(
[self.idx2cls_dict[label] for label in label_classes])
tmp_label_boxes_3d = label_boxes_3d.copy()
# expand by 0.1, so as to cover context information
tmp_label_boxes_3d[:,
3:-1] += cfg.TRAIN.AUGMENTATIONS.EXPAND_DIMS_LENGTH
points_mask = check_inside_points(
points, tmp_label_boxes_3d) # [pts_num, gt_num]
pts_num_inside_box = np.sum(points_mask, axis=0) # gt_num
valid_box_idx = np.where(
pts_num_inside_box >= cfg.DATASET.MIN_POINTS_NUM)[0]
if len(valid_box_idx) == 0: return None
valid_label_boxes_3d = label_boxes_3d[valid_box_idx, :]
valid_label_classes = label_class_names[valid_box_idx]
sample_dicts = []
for index, i in enumerate(valid_box_idx):
cur_points_mask = points_mask[:, i]
cur_points_idx = np.where(cur_points_mask)[0]
cur_inside_points = points[cur_points_idx, :]
sample_dict = {
maps_dict.KEY_SAMPLED_GT_POINTS: cur_inside_points,
maps_dict.KEY_SAMPLED_GT_LABELS_3D:
valid_label_boxes_3d[index],
maps_dict.KEY_SAMPLED_GT_CLSES: valid_label_classes[index],
}
sample_dicts.append(sample_dict)
return sample_dicts
def vote_targets_torch(self, vote_base, gt_boxes_3d):
""" Generating vote_targets for each vote_base point
vote_base: [bs, points_num, 3]
gt_boxes_3d: [bs, gt_num, 7]
Return:
vote_mask: [bs, points_num]
vote_target: [bs, points_num, 3]
"""
bs, points_num, _ = vote_base.shape
vote_mask = torch.zeros((bs, points_num)).float().to(vote_base.device)
vote_target = torch.zeros(
(bs, points_num, 3)).float().to(vote_base.device)
for i in range(bs):
cur_vote_base = vote_base[i]
cur_gt_boxes_3d = gt_boxes_3d[i]
filter_idx = torch.where(
torch.any(torch.not_equal(cur_gt_boxes_3d, 0),
dim=-1))[0].to(vote_base.device)
cur_gt_boxes_3d = cur_gt_boxes_3d[filter_idx]
cur_vote_base_numpy = cur_vote_base.cpu().detach().numpy()
cur_expand_boxes_3d_numpy = cur_gt_boxes_3d.cpu().detach().numpy()
cur_expand_boxes_3d_numpy[:, 3:
-1] += cfg.TRAIN.AUGMENTATIONS.EXPAND_DIMS_LENGTH
cur_points_mask = check_inside_points(
cur_vote_base_numpy,
cur_expand_boxes_3d_numpy) # [pts_num, gt_num]
cur_vote_mask = np.max(cur_points_mask, axis=1).astype(np.float32)
vote_mask[i] = torch.from_numpy(cur_vote_mask).float().to(
vote_base.device)
cur_vote_target_idx = np.argmax(cur_points_mask,
axis=1) # [pts_num]
cur_vote_target_idx = torch.from_numpy(
cur_vote_target_idx).long().to(vote_base.device)
cur_vote_target = cur_gt_boxes_3d[cur_vote_target_idx]
cur_vote_target[:,
1] = cur_vote_target[:,
1] - cur_vote_target[:,
4] / 2.
cur_vote_target = cur_vote_target[:, :3] - cur_vote_base
vote_target[i] = cur_vote_target
return vote_mask, vote_target
def vote_targets_np(vote_base, gt_boxes_3d):
""" Generating vote_targets for each vote_base point
vote_base: [bs, points_num, 3]
gt_boxes_3d: [bs, gt_num, 7]
Return:
vote_mask: [bs, points_num]
vote_target: [bs, points_num, 3]
"""
bs, points_num, _ = vote_base.shape
vote_mask = np.zeros([bs, points_num], dtype=np.float32)
vote_target = np.zeros([bs, points_num, 3], dtype=np.float32)
for i in range(bs):
cur_vote_base = vote_base[i]
cur_gt_boxes_3d = gt_boxes_3d[i]
filter_idx = np.where(np.any(np.not_equal(cur_gt_boxes_3d, 0),
axis=-1))[0]
cur_gt_boxes_3d = cur_gt_boxes_3d[filter_idx]
cur_expand_boxes_3d = cur_gt_boxes_3d.copy()
cur_expand_boxes_3d[:,
3:-1] += cfg.TRAIN.AUGMENTATIONS.EXPAND_DIMS_LENGTH
cur_points_mask = check_inside_points(
cur_vote_base, cur_expand_boxes_3d) # [pts_num, gt_num]
cur_vote_mask = np.max(cur_points_mask, axis=1).astype(np.float32)
vote_mask[i] = cur_vote_mask
cur_vote_target_idx = np.argmax(cur_points_mask, axis=1) # [pts_num]
cur_vote_target = cur_gt_boxes_3d[cur_vote_target_idx]
cur_vote_target[:,
1] = cur_vote_target[:, 1] - cur_vote_target[:, 4] / 2.
cur_vote_target = cur_vote_target[:, :3] - cur_vote_base
vote_target[i] = cur_vote_target
return vote_mask, vote_target
def generate_mixup_sample(self, sample_dict):
""" This function is bound for generating mixup dataset """
all_boxes_3d = sample_dict[maps_dict.KEY_LABEL_BOXES_3D]
all_boxes_classes = sample_dict[maps_dict.KEY_LABEL_CLASSES]
point_cloud_path = sample_dict[maps_dict.KEY_POINT_CLOUD]
# then we first cast all_boxes_3d to kitti format
all_boxes_3d = cast_box_3d_to_kitti_format(all_boxes_3d)
# load points
points = np.fromfile(point_cloud_path, dtype=np.float32).reshape((-1, 5))
points = cast_points_to_kitti(points)
points[:, 3] /= 255
points[:, 4] = 0 # timestamp is zero
points_mask = check_inside_points(points, all_boxes_3d) # [pts_num, gt_num]
points_masks_num = np.sum(points_masks, axis=0) # [gt_num]
valid_box_idx = np.where(points_masks_num >= cfg.DATASET.MIN_POINTS_NUM)[0]
if len(valid_box_idx) == 0:
return None
valid_label_boxes_3d = all_boxes_3d[valid_box_idx]
valid_label_classes = all_boxes_classes[valid_box_idx]
sample_dicts = []
for index, i in enumerate(valid_box_idx):
cur_points_mask = points_mask[:, i]
cur_points_idx = np.where(cur_points_mask)[0]
cur_inside_points = points[cur_points_idx, :]
sample_dict = {
# 0 timestamp and /255 reflectance
maps_dict.KEY_SAMPLED_GT_POINTS: cur_inside_points, # kitti format points
maps_dict.KEY_SAMPLED_GT_LABELS_3D: valid_label_boxes_3d[index],
maps_dict.KEY_SAMPLED_GT_CLSES: valid_label_classes[index],
}
sample_dicts.append(sample_dict)
return sample_dicts
def preprocess_samples(self, indices):
sample_dicts = []
biggest_label_num = 0
for sample_idx in indices:
sample_id = int(self.idx_list[sample_idx])
img = self.kitti_object.get_image(sample_id)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image_shape = img.shape
calib = self.kitti_object.get_calibration(sample_id)
points = self.kitti_object.get_lidar(sample_id)
points_intensity = points[:, 3:]
points = points[:, :3]
# filter out this, first cast it to rect
points = calib.project_velo_to_rect(points)
img_points_filter = get_point_filter_in_image(
points, calib, image_shape[0], image_shape[1])
voxelnet_points_filter = get_point_filter(points, self.extents)
img_points_filter = np.logical_and(img_points_filter,
voxelnet_points_filter)
img_points_filter = np.where(img_points_filter)[0]
points = points[img_points_filter]
points_intensity = points_intensity[img_points_filter]
if self.img_list in ['train', 'val', 'trainval'
] and cfg.TEST.WITH_GT:
# then we also need to preprocess groundtruth
objs = self.kitti_object.get_label_objects(sample_id)
filtered_obj_list = [
obj for obj in objs if obj.type in self.cls_list
]
if len(filtered_obj_list) == 0:
# continue if no obj
return None, biggest_label_num
# then is time to generate anchors
label_boxes_3d = np.array(
[object_label_to_box_3d(obj) for obj in filtered_obj_list])
label_boxes_3d = np.reshape(label_boxes_3d, [-1, 7])
label_classes = np.array(
[self.cls2idx_dict[obj.type] for obj in filtered_obj_list],
np.int)
# then calculate sem_labels and sem_dists
tmp_label_boxes_3d = label_boxes_3d.copy()
# expand by 0.1, so as to cover context information
tmp_label_boxes_3d[:, 3:
-1] += cfg.TRAIN.AUGMENTATIONS.EXPAND_DIMS_LENGTH
points_mask = check_inside_points(
points, tmp_label_boxes_3d) # [pts_num, gt_num]
points_cls_index = np.argmax(points_mask, axis=1) # [pts_num]
points_cls_index = label_classes[points_cls_index] # [pts_num]
sem_labels = np.max(points_mask,
axis=1) * points_cls_index # [pts_num]
sem_labels = sem_labels.astype(np.int)
sem_dists = np.ones_like(sem_labels).astype(np.float32)
else:
sem_labels = np.ones([points.shape[0]], dtype=np.int)
sem_dists = np.ones([points.shape[0]], dtype=np.float32)
points = np.concatenate([points, points_intensity], axis=-1)
if np.sum(sem_labels) == 0:
return None, biggest_label_num
# finally return the sealed result and save as npy file
if self.img_list in ['train', 'val', 'trainval'
] and cfg.TEST.WITH_GT:
sample_dict = {
maps_dict.KEY_LABEL_BOXES_3D: label_boxes_3d,
maps_dict.KEY_LABEL_CLASSES: label_classes,
maps_dict.KEY_LABEL_SEMSEG: sem_labels,
maps_dict.KEY_LABEL_DIST: sem_dists,
maps_dict.KEY_POINT_CLOUD: points,
maps_dict.KEY_STEREO_CALIB: calib,
maps_dict.KEY_SAMPLE_NAME: sample_id,
maps_dict.KEY_LABEL_NUM: len(label_boxes_3d)
}
biggest_label_num = max(len(label_boxes_3d), biggest_label_num)
else:
# img_list is test
sample_dict = {
maps_dict.KEY_LABEL_SEMSEG: sem_labels,
maps_dict.KEY_LABEL_DIST: sem_dists,
maps_dict.KEY_POINT_CLOUD: points,
maps_dict.KEY_STEREO_CALIB: calib,
maps_dict.KEY_SAMPLE_NAME: sample_id
}
sample_dicts.append(sample_dict)
return sample_dicts, biggest_label_num
def __mask_assign_targets_anchors_np(self, batch_points, batch_anchors_3d,
batch_gt_boxes_3d, batch_gt_labels,
minibatch_size, positive_rate,
pos_iou, neg_iou,
effective_sample_range, valid_mask):
""" Mask assign targets function
batch_points: [bs, points_num, 3]
batch_anchors_3d: [bs, points_num, cls_num, 7]
batch_gt_boxes_3d: [bs, gt_num, 7]
batch_gt_labels: [bs, gt_num]
valid_mask: [bs, points_num, cls_num]
return:
assigned_idx: [bs, points_num, cls_num], int32, the index of groundtruth
assigned_pmask: [bs, points_num, cls_num], float32
assigned_nmask: [bs, points_num, cls_num], float32
"""
bs, pts_num, cls_num, _ = batch_anchors_3d.shape
positive_size = int(minibatch_size * positive_rate)
batch_assigned_idx = np.zeros([bs, pts_num, cls_num], np.int32)
batch_assigned_pmask = np.zeros([bs, pts_num, cls_num], np.float32)
batch_assigned_nmask = np.zeros([bs, pts_num, cls_num], np.float32)
for i in range(bs):
cur_points = batch_points[i]
cur_anchors_3d = batch_anchors_3d[i] # [pts_num, cls_num, 3/7]
cur_valid_mask = valid_mask[i] # [pts_num, cls_num]
# gt_num
cur_gt_labels = batch_gt_labels[i] # [gt_num]
cur_gt_boxes_3d = batch_gt_boxes_3d[i] # [gt_num, 7]
# first filter gt_boxes
filter_idx = np.where(
np.any(np.not_equal(cur_gt_boxes_3d, 0), axis=-1))[0]
cur_gt_labels = cur_gt_labels[filter_idx]
cur_gt_boxes_3d = cur_gt_boxes_3d[filter_idx]
points_mask = check_inside_points(
cur_points, cur_gt_boxes_3d) # [pts_num, gt_num]
sampled_gt_idx = np.argmax(points_mask, axis=-1) # [pts_num]
# used for label_mask
assigned_gt_label = cur_gt_labels[sampled_gt_idx] # [pts_num]
assigned_gt_label = assigned_gt_label - 1 # 1... -> 0...
# used for dist_mask
assigned_gt_boxes = cur_gt_boxes_3d[sampled_gt_idx] # [pts_num, 7]
# then calc the distance between anchors and assigned_boxes
dist = np.linalg.norm(cur_anchors_3d[:, :, :3] -
assigned_gt_boxes[:, np.newaxis, :3],
axis=-1) # [pts_num, cls_num]
filtered_assigned_idx = filter_idx[sampled_gt_idx] # [pts_num]
filtered_assigned_idx = np.tile(
np.reshape(filtered_assigned_idx, [pts_num, 1]), [1, cls_num])
batch_assigned_idx[i] = filtered_assigned_idx
if cls_num == 1: # anchor_free
label_mask = np.ones([pts_num, cls_num], dtype=np.float32)
else: # multiple anchors
label_mask = np.tile(
np.reshape(np.arange(cls_num), [1, cls_num]), [pts_num, 1])
label_mask = np.equal(label_mask,
assigned_gt_label[:, np.newaxis]).astype(
np.float32)
pmask = np.max(points_mask, axis=1) > 0
dist_mask = np.less_equal(
dist, effective_sample_range) # pts_num, cls_num
pmask = np.logical_and(pmask[:, np.newaxis],
dist_mask).astype(np.float32)
pmask = pmask * label_mask
pmask = pmask * cur_valid_mask
nmask = np.max(points_mask, axis=1) == 0
nmask = np.tile(np.reshape(nmask, [pts_num, 1]), [1, cls_num])
nmask = nmask * label_mask
nmask = nmask * cur_valid_mask
# then randomly sample
if minibatch_size != -1:
pts_pmask = np.any(pmask, axis=1) # pts_num
pts_nmask = np.any(nmask, axis=1) # [pts_num]
positive_inds = np.where(pts_pmask)[0]
cur_positive_num = np.minimum(len(positive_inds),
positive_size)
if cur_positive_num > 0:
positive_inds = np.random.choice(positive_inds,
cur_positive_num,
replace=False)
pts_pmask = np.zeros_like(pts_pmask)
pts_pmask[positive_inds] = 1
cur_negative_num = minibatch_size - cur_positive_num
negative_inds = np.where(pts_nmask)[0]
cur_negative_num = np.minimum(len(negative_inds),
cur_negative_num)
if cur_negative_num > 0:
negative_inds = np.random.choice(negative_inds,
cur_negative_num,
replace=False)
pts_nmask = np.zeros_like(pts_nmask)
pts_nmask[negative_inds] = 1
pmask = pmask * pts_pmask[:, np.newaxis]
nmask = nmask * pts_nmask[:, np.newaxis]
batch_assigned_pmask[i] = pmask
batch_assigned_nmask[i] = nmask
return batch_assigned_idx, batch_assigned_pmask, batch_assigned_nmask
def __mask_assign_targets_anchors_torch(
self, batch_points, batch_anchors_3d, batch_gt_boxes_3d,
batch_gt_labels, minibatch_size, positive_rate, pos_iou, neg_iou,
effective_sample_range, valid_mask):
""" Mask assign targets function
batch_points: [bs, points_num, 3]
batch_anchors_3d: [bs, points_num, cls_num, 7]
batch_gt_boxes_3d: [bs, gt_num, 7]
batch_gt_labels: [bs, gt_num]
valid_mask: [bs, points_num, cls_num]
return:
assigned_idx: [bs, points_num, cls_num], int32, the index of groundtruth
assigned_pmask: [bs, points_num, cls_num], float32
assigned_nmask: [bs, points_num, cls_num], float32
"""
bs, pts_num, cls_num, _ = batch_anchors_3d.shape
positive_size = int(minibatch_size * positive_rate)
batch_assigned_idx = torch.zeros([bs, pts_num, cls_num
]).long().to(batch_points.device)
batch_assigned_pmask = torch.zeros([bs, pts_num, cls_num
]).float().to(batch_points.device)
batch_assigned_nmask = torch.zeros([bs, pts_num, cls_num
]).float().to(batch_points.device)
for i in range(bs):
cur_points = batch_points[i]
cur_anchors_3d = batch_anchors_3d[i] # [pts_num, cls_num, 3/7]
cur_valid_mask = valid_mask[i] # [pts_num, cls_num]
# gt_num
cur_gt_labels = batch_gt_labels[i] # [gt_num]
cur_gt_boxes_3d = batch_gt_boxes_3d[i] # [gt_num, 7]
# first filter gt_boxes
filter_idx = torch.where(
torch.any(torch.not_equal(cur_gt_boxes_3d, 0),
dim=-1))[0].to(cur_gt_labels.device)
cur_gt_labels = cur_gt_labels[filter_idx]
cur_gt_boxes_3d = cur_gt_boxes_3d[filter_idx]
cur_points_numpy = cur_points.cpu().detach().numpy()
cur_gt_boxes_3d_numpy = cur_gt_boxes_3d.cpu().detach().numpy()
points_mask_numpy = check_inside_points(
cur_points_numpy, cur_gt_boxes_3d_numpy) # [pts_num, gt_num]
points_mask = torch.from_numpy(points_mask_numpy).int().to(
cur_points.device)
sampled_gt_idx_numpy = np.argmax(points_mask_numpy, axis=-1)
sampled_gt_idx = torch.from_numpy(sampled_gt_idx_numpy).long().to(
cur_points.device) # [pts_num]
# used for label_mask
assigned_gt_label = cur_gt_labels[sampled_gt_idx] # [pts_num]
assigned_gt_label = assigned_gt_label - 1 # 1... -> 0...
# used for dist_mask
assigned_gt_boxes = cur_gt_boxes_3d[sampled_gt_idx] # [pts_num, 7]
# then calc the distance between anchors and assigned_boxes
# dist = cur_anchors_3d[:, :, :3] - assigned_gt_boxes[:, 0:3].unsqueeze(dim=1).repeat((1, cur_anchors_3d.shape[1], 1))
# dist = torch.sqrt(torch.sum(dist * dist, dim=-1))
dist = torch.linalg.norm(
cur_anchors_3d[:, :, :3] -
assigned_gt_boxes[:, 0:3].unsqueeze(dim=1).repeat(
(1, cur_anchors_3d.shape[1], 1)),
dim=-1)
filtered_assigned_idx = filter_idx[sampled_gt_idx] # [pts_num]
filtered_assigned_idx = filtered_assigned_idx.view(pts_num,
1).repeat(
(1,
cls_num))
batch_assigned_idx[i] = filtered_assigned_idx
# then we generate pos/neg mask
if cls_num == 1: # anchor_free
label_mask = torch.ones(
(pts_num, cls_num)).float().to(points_mask.device)
else: # multiple anchors
label_mask = np.tile(
np.reshape(np.arange(cls_num), [1, cls_num]), [pts_num, 1])
label_mask = np.equal(label_mask,
assigned_gt_label[:, np.newaxis]).astype(
np.float32)
pmask = torch.max(points_mask, dim=1)[0] > 0
dist_mask = torch.less_equal(
dist, effective_sample_range) # pts_num, cls_num
pmask = torch.logical_and(pmask.unsqueeze(-1), dist_mask)
pmask = pmask.float() * label_mask
pmask = pmask * cur_valid_mask
nmask = torch.max(points_mask, dim=1)[0] == 0
nmask = nmask.view(pts_num, 1).repeat((1, cls_num))
nmask = nmask.float() * label_mask
nmask = nmask * cur_valid_mask
# then randomly sample
if minibatch_size != -1:
pts_pmask = np.any(pmask, axis=1) # pts_num
pts_nmask = np.any(nmask, axis=1) # [pts_num]
positive_inds = np.where(pts_pmask)[0]
cur_positive_num = np.minimum(len(positive_inds),
positive_size)
if cur_positive_num > 0:
positive_inds = np.random.choice(positive_inds,
cur_positive_num,
replace=False)
pts_pmask = np.zeros_like(pts_pmask)
pts_pmask[positive_inds] = 1
cur_negative_num = minibatch_size - cur_positive_num
negative_inds = np.where(pts_nmask)[0]
cur_negative_num = np.minimum(len(negative_inds),
cur_negative_num)
if cur_negative_num > 0:
negative_inds = np.random.choice(negative_inds,
cur_negative_num,
replace=False)
pts_nmask = np.zeros_like(pts_nmask)
pts_nmask[negative_inds] = 1
pmask = pmask * pts_pmask[:, np.newaxis]
nmask = nmask * pts_nmask[:, np.newaxis]
batch_assigned_pmask[i] = pmask
batch_assigned_nmask[i] = nmask
return batch_assigned_idx, batch_assigned_pmask, batch_assigned_nmask
def iou_assign_targets_anchors_np(batch_iou_matrix, batch_points,
batch_anchors_3d, batch_gt_boxes_3d,
batch_gt_labels, minibatch_size,
positive_rate, pos_iou, neg_iou,
effective_sample_range, valid_mask):
""" IoU assign targets function
batch_iou_matrix: [bs, points_num, cls_num, gt_num]
batch_points: [bs, points_num, 3]
batch_anchors_3d: [bs, points_num, cls_num, 7]
batch_gt_boxes_3d: [bs, gt_num, 7]
batch_gt_labels: [bs, gt_num]
valid_mask: [bs, points_num, cls_num]
return:
assigned_idx: [bs, points_num, cls_num], int32, the index of groundtruth
assigned_pmask: [bs, points_num, cls_num], float32
assigned_nmask: [bs, points_num, cls_num], float32
"""
bs, pts_num, cls_num, gt_num = batch_iou_matrix.shape
positive_size = int(minibatch_size * positive_rate)
batch_assigned_idx = np.zeros([bs, pts_num, cls_num], np.int32)
batch_assigned_pmask = np.zeros([bs, pts_num, cls_num], np.float32)
batch_assigned_nmask = np.zeros([bs, pts_num, cls_num], np.float32)
for i in range(bs):
# first calc the 3d iou matrix or 2d iou
# pts_num, cls_num, 7
cur_points = batch_points[i]
cur_anchors_3d = batch_anchors_3d[i] # [pts_num, cls_num, 7]
cur_valid_mask = valid_mask[i]
# gt_num
cur_gt_labels = batch_gt_labels[i] # [gt_num]
cur_gt_boxes_3d = batch_gt_boxes_3d[i] # [gt_num, 7]
iou_matrix = batch_iou_matrix[i] # [pts_num, cls_num, gt_num]
# first filter gt_boxes
filter_idx = np.where(np.any(np.not_equal(cur_gt_boxes_3d, 0),
axis=-1))[0]
cur_gt_labels = cur_gt_labels[filter_idx]
cur_gt_boxes_3d = cur_gt_boxes_3d[filter_idx]
iou_matrix = iou_matrix[:, :, filter_idx]
# first we check whether a point is within a box
points_mask = check_inside_points(cur_points,
cur_gt_boxes_3d) # [pts_num, gt_num]
sampled_gt_idx = np.argmax(points_mask, axis=-1) # [pts_num]
# used for generating label_mask
assigned_gt_label = cur_gt_labels[sampled_gt_idx] # [pts_num]
assigned_gt_label = assigned_gt_label - 1 # 1... -> 0...
# used for generating dist_mask
assigned_gt_boxes = cur_gt_boxes_3d[sampled_gt_idx] # [pts_num, 7]
# then calc the distance between anchors and assigned_boxes
dist = np.linalg.norm(cur_anchors_3d[:, :, :3] -
assigned_gt_boxes[:, np.newaxis, :3],
axis=-1) # [pts_num, cls_num]
# then we get assigned_idx by whether a point is within an object
filtered_assigned_idx = filter_idx[sampled_gt_idx] # [pts_num]
filtered_assigned_idx = np.tile(
np.reshape(filtered_assigned_idx, [pts_num, 1]), [1, cls_num])
batch_assigned_idx[i] = filtered_assigned_idx
# then we generate pos/neg mask
assigned_idx = np.tile(np.reshape(sampled_gt_idx, [pts_num, 1, 1]),
[1, cls_num, 1])
iou_mask = np.tile(
np.reshape(np.arange(len(filter_idx)),
[1, 1, len(filter_idx)]),
[pts_num, cls_num, 1]) # [pts_num, cls_num, len(filter_idx)]
iou_matrix = np.sum(
np.equal(iou_mask, assigned_idx).astype(np.float32) * iou_matrix,
axis=-1) # [pts_num, cls_num]
if cls_num > 1:
label_mask = np.tile(np.reshape(np.arange(cls_num), [1, cls_num]),
[pts_num, 1])
label_mask = np.equal(label_mask,
assigned_gt_label[:, np.newaxis]).astype(
np.float32)
else:
label_mask = np.ones([pts_num, cls_num], dtype=np.float32)
iou_matrix = iou_matrix * label_mask + (1 - label_mask) * np.ones_like(
iou_matrix) * -1 # count and ignored
pmask = np.greater_equal(iou_matrix, pos_iou) # [pts_num, gt_num]
dist_mask = np.less_equal(dist, effective_sample_range)
pmask = np.logical_and(pmask, dist_mask).astype(np.float32)
nmask = np.logical_and(np.less(iou_matrix, neg_iou),
np.greater_equal(iou_matrix,
0.05)).astype(np.float32)
pmask = pmask * cur_valid_mask
nmask = nmask * cur_valid_mask
# finally let's randomly choice some points
if minibatch_size != -1:
pts_pmask = np.any(pmask, axis=1) # [pts_num]
pts_nmask = np.any(nmask, axis=1) # [pts_num]
positive_inds = np.where(pts_pmask)[0]
cur_positive_num = np.minimum(len(positive_inds), positive_size)
if cur_positive_num > 0:
positive_inds = np.random.choice(positive_inds,
cur_positive_num,
replace=False)
pts_pmask = np.zeros_like(pts_pmask)
pts_pmask[positive_inds] = 1
cur_negative_num = minibatch_size - cur_positive_num
negative_inds = np.where(pts_nmask)[0]
cur_negative_num = np.minimum(len(negative_inds), cur_negative_num)
if cur_negative_num > 0:
negative_inds = np.random.choice(negative_inds,
cur_negative_num,
replace=False)
pts_nmask = np.zeros_like(pts_nmask)
pts_nmask[negative_inds] = 1
pmask = pmask * pts_pmask[:, np.newaxis]
nmask = nmask * pts_nmask[:, np.newaxis]
batch_assigned_pmask[i] = pmask
batch_assigned_nmask[i] = nmask
return batch_assigned_idx, batch_assigned_pmask, batch_assigned_nmask
