Browse Source

update o2o

yjh0410 9 tháng trước cách đây
mục cha
commit
20796ad3c6
1 tập tin đã thay đổi với 5 bổ sung0 xóa
  1. 5 0
      yolo/models/yolov10/yolov10.py

+ 5 - 0
yolo/models/yolov10/yolov10.py

@@ -120,6 +120,11 @@ class Yolov10(nn.Module):
         scores, labels, bboxes = multiclass_nms(
             scores, labels, bboxes, self.nms_thresh, self.num_classes)
         
+        # keep top-300 results
+        scores = scores[:300]
+        bboxes = bboxes[:300]
+        labels = labels[:300]
+        
         return bboxes, scores, labels
     
     def forward(self, x):