From dbf759d925bf724e6cf8d446f9cae18098b681b0 Mon Sep 17 00:00:00 2001 From: kangsanha <68986305+sanha9999@users.noreply.github.com> Date: Tue, 4 Jun 2024 19:08:46 +0900 Subject: [PATCH] fix onnx AssertionError --- ultralytics/models/yolov10/val.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ultralytics/models/yolov10/val.py b/ultralytics/models/yolov10/val.py index 1106b9eb..77450b44 100644 --- a/ultralytics/models/yolov10/val.py +++ b/ultralytics/models/yolov10/val.py @@ -14,7 +14,12 @@ 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) \ No newline at end of file + if preds.shape[-1] == 6: + pass + else: + preds = preds.transpose(-1, -2) + bboxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, preds.shape[-1]-4) + bboxes = ops.xywh2xyxy(bboxes) + preds = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1) + + return preds