Fix NMS Return Index (#2196)

* fix nms index return * fix
open-mmlab · Jun 25, 2023 · 2f2ada9 · 2f2ada9
1 parent 90f17f2
commit 2f2ada9
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mmdeploy/mmcv/ops/nms.py b/mmdeploy/mmcv/ops/nms.py
@@ -345,7 +345,11 @@ def _multiclass_nms_single(boxes: Tensor,
     labels = labels[:, topk_inds, ...]
 
     if output_index:
-        bbox_index = pre_topk_inds[None, topk_inds]
+        bbox_index = box_inds.unsqueeze(0)
+        if pre_top_k > 0:
+            bbox_index = pre_topk_inds[None, box_inds]
+        if keep_top_k > 0:
+            bbox_index = bbox_index[:, topk_inds[:-1]]
         return dets, labels, bbox_index
     else:
         return dets, labels