Function batched_nms

monai/data/box_utils.py:1162–1207

Performs non-maximum suppression in a batched fashion. Each labels value correspond to a category, and NMS will not be applied between elements of different categories. Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/nms.py Args: boxes: boun

(
    boxes: NdarrayOrTensor,
    scores: NdarrayOrTensor,
    labels: NdarrayOrTensor,
    nms_thresh: float,
    max_proposals: int = -1,
    box_overlap_metric: Callable = box_iou,
)

def batched_nms(
    boxes: NdarrayOrTensor,
    scores: NdarrayOrTensor,
    labels: NdarrayOrTensor,
    nms_thresh: float,
    max_proposals: int = -1,
    box_overlap_metric: Callable = box_iou,
) -> NdarrayOrTensor:
    """
    Performs non-maximum suppression in a batched fashion.
    Each labels value correspond to a category, and NMS will not be applied between elements of different categories.

    Adapted from https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/nms.py

    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.
        labels: indices of the categories for each one of the boxes. sized(N,), value range is (0, num_classes)
        nms_thresh: threshold of NMS. Discards all overlapping boxes with box_overlap > nms_thresh.
        max_proposals: maximum number of boxes it keeps.
            If ``max_proposals`` = -1, there is no limit on the number of boxes that are kept.
        box_overlap_metric: the metric to compute overlap between boxes.

    Returns:
        Indexes of ``boxes`` that are kept after NMS.
    """
    # returns empty array if boxes is empty
    if boxes.shape[0] == 0:
        return convert_to_dst_type(src=np.array([]), dst=boxes, dtype=torch.long)[0]

    # convert numpy to tensor if needed
    boxes_t, *_ = convert_data_type(boxes, torch.Tensor, dtype=torch.float32)
    scores_t, *_ = convert_to_dst_type(scores, boxes_t)
    labels_t, *_ = convert_to_dst_type(labels, boxes_t, dtype=torch.long)

    # strategy: in order to perform NMS independently per class.
    # we add an offset to all the boxes. The offset is dependent
    # only on the class idx, and is large enough so that boxes
    # from different classes do not overlap
    max_coordinate = boxes_t.max()
    offsets = labels_t.to(boxes_t) * (max_coordinate + 1)
    boxes_for_nms = boxes + offsets[:, None]
    keep = non_max_suppression(boxes_for_nms, scores_t, nms_thresh, max_proposals, box_overlap_metric)

    # convert tensor back to numpy if needed
    return convert_to_dst_type(src=keep, dst=boxes, dtype=keep.dtype)[0]