From cbf4e0db57f9809a70a005dbbdabea3f85907b18 Mon Sep 17 00:00:00 2001 From: Ilija Date: Thu, 10 Oct 2024 15:11:38 +0200 Subject: [PATCH 1/2] batched_nms using proper score values --- yolo/utils/bounding_box_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/utils/bounding_box_utils.py b/yolo/utils/bounding_box_utils.py index 12d95c583..c626bc1a4 100644 --- a/yolo/utils/bounding_box_utils.py +++ b/yolo/utils/bounding_box_utils.py @@ -388,7 +388,7 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt valid_box = bbox[valid_mask.repeat(1, 1, 4)].view(-1, 4) batch_idx, *_ = torch.where(valid_mask) - nms_idx = batched_nms(valid_box, valid_cls, batch_idx, nms_cfg.min_iou) + nms_idx = batched_nms(valid_box, valid_con, batch_idx, nms_cfg.min_iou) predicts_nms = [] for idx in range(cls_dist.size(0)): instance_idx = nms_idx[idx == batch_idx[nms_idx]] From a6f4502a2ed622c664a522378a754fb131a2ac3d Mon Sep 17 00:00:00 2001 From: Ilija Date: Thu, 10 Oct 2024 16:19:58 +0200 Subject: [PATCH 2/2] readability update --- yolo/utils/bounding_box_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/utils/bounding_box_utils.py b/yolo/utils/bounding_box_utils.py index c626bc1a4..06a3173d4 100644 --- a/yolo/utils/bounding_box_utils.py +++ b/yolo/utils/bounding_box_utils.py @@ -388,7 +388,7 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt valid_box = bbox[valid_mask.repeat(1, 1, 4)].view(-1, 4) batch_idx, *_ = torch.where(valid_mask) - nms_idx = batched_nms(valid_box, valid_con, batch_idx, nms_cfg.min_iou) + nms_idx = batched_nms(boxes=valid_box, scores=valid_con, idxs=batch_idx, iou_threshold=nms_cfg.min_iou) predicts_nms = [] for idx in range(cls_dist.size(0)): instance_idx = nms_idx[idx == batch_idx[nms_idx]]