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