resolve conflict

This commit is contained in:
wa22 2024-06-04 19:15:37 +08:00
parent dbf759d925
commit 975ef03f84

View File

@ -14,12 +14,11 @@ class YOLOv10DetectionValidator(DetectionValidator):
if isinstance(preds, (list, tuple)): if isinstance(preds, (list, tuple)):
preds = preds[0] preds = preds[0]
# Acknowledgement: Thanks to sanha9999 in #190 and #181!
if preds.shape[-1] == 6: if preds.shape[-1] == 6:
pass return preds
else: else:
preds = preds.transpose(-1, -2) preds = preds.transpose(-1, -2)
bboxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, preds.shape[-1]-4) boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
bboxes = ops.xywh2xyxy(bboxes) bboxes = ops.xywh2xyxy(boxes)
preds = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1) return torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
return preds