Skip to content

Commit

Permalink
fix for predicting with batch_size > 1 (#47)
Browse files Browse the repository at this point in the history
* fix for predicting with batch_size > 1
  • Loading branch information
tfriedel authored May 26, 2024
1 parent 7a8f964 commit 799ff3b
Showing 1 changed file with 2 additions and 4 deletions.
6 changes: 2 additions & 4 deletions ultralytics/models/yolov10/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ def postprocess(self, preds, img, orig_imgs):

if isinstance(preds, (list, tuple)):
preds = preds[0]

if preds.shape[-1] == 6:
pass
else:
Expand All @@ -22,9 +22,7 @@ def postprocess(self, preds, img, orig_imgs):

mask = preds[..., 4] > self.args.conf

b, _, c = preds.shape
preds = preds.view(-1, preds.shape[-1])[mask.view(-1)]
preds = preds.view(b, -1, c)
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
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
Expand Down

0 comments on commit 799ff3b

Please sign in to comment.