diff --git a/ultralytics/models/yolov10/predict.py b/ultralytics/models/yolov10/predict.py index 29f21a9c..e3670f20 100644 --- a/ultralytics/models/yolov10/predict.py +++ b/ultralytics/models/yolov10/predict.py @@ -12,32 +12,25 @@ class YOLOv10DetectionPredictor(DetectionPredictor): if isinstance(preds, (list, tuple)): preds = preds[0] - if preds.shape[-1] != 6: + 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, 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 - ) + preds = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1) mask = preds[..., 4] > self.args.conf - # Filter predictions using the mask and keep batch dimension - filtered_preds = [p[mask[idx]] for idx, p in enumerate(preds)] + preds = [p[mask[idx]] for idx, p in enumerate(preds)] - if not isinstance( - orig_imgs, list - ): # input images are a torch.Tensor, not a list + if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) results = [] - for i, pred in enumerate(filtered_preds): + for i, pred in enumerate(preds): orig_img = orig_imgs[i] pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) img_path = self.batch[0][i] - results.append( - Results(orig_img, path=img_path, names=self.model.names, boxes=pred) - ) + results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred)) return results