Update nms_rotated_wrapper.py

过滤掉小box的score
little-white001 · Feb 15, 2022 · fd7dcf8 · fd7dcf8
1 parent 9d86536
commit fd7dcf8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils/nms_rotated/nms_rotated_wrapper.py b/utils/nms_rotated/nms_rotated_wrapper.py
@@ -36,6 +36,7 @@ def obb_nms(dets, scores, iou_thr, device_id=None):
             ori_inds = torch.arange(dets_th.size(0)) # 0 ~ n-1
             ori_inds = ori_inds[~too_small]
             dets_th = dets_th[~too_small] # (n_filter, 5)
+            scores = scores[~too_small] 
 
             inds = nms_rotated_ext.nms_rotated(dets_th, scores, iou_thr)
             inds = ori_inds[inds]
@@ -74,4 +75,4 @@ def poly_nms(dets, iou_thr, device_id=None):
                                     [136.6, 111.6, 200, 100, 120],
                                     [100, 100, 141.4, 141.4, 45],
                                     [100, 100, 141.4, 141.4, 135]))
-
+