🐛 [Update] NMS, enable multiclass result

yolo/utils/bounding_box_utils.py

@@ -411,15 +411,11 @@ def create_converter(model_version: str = "v9-c", *args, **kwargs) -> Union[Anc2
 def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor] = None):
     cls_dist = cls_dist.sigmoid() * (1 if confidence is None else confidence)
 
-    # filter class by confidence
-    cls_val, cls_idx = cls_dist.max(dim=-1, keepdim=True)
-    valid_mask = cls_val > nms_cfg.min_confidence
-    valid_cls = cls_idx[valid_mask].float()
-    valid_con = cls_val[valid_mask].float()
-    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)
+    batch_idx, valid_grid, valid_cls = torch.where(cls_dist > nms_cfg.min_confidence)
+    valid_con = cls_dist[batch_idx, valid_grid, valid_cls]
+    valid_box = bbox[batch_idx, valid_grid]
+
+    nms_idx = batched_nms(valid_box, valid_con, batch_idx + valid_cls * bbox.size(0), nms_cfg.min_iou)
     predicts_nms = []
     for idx in range(cls_dist.size(0)):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]