def get_direction_target(reg_targets, one_hot=True, dir_offset=0, num_bins=2): rot_gt = reg_targets[..., 6] offset_rot = box_torch_ops.limit_period(rot_gt - dir_offset, 0, 2 * np.pi) dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / num_bins)).long() dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=num_bins - 1) if one_hot: dir_cls_targets = torchplus.nn.one_hot(dir_cls_targets, num_bins, dtype=reg_targets.dtype) return dir_cls_targets
def get_direction_target(anchors, reg_targets, one_hot=True, dir_offset=0, num_bins=2): batch_size = reg_targets.shape[0] anchors = anchors.view(batch_size, -1, anchors.shape[-1]) rot_gt = reg_targets[..., 6] + anchors[..., 6] offset_rot = box_torch_ops.limit_period(rot_gt - dir_offset, 0, 2 * np.pi) dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / num_bins)).long() dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=num_bins - 1) if one_hot: dir_cls_targets = torchplus.nn.one_hot( dir_cls_targets, num_bins, dtype=anchors.dtype) return dir_cls_targets
def predict(self, example, preds_dict): """start with v1.6.0, this function don't contain any kitti-specific code. Returns: predict: list of pred_dict. pred_dict: { box3d_lidar: [N, 7] 3d box. scores: [N] label_preds: [N] metadata: meta-data which contains dataset-specific information. for kitti, it contains image idx (label idx), for nuscenes, sample_token is saved in it. } """ batch_size = example['anchors'].shape[0] if "metadata" not in example or len(example["metadata"]) == 0: meta_list = [None] * batch_size else: meta_list = example["metadata"] batch_anchors = example["anchors"].view(batch_size, -1, example["anchors"].shape[-1]) if "anchors_mask" not in example: batch_anchors_mask = [None] * batch_size else: batch_anchors_mask = example["anchors_mask"].view(batch_size, -1) t = time.time() batch_box_preds = preds_dict["box_preds"] batch_cls_preds = preds_dict["cls_preds"] batch_box_preds = batch_box_preds.view(batch_size, -1, self._box_coder.code_size) num_class_with_bg = self._num_class if not self._encode_background_as_zeros: num_class_with_bg = self._num_class + 1 batch_cls_preds = batch_cls_preds.view(batch_size, -1, num_class_with_bg) batch_box_preds = self._box_coder.decode_torch(batch_box_preds, batch_anchors) if self._use_direction_classifier: batch_dir_preds = preds_dict["dir_cls_preds"] batch_dir_preds = batch_dir_preds.view(batch_size, -1, self._num_direction_bins) else: batch_dir_preds = [None] * batch_size predictions_dicts = [] post_center_range = None if len(self._post_center_range) > 0: post_center_range = torch.tensor( self._post_center_range, dtype=batch_box_preds.dtype, device=batch_box_preds.device).float() for box_preds, cls_preds, dir_preds, a_mask, meta in zip( batch_box_preds, batch_cls_preds, batch_dir_preds, batch_anchors_mask, meta_list): if a_mask is not None: box_preds = box_preds[a_mask] cls_preds = cls_preds[a_mask] box_preds = box_preds.float() cls_preds = cls_preds.float() if self._use_direction_classifier: if a_mask is not None: dir_preds = dir_preds[a_mask] dir_labels = torch.max(dir_preds, dim=-1)[1] if self._encode_background_as_zeros: # this don't support softmax assert self._use_sigmoid_score is True total_scores = torch.sigmoid(cls_preds) else: # encode background as first element in one-hot vector if self._use_sigmoid_score: total_scores = torch.sigmoid(cls_preds)[..., 1:] else: total_scores = F.softmax(cls_preds, dim=-1)[..., 1:] # Apply NMS in birdeye view if self._use_rotate_nms: nms_func = box_torch_ops.rotate_nms else: nms_func = box_torch_ops.nms feature_map_size_prod = batch_box_preds.shape[ 1] // self.target_assigner.num_anchors_per_location if self._multiclass_nms: assert self._encode_background_as_zeros is True boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]] if not self._use_rotate_nms: box_preds_corners = box_torch_ops.center_to_corner_box2d( boxes_for_nms[:, :2], boxes_for_nms[:, 2:4], boxes_for_nms[:, 4]) boxes_for_nms = box_torch_ops.corner_to_standup_nd( box_preds_corners) selected_boxes, selected_labels, selected_scores = [], [], [] selected_dir_labels = [] scores = total_scores boxes = boxes_for_nms selected_per_class = [] score_threshs = self._nms_score_thresholds pre_max_sizes = self._nms_pre_max_sizes post_max_sizes = self._nms_post_max_sizes iou_thresholds = self._nms_iou_thresholds for class_idx, score_thresh, pre_ms, post_ms, iou_th in zip( range(self._num_class), score_threshs, pre_max_sizes, post_max_sizes, iou_thresholds): if self._nms_class_agnostic: class_scores = total_scores.view( feature_map_size_prod, -1, self._num_class)[..., class_idx] class_scores = class_scores.contiguous().view(-1) class_boxes_nms = boxes.view(-1, boxes_for_nms.shape[-1]) class_boxes = box_preds class_dir_labels = dir_labels else: anchors_range = self.target_assigner.anchors_range( class_idx) class_scores = total_scores.view( -1, self._num_class)[anchors_range[0]:anchors_range[1], class_idx] class_boxes_nms = boxes.view( -1, boxes_for_nms.shape[-1])[ anchors_range[0]:anchors_range[1], :] class_scores = class_scores.contiguous().view(-1) class_boxes_nms = class_boxes_nms.contiguous().view( -1, boxes_for_nms.shape[-1]) class_boxes = box_preds.view(-1, box_preds.shape[-1])[ anchors_range[0]:anchors_range[1], :] class_boxes = class_boxes.contiguous().view( -1, box_preds.shape[-1]) if self._use_direction_classifier: class_dir_labels = dir_labels.view( -1)[anchors_range[0]:anchors_range[1]] class_dir_labels = class_dir_labels.contiguous( ).view(-1) if score_thresh > 0.0: class_scores_keep = class_scores >= score_thresh if class_scores_keep.shape[0] == 0: selected_per_class.append(None) continue class_scores = class_scores[class_scores_keep] if class_scores.shape[0] != 0: if score_thresh > 0.0: class_boxes_nms = class_boxes_nms[ class_scores_keep] class_boxes = class_boxes[class_scores_keep] class_dir_labels = class_dir_labels[ class_scores_keep] keep = nms_func(class_boxes_nms, class_scores, pre_ms, post_ms, iou_th) if keep.shape[0] != 0: selected_per_class.append(keep) else: selected_per_class.append(None) else: selected_per_class.append(None) selected = selected_per_class[-1] if selected is not None: selected_boxes.append(class_boxes[selected]) selected_labels.append( torch.full([class_boxes[selected].shape[0]], class_idx, dtype=torch.int64, device=box_preds.device)) if self._use_direction_classifier: selected_dir_labels.append( class_dir_labels[selected]) selected_scores.append(class_scores[selected]) selected_boxes = torch.cat(selected_boxes, dim=0) selected_labels = torch.cat(selected_labels, dim=0) selected_scores = torch.cat(selected_scores, dim=0) if self._use_direction_classifier: selected_dir_labels = torch.cat(selected_dir_labels, dim=0) else: # get highest score per prediction, than apply nms # to remove overlapped box. if num_class_with_bg == 1: top_scores = total_scores.squeeze(-1) top_labels = torch.zeros(total_scores.shape[0], device=total_scores.device, dtype=torch.long) else: top_scores, top_labels = torch.max(total_scores, dim=-1) if self._nms_score_thresholds[0] > 0.0: top_scores_keep = top_scores >= self._nms_score_thresholds[ 0] top_scores = top_scores.masked_select(top_scores_keep) if top_scores.shape[0] != 0: if self._nms_score_thresholds[0] > 0.0: box_preds = box_preds[top_scores_keep] if self._use_direction_classifier: dir_labels = dir_labels[top_scores_keep] top_labels = top_labels[top_scores_keep] boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]] if not self._use_rotate_nms: box_preds_corners = box_torch_ops.center_to_corner_box2d( boxes_for_nms[:, :2], boxes_for_nms[:, 2:4], boxes_for_nms[:, 4]) boxes_for_nms = box_torch_ops.corner_to_standup_nd( box_preds_corners) # the nms in 3d detection just remove overlap boxes. selected = nms_func( boxes_for_nms, top_scores, pre_max_size=self._nms_pre_max_sizes[0], post_max_size=self._nms_post_max_sizes[0], iou_threshold=self._nms_iou_thresholds[0], ) else: selected = [] # if selected is not None: selected_boxes = box_preds[selected] if self._use_direction_classifier: selected_dir_labels = dir_labels[selected] selected_labels = top_labels[selected] selected_scores = top_scores[selected] # finally generate predictions. if selected_boxes.shape[0] != 0: box_preds = selected_boxes scores = selected_scores label_preds = selected_labels if self._use_direction_classifier: dir_labels = selected_dir_labels period = (2 * np.pi / self._num_direction_bins) dir_rot = box_torch_ops.limit_period( box_preds[..., 6] - self._dir_offset, self._dir_limit_offset, period) box_preds[ ..., 6] = dir_rot + self._dir_offset + period * dir_labels.to( box_preds.dtype) final_box_preds = box_preds final_scores = scores final_labels = label_preds if post_center_range is not None: mask = (final_box_preds[:, :3] >= post_center_range[:3]).all(1) mask &= (final_box_preds[:, :3] <= post_center_range[3:]).all(1) predictions_dict = { "box3d_lidar": final_box_preds[mask], "scores": final_scores[mask], "label_preds": label_preds[mask], "metadata": meta, } else: predictions_dict = { "box3d_lidar": final_box_preds, "scores": final_scores, "label_preds": label_preds, "metadata": meta, } else: dtype = batch_box_preds.dtype device = batch_box_preds.device predictions_dict = { "box3d_lidar": torch.zeros([0, box_preds.shape[-1]], dtype=dtype, device=device), "scores": torch.zeros([0], dtype=dtype, device=device), "label_preds": torch.zeros([0], dtype=top_labels.dtype, device=device), "metadata": meta, } predictions_dicts.append(predictions_dict) return predictions_dicts
def run_nms(self, box_preds, cls_preds, dir_preds, a_mask, meta, bbp_shape_1, bbp_dtype, bbp_device, post_center_range): if a_mask is not None: box_preds = box_preds[a_mask] cls_preds = cls_preds[a_mask] box_preds = box_preds.float() cls_preds = cls_preds.float() if self._use_direction_classifier: if a_mask is not None: dir_preds = dir_preds[a_mask] dir_labels = torch.max(dir_preds, dim=-1)[1] if self._encode_background_as_zeros: # this don't support softmax assert self._use_sigmoid_score is True total_scores = torch.sigmoid(cls_preds) else: # encode background as first element in one-hot vector if self._use_sigmoid_score: total_scores = torch.sigmoid(cls_preds)[..., 1:] else: total_scores = F.softmax(cls_preds, dim=-1)[..., 1:] # Apply NMS in birdeye view if self._use_rotate_nms: nms_func = box_torch_ops.rotate_nms else: nms_func = box_torch_ops.nms feature_map_size_prod = bbp_shape_1 // \ self.target_assigner.num_anchors_per_location if self._multiclass_nms: assert self._encode_background_as_zeros is True boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]] if not self._use_rotate_nms: box_preds_corners = box_torch_ops.center_to_corner_box2d( boxes_for_nms[:, :2], boxes_for_nms[:, 2:4], boxes_for_nms[:, 4]) boxes_for_nms = box_torch_ops.corner_to_standup_nd( box_preds_corners) selected_boxes, selected_labels, selected_scores = [], [], [] selected_dir_labels = [] scores = total_scores boxes = boxes_for_nms selected_per_class = [] score_threshs = self._nms_score_thresholds pre_max_sizes = self._nms_pre_max_sizes post_max_sizes = self._nms_post_max_sizes iou_thresholds = self._nms_iou_thresholds for class_idx, score_thresh, pre_ms, post_ms, iou_th in zip( range(self._num_class), score_threshs, pre_max_sizes, post_max_sizes, iou_thresholds): if self._nms_class_agnostic: class_scores = total_scores.view( feature_map_size_prod, -1, self._num_class)[..., class_idx] class_scores = class_scores.contiguous().view(-1) class_boxes_nms = boxes.view(-1, boxes_for_nms.shape[-1]) class_boxes = box_preds class_dir_labels = dir_labels else: anchors_range = self.target_assigner.anchors_range(class_idx) class_scores = total_scores.view( -1, self._num_class)[anchors_range[0]:anchors_range[1], class_idx] class_boxes_nms = boxes.view(-1, boxes_for_nms.shape[-1])[anchors_range[0]:anchors_range[1], :] class_scores = class_scores.contiguous().view(-1) class_boxes_nms = class_boxes_nms.contiguous().view( -1, boxes_for_nms.shape[-1]) class_boxes = box_preds.view(-1, box_preds.shape[-1])[anchors_range[0]:anchors_range[1], :] class_boxes = class_boxes.contiguous().view( -1, box_preds.shape[-1]) if self._use_direction_classifier: class_dir_labels = dir_labels.view(-1)[anchors_range[0]:anchors_range[1]] class_dir_labels = class_dir_labels.contiguous( ).view(-1) if score_thresh > 0.0: class_scores_keep = class_scores >= score_thresh if class_scores_keep.shape[0] == 0: selected_per_class.append(None) continue class_scores = class_scores[class_scores_keep] if class_scores.shape[0] != 0: if score_thresh > 0.0: class_boxes_nms = class_boxes_nms[ class_scores_keep] class_boxes = class_boxes[class_scores_keep] class_dir_labels = class_dir_labels[ class_scores_keep] keep = nms_func(class_boxes_nms, class_scores, pre_ms, post_ms, iou_th) if keep.shape[0] != 0: selected_per_class.append(keep) else: selected_per_class.append(None) else: selected_per_class.append(None) selected = selected_per_class[-1] if selected is not None: selected_boxes.append(class_boxes[selected]) selected_labels.append( torch.full([class_boxes[selected].shape[0]], class_idx, dtype=torch.int64, device=box_preds.device)) if self._use_direction_classifier: selected_dir_labels.append( class_dir_labels[selected]) selected_scores.append(class_scores[selected]) selected_boxes = torch.cat(selected_boxes, dim=0) selected_labels = torch.cat(selected_labels, dim=0) selected_scores = torch.cat(selected_scores, dim=0) if self._use_direction_classifier: selected_dir_labels = torch.cat(selected_dir_labels, dim=0) else: # get highest score per prediction, than apply nms # to remove overlapped box. if self._num_class_with_bg == 1: top_scores = total_scores.squeeze(-1) top_labels = torch.zeros( total_scores.shape[0], device=total_scores.device, dtype=torch.long) else: top_scores, top_labels = torch.max( total_scores, dim=-1) if self._nms_score_thresholds[0] > 0.0: top_scores_keep = top_scores >= self._nms_score_thresholds[0] top_scores = top_scores.masked_select(top_scores_keep) if top_scores.shape[0] != 0: if self._nms_score_thresholds[0] > 0.0: box_preds = box_preds[top_scores_keep] if self._use_direction_classifier: dir_labels = dir_labels[top_scores_keep] top_labels = top_labels[top_scores_keep] boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]] if not self._use_rotate_nms: box_preds_corners = box_torch_ops.center_to_corner_box2d( boxes_for_nms[:, :2], boxes_for_nms[:, 2:4], boxes_for_nms[:, 4]) boxes_for_nms = box_torch_ops.corner_to_standup_nd( box_preds_corners) # the nms in 3d detection just remove overlap boxes. selected = nms_func( boxes_for_nms, top_scores, pre_max_size=self._nms_pre_max_sizes[0], post_max_size=self._nms_post_max_sizes[0], iou_threshold=self._nms_iou_thresholds[0], ) else: selected = [] # if selected is not None: selected_boxes = box_preds[selected] if self._use_direction_classifier: selected_dir_labels = dir_labels[selected] selected_labels = top_labels[selected] selected_scores = top_scores[selected] # finally generate predictions. if selected_boxes.shape[0] != 0: box_preds = selected_boxes scores = selected_scores label_preds = selected_labels if self._use_direction_classifier: dir_labels = selected_dir_labels period = (2 * np.pi / self._num_direction_bins) dir_rot = box_torch_ops.limit_period( box_preds[..., 6] - self._dir_offset, self._dir_limit_offset, period) box_preds[ ..., 6] = dir_rot + self._dir_offset + period * dir_labels.to( box_preds.dtype) final_box_preds = box_preds final_scores = scores final_labels = label_preds if post_center_range is not None: mask = (final_box_preds[:, :3] >= post_center_range[:3]).all(1) mask &= (final_box_preds[:, :3] <= post_center_range[3:]).all(1) predictions_dict = { "box3d_lidar": final_box_preds[mask], "scores": final_scores[mask], "label_preds": label_preds[mask], "metadata": meta, } else: predictions_dict = { "box3d_lidar": final_box_preds, "scores": final_scores, "label_preds": label_preds, "metadata": meta, } else: dtype = bbp_dtype device = bbp_device predictions_dict = { "box3d_lidar": torch.zeros([0, box_preds.shape[-1]], dtype=dtype, device=device), "scores": torch.zeros([0], dtype=dtype, device=device), "label_preds": torch.zeros([0], dtype=top_labels.dtype, device=device), "metadata": meta, } return predictions_dict