Non-maximum suppression (NMS). Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` scores: prediction scores of the boxes, sized (N,). This function keeps boxes with higher scores. nms_thresh: threshold of
(
boxes: NdarrayOrTensor,
scores: NdarrayOrTensor,
nms_thresh: float,
max_proposals: int = -1,
box_overlap_metric: Callable = box_iou,
)
| 1085 | |
| 1086 | |
| 1087 | def non_max_suppression( |
| 1088 | boxes: NdarrayOrTensor, |
| 1089 | scores: NdarrayOrTensor, |
| 1090 | nms_thresh: float, |
| 1091 | max_proposals: int = -1, |
| 1092 | box_overlap_metric: Callable = box_iou, |
| 1093 | ) -> NdarrayOrTensor: |
| 1094 | """ |
| 1095 | Non-maximum suppression (NMS). |
| 1096 | |
| 1097 | Args: |
| 1098 | boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` |
| 1099 | scores: prediction scores of the boxes, sized (N,). This function keeps boxes with higher scores. |
| 1100 | nms_thresh: threshold of NMS. Discards all overlapping boxes with box_overlap > nms_thresh. |
| 1101 | max_proposals: maximum number of boxes it keeps. |
| 1102 | If ``max_proposals`` = -1, there is no limit on the number of boxes that are kept. |
| 1103 | box_overlap_metric: the metric to compute overlap between boxes. |
| 1104 | |
| 1105 | Returns: |
| 1106 | Indexes of ``boxes`` that are kept after NMS. |
| 1107 | |
| 1108 | Example: |
| 1109 | .. code-block:: python |
| 1110 | |
| 1111 | boxes = torch.ones(10,6) |
| 1112 | scores = torch.ones(10) |
| 1113 | keep = non_max_suppression(boxes, scores, num_thresh=0.1) |
| 1114 | boxes_after_nms = boxes[keep] |
| 1115 | """ |
| 1116 | |
| 1117 | # returns empty array if boxes is empty |
| 1118 | if boxes.shape[0] == 0: |
| 1119 | return convert_to_dst_type(src=np.array([]), dst=boxes, dtype=torch.long)[0] |
| 1120 | |
| 1121 | if boxes.shape[0] != scores.shape[0]: |
| 1122 | raise ValueError( |
| 1123 | f"boxes and scores should have same length, got boxes shape {boxes.shape}, scores shape {scores.shape}" |
| 1124 | ) |
| 1125 | |
| 1126 | # convert numpy to tensor if needed |
| 1127 | boxes_t, *_ = convert_data_type(boxes, torch.Tensor) |
| 1128 | scores_t, *_ = convert_to_dst_type(scores, boxes_t) |
| 1129 | |
| 1130 | # sort boxes in descending order according to the scores |
| 1131 | # use stable=True to ensure deterministic ordering when scores are equal |
| 1132 | sort_idxs = torch.argsort(scores_t, dim=0, descending=True, stable=True) |
| 1133 | boxes_sort = deepcopy(boxes_t)[sort_idxs, :] |
| 1134 | |
| 1135 | # initialize the list of picked indexes |
| 1136 | pick = [] |
| 1137 | idxs = torch.Tensor(list(range(boxes_sort.shape[0]))).to(device=boxes_t.device, dtype=torch.long) |
| 1138 | |
| 1139 | # keep looping while some indexes still remain in the indexes list |
| 1140 | while len(idxs) > 0: |
| 1141 | # pick the first index in the indexes list and add the index value to the list of picked indexes |
| 1142 | i = int(idxs[0].item()) |
| 1143 | pick.append(i) |
| 1144 | if len(pick) >= max_proposals >= 1: |
searching dependent graphs…