xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix onnx validation (#179)

Browse Source
This commit is contained in:
leonnil 2024-06-04 10:29:21 +00:00
parent e5f2f03c55
commit ba0e6cbc1a
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
ultralytics/models/yolov10/val.py
View File

@ -14,7 +14,10 @@ class YOLOv10DetectionValidator(DetectionValidator):
if isinstance(preds, (list, tuple)):
preds = preds[0]
preds = preds.transpose(-1, -2)
boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
bboxes = ops.xywh2xyxy(boxes)
return torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
if preds.shape[-1] == 6:
return preds
else:
preds = preds.transpose(-1, -2)
boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
bboxes = ops.xywh2xyxy(boxes)
return torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)