🔧 [Add] MaxBox for NMS output

yolo/config/config.py

@@ -102,8 +102,9 @@ class EMAConfig:
 
 @dataclass
 class NMSConfig:
-    min_confidence: int
-    min_iou: int
+    min_confidence: float
+    min_iou: float
+    max_bbox: int
 
 
 @dataclass

yolo/config/task/inference.yaml

@@ -8,4 +8,5 @@ data:
 nms:
   min_confidence: 0.5
   min_iou: 0.5
+  max_bbox: 300
 save_predict: True

yolo/config/task/validation.yaml

@@ -10,3 +10,4 @@ data:
 nms:
   min_confidence: 0.0001
   min_iou: 0.7
+  max_bbox: 1000

yolo/utils/bounding_box_utils.py

@@ -424,7 +424,7 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Opt
             [valid_cls[instance_idx][:, None], valid_box[instance_idx], valid_con[instance_idx][:, None]], dim=-1
         )
 
-        predicts_nms.append(predict_nms)
+        predicts_nms.append(predict_nms[: nms_cfg.max_bbox])
     return predicts_nms