Repository: Project-MONAI/MONAI

Issue:
Box selector
**Is your feature request related to a problem? Please describe.**
A box selector for detection postprocessing is required

**Describe the solution you'd like**
Add box selector

**Describe alternatives you've considered**
N/A

**Additional context**
1. Top box selector
2. Non-maximum suppression selector
add batch_nms to box util
**Is your feature request related to a problem? Please describe.**
add batch_nms to box util

**Describe the solution you'd like**
- Please refer to batch_nms in
https://github.com/pytorch/vision/blob/main/torchvision/ops/boxes.py
- update NMS in detector into batch_nms

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.



Discussion:

Resolving patch:
diff --git a/monai/apps/detection/utils/box_coder.py b/monai/apps/detection/utils/box_coder.py
--- a/monai/apps/detection/utils/box_coder.py
+++ b/monai/apps/detection/utils/box_coder.py
@@ -56,15 +56,7 @@
 import torch
 from torch import Tensor
 
-from monai.data.box_utils import (
-    COMPUTE_DTYPE,
-    CenterSizeMode,
-    CornerCornerModeTypeB,
-    StandardMode,
-    convert_box_mode,
-    convert_box_to_standard_mode,
-    is_valid_box_values,
-)
+from monai.data.box_utils import COMPUTE_DTYPE, CenterSizeMode, StandardMode, convert_box_mode, is_valid_box_values
 from monai.utils.module import look_up_option
 
 
@@ -207,27 +199,27 @@ def decode_single(self, rel_codes: Tensor, reference_boxes: Tensor) -> Tensor:
         From a set of original boxes and encoded relative box offsets,
 
         Args:
-            rel_codes: encoded boxes, Nx4 or Nx6 torch tensor.
+            rel_codes: encoded boxes, Nx(4*num_box_reg) or Nx(6*num_box_reg) torch tensor.
             reference_boxes: reference boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode``
 
         Return:
-            decoded boxes, Nx4 or Nx6 torch tensor. The box mode will to be ``StandardMode``
+            decoded boxes, Nx(4*num_box_reg) or Nx(6*num_box_reg) torch tensor. The box mode will to be ``StandardMode``
         """
         reference_boxes = reference_boxes.to(rel_codes.dtype)
+        offset = reference_boxes.shape[-1]
 
         pred_boxes = []
         boxes_cccwhd = convert_box_mode(reference_boxes, src_mode=StandardMode, dst_mode=CenterSizeMode)
         for axis in range(self.spatial_dims):
             whd_axis = boxes_cccwhd[:, axis + self.spatial_dims]
             ctr_xyz_axis = boxes_cccwhd[:, axis]
-            dxyz_axis = rel_codes[:, axis] / self.weights[axis]
-            dwhd_axis = rel_codes[:, self.spatial_dims + axis] / self.weights[axis + self.spatial_dims]
-
+            dxyz_axis = rel_codes[:, axis::offset] / self.weights[axis]
+            dwhd_axis = rel_codes[:, self.spatial_dims + axis :: offset] / self.weights[axis + self.spatial_dims]
             # Prevent sending too large values into torch.exp()
             dwhd_axis = torch.clamp(dwhd_axis.to(COMPUTE_DTYPE), max=self.boxes_xform_clip)
 
-            pred_ctr_xyx_axis = dxyz_axis * whd_axis + ctr_xyz_axis
-            pred_whd_axis = torch.exp(dwhd_axis) * whd_axis
+            pred_ctr_xyx_axis = dxyz_axis * whd_axis[:, None] + ctr_xyz_axis[:, None]
+            pred_whd_axis = torch.exp(dwhd_axis) * whd_axis[:, None]
             pred_whd_axis = pred_whd_axis.to(dxyz_axis.dtype)
 
             # When convert float32 to float16, Inf or Nan may occur
@@ -242,5 +234,6 @@ def decode_single(self, rel_codes: Tensor, reference_boxes: Tensor) -> Tensor:
             pred_boxes.append(pred_ctr_xyx_axis - c_to_c_whd_axis)
             pred_boxes.append(pred_ctr_xyx_axis + c_to_c_whd_axis)
 
-        pred_boxes_xxyyzz = torch.stack(pred_boxes, dim=1)
-        return convert_box_to_standard_mode(pred_boxes_xxyyzz, mode=CornerCornerModeTypeB)  # type: ignore
+        pred_boxes = pred_boxes[::2] + pred_boxes[1::2]
+        pred_boxes_final = torch.stack(pred_boxes, dim=2).flatten(1)
+        return pred_boxes_final
diff --git a/monai/apps/detection/utils/box_selector.py b/monai/apps/detection/utils/box_selector.py
new file mode 100644
--- /dev/null
+++ b/monai/apps/detection/utils/box_selector.py
@@ -0,0 +1,216 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# =========================================================================
+# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py
+# which has the following license...
+# https://github.com/pytorch/vision/blob/main/LICENSE
+
+# BSD 3-Clause License
+
+# Copyright (c) Soumith Chintala 2016,
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+
+"""
+Part of this script is adapted from
+https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py
+"""
+
+from typing import Callable, List, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from monai.data.box_utils import batched_nms, box_iou, clip_boxes_to_image
+from monai.transforms.utils_pytorch_numpy_unification import floor_divide
+
+
+class BoxSelector:
+    """
+    Box selector which selects the predicted boxes.
+    The box selection is performed with the following steps:
+
+    #. For each level, discard boxes with scores less than self.score_thresh.
+    #. For each level, keep boxes with top self.topk_candidates_per_level scores.
+    #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overapping threshold nms_thresh.
+    #. For the whole image, keep boxes with top self.detections_per_img scores.
+
+    Args:
+        apply_sigmoid: whether to apply sigmoid to get scores from classification logits
+        score_thresh: no box with scores less than score_thresh will be kept
+        topk_candidates_per_level: max number of boxes to keep for each level
+        nms_thresh: box overlapping threshold for NMS
+        detections_per_img: max number of boxes to keep for each image
+
+    Example:
+
+        .. code-block:: python
+
+            input_param = {
+                "apply_sigmoid": True,
+                "score_thresh": 0.1,
+                "topk_candidates_per_level": 2,
+                "nms_thresh": 0.1,
+                "detections_per_img": 5,
+            }
+            box_selector = BoxSelector(**input_param)
+            boxes = [torch.randn([3,6]), torch.randn([7,6])]
+            logits = [torch.randn([3,3]), torch.randn([7,3])]
+            spatial_size = (8,8,8)
+            selected_boxes, selected_scores, selected_labels = box_selector.select_boxes_per_image(
+                boxes, logits, spatial_size
+            )
+    """
+
+    def __init__(
+        self,
+        box_overlap_metric: Callable = box_iou,
+        apply_sigmoid: bool = True,
+        score_thresh=0.05,
+        topk_candidates_per_level=1000,
+        nms_thresh=0.5,
+        detections_per_img=300,
+    ):
+        self.box_overlap_metric = box_overlap_metric
+
+        self.apply_sigmoid = apply_sigmoid
+        self.score_thresh = score_thresh
+        self.topk_candidates_per_level = topk_candidates_per_level
+        self.nms_thresh = nms_thresh
+        self.detections_per_img = detections_per_img
+
+    def select_top_score_idx_per_level(self, logits: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Select indice with highest scores.
+
+        The indice selection is performed with the following steps:
+
+        #. If self.apply_sigmoid, get scores by applying sigmoid to logits. Otherwise, use logits as scores.
+        #. Discard indice with scores less than self.score_thresh
+        #. Keep indice with top self.topk_candidates_per_level scores
+
+        Args:
+            logits: predicted classification logits, Tensor sized (N, num_classes)
+
+        Return:
+            - topk_idxs: selected M indices, Tensor sized (M, )
+            - selected_scores: selected M scores, Tensor sized (M, )
+            - selected_labels: selected M labels, Tensor sized (M, )
+        """
+        num_classes = logits.shape[-1]
+
+        # apply sigmoid to classification logits if asked
+        if self.apply_sigmoid:
+            scores = torch.sigmoid(logits.to(torch.float32)).flatten()
+        else:
+            scores = logits.flatten()
+
+        # remove low scoring boxes
+        keep_idxs = scores > self.score_thresh
+        scores = scores[keep_idxs]
+        flatten_topk_idxs = torch.where(keep_idxs)[0]
+
+        # keep only topk scoring predictions
+        num_topk = min(self.topk_candidates_per_level, flatten_topk_idxs.size(0))
+        selected_scores, idxs = scores.to(torch.float32).topk(
+            num_topk
+        )  # half precision not implemented for cpu float16
+        flatten_topk_idxs = flatten_topk_idxs[idxs]
+
+        selected_labels = flatten_topk_idxs % num_classes
+
+        topk_idxs = floor_divide(flatten_topk_idxs, num_classes)
+        return topk_idxs, selected_scores, selected_labels  # type: ignore
+
+    def select_boxes_per_image(
+        self, boxes_list: List[Tensor], logits_list: List[Tensor], spatial_size: Union[List[int], Tuple[int]]
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Postprocessing to generate detection result from classification logits and boxes.
+
+        The box selection is performed with the following steps:
+
+        #. For each level, discard boxes with scores less than self.score_thresh.
+        #. For each level, keep boxes with top self.topk_candidates_per_level scores.
+        #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overapping threshold nms_thresh.
+        #. For the whole image, keep boxes with top self.detections_per_img scores.
+
+        Args:
+            boxes_list: list of predicted boxes from a single image,
+                each element i is a Tensor sized (N_i, 2*spatial_dims)
+            logits_list: list of predicted classification logits from a single image,
+                each element i is a Tensor sized (N_i, num_classes)
+            spatial_size: spatial size of the image
+
+        Return:
+            - selected boxes, Tensor sized (P, 2*spatial_dims)
+            - selected_scores, Tensor sized (P, )
+            - selected_labels, Tensor sized (P, )
+        """
+
+        if len(boxes_list) != len(logits_list):
+            raise ValueError(
+                "len(boxes_list) should equal to len(logits_list). "
+                f"Got len(boxes_list)={len(boxes_list)}, len(logits_list)={len(logits_list)}"
+            )
+
+        image_boxes = []
+        image_scores = []
+        image_labels = []
+
+        boxes_dtype = boxes_list[0].dtype
+        logits_dtype = logits_list[0].dtype
+
+        for boxes_per_level, logits_per_level in zip(boxes_list, logits_list):
+            # select topk boxes for each level
+            topk_idxs: Tensor
+            topk_idxs, scores_per_level, labels_per_level = self.select_top_score_idx_per_level(logits_per_level)
+            boxes_per_level = boxes_per_level[topk_idxs]
+
+            keep: Tensor
+            boxes_per_level, keep = clip_boxes_to_image(boxes_per_level, spatial_size, remove_empty=True)  # type: ignore
+            image_boxes.append(boxes_per_level)
+            image_scores.append(scores_per_level[keep])
+            image_labels.append(labels_per_level[keep])
+
+        image_boxes_t: Tensor = torch.cat(image_boxes, dim=0)
+        image_scores_t: Tensor = torch.cat(image_scores, dim=0)
+        image_labels_t: Tensor = torch.cat(image_labels, dim=0)
+
+        # non-maximum suppression on detected boxes from all levels
+        keep_t: Tensor = batched_nms(  # type: ignore
+            image_boxes_t,
+            image_scores_t,
+            image_labels_t,
+            self.nms_thresh,
+            box_overlap_metric=self.box_overlap_metric,
+            max_proposals=self.detections_per_img,
+        )
+
+        selected_boxes = image_boxes_t[keep_t].to(boxes_dtype)
+        selected_scores = image_scores_t[keep_t].to(logits_dtype)
+        selected_labels = image_labels_t[keep_t]
+
+        return selected_boxes, selected_scores, selected_labels
diff --git a/monai/data/box_utils.py b/monai/data/box_utils.py
--- a/monai/data/box_utils.py
+++ b/monai/data/box_utils.py
@@ -1023,8 +1023,7 @@ def non_max_suppression(
     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 NMS. For boxes with overlap more than nms_thresh,
-            we only keep the one with the highest score.
+        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.
@@ -1046,7 +1045,7 @@ def non_max_suppression(
             f"boxes and scores should have same length, got boxes shape {boxes.shape}, scores shape {scores.shape}"
         )
 
-    # convert tensor to numpy if needed
+    # convert numpy to tensor if needed
     boxes_t, *_ = convert_data_type(boxes, torch.Tensor)
     scores_t, *_ = convert_to_dst_type(scores, boxes_t)
 
@@ -1077,5 +1076,53 @@ def non_max_suppression(
     # return only the bounding boxes that were picked using the integer data type
     pick_idx = sort_idxs[pick]
 
-    # convert numpy back to tensor if needed
+    # convert tensor back to numpy if needed
     return convert_to_dst_type(src=pick_idx, dst=boxes, dtype=pick_idx.dtype)[0]
+
+
+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]


Test patch:
diff --git a/tests/test_detector_boxselector.py b/tests/test_detector_boxselector.py
new file mode 100644
--- /dev/null
+++ b/tests/test_detector_boxselector.py
@@ -0,0 +1,65 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.apps.detection.utils.box_selector import BoxSelector
+from tests.utils import assert_allclose
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+num_anchors = 7
+
+TEST_CASE = []
+TEST_CASE.append(
+    [  # 2D
+        {
+            "apply_sigmoid": False,
+            "score_thresh": 0.1,
+            "topk_candidates_per_level": 2,
+            "nms_thresh": 0.1,
+            "detections_per_img": 5,
+        },
+        [torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 10]]).to(torch.float16)],
+        [torch.tensor([[0.1, 0.6], [0.2, 0.2]])],
+        (20, 20, 20),
+        torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 10]]),
+    ]
+)
+TEST_CASE.append(
+    [
+        {
+            "apply_sigmoid": False,
+            "score_thresh": 0.1,
+            "topk_candidates_per_level": 1,
+            "nms_thresh": 0.1,
+            "detections_per_img": 5,
+        },
+        [torch.tensor([[1, 2, 3, 2, 3, 4]]).to(torch.float32), torch.tensor([[5, 6, 7, 8, 9, 10]]).to(torch.float32)],
+        [torch.tensor([[0.3, 0.6]]), torch.tensor([[0.2, 0.2]])],
+        (20, 20, 8),
+        torch.tensor([[1, 2, 3, 2, 3, 4], [5, 6, 7, 8, 9, 8]]),
+    ]
+)
+
+
+class TestBoxSelector(unittest.TestCase):
+    @parameterized.expand(TEST_CASE)
+    def test_box_selector(self, input_param, boxes, logits, image_shape, expected_results):
+        box_selector = BoxSelector(**input_param)
+        result = box_selector.select_boxes_per_image(boxes, logits, image_shape)
+        assert_allclose(result[0], expected_results, type_test=True, device_test=False, atol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()