Non-Maximum Suppression for 3D Euler boxes. Additionally, only the top-k boxes will be kept for each category to avoid redundant boxes in the visualization. Args: pred_results (mmengine.structures.instance_data.InstanceData): Results predicted by the model io
(pred_results, iou_thr=0.15, score_thr=0.075, topk_per_class=10)
| 82 | |
| 83 | |
| 84 | def nms_filter(pred_results, iou_thr=0.15, score_thr=0.075, topk_per_class=10): |
| 85 | """Non-Maximum Suppression for 3D Euler boxes. Additionally, only the top-k |
| 86 | boxes will be kept for each category to avoid redundant boxes in the |
| 87 | visualization. |
| 88 | |
| 89 | Args: |
| 90 | pred_results (mmengine.structures.instance_data.InstanceData): |
| 91 | Results predicted by the model |
| 92 | iou_thr (float): IoU thresholds for NMS. Defaults to 0.15. |
| 93 | score_thr (float): Score thresholds. |
| 94 | Instances with scores below thresholds will not be kept. |
| 95 | Defaults to 0.075. |
| 96 | topk_per_class (int): Number of instances kept per category. |
| 97 | |
| 98 | Returns: |
| 99 | boxes (numpy.ndarray[float]): filtered instances, shape (N,9) |
| 100 | labels (numpy.ndarray[int]): filtered labels, shape (N,) |
| 101 | """ |
| 102 | boxes = pred_results.bboxes_3d |
| 103 | boxes_tensor = boxes.tensor.cpu().numpy() |
| 104 | iou = boxes.overlaps(boxes, boxes, eps=1e-5) |
| 105 | score = pred_results.scores_3d.cpu().numpy() |
| 106 | label = pred_results.labels_3d.cpu().numpy() |
| 107 | selected_per_class = dict() |
| 108 | |
| 109 | n = boxes_tensor.shape[0] |
| 110 | idx = list(range(n)) |
| 111 | idx.sort(key=lambda x: score[x], reverse=True) |
| 112 | selected_idx = [] |
| 113 | for i in idx: |
| 114 | if selected_per_class.get(label[i], 0) >= topk_per_class: |
| 115 | continue |
| 116 | if score[i] < score_thr: |
| 117 | continue |
| 118 | bo = False |
| 119 | for j in selected_idx: |
| 120 | if iou[i][j] > iou_thr: |
| 121 | bo = True |
| 122 | break |
| 123 | if not bo: |
| 124 | selected_idx.append(i) |
| 125 | if label[i] not in selected_per_class: |
| 126 | selected_per_class[label[i]] = 1 |
| 127 | else: |
| 128 | selected_per_class[label[i]] += 1 |
| 129 | |
| 130 | return boxes_tensor[selected_idx], label[selected_idx] |
| 131 | |
| 132 | |
| 133 | def main(args): |