Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ultralytics Assistant 2024-09-06 03:54:35 +08:00 committed by GitHub
parent 95d54828bb
commit ac2c2be8f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 45 additions and 62 deletions

View file

@ -53,7 +53,7 @@ class ClassificationPredictor(BasePredictor):
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 pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
results.append(Results(orig_img, path=img_path, names=self.model.names, probs=pred))
return results
return [
Results(orig_img, path=img_path, names=self.model.names, probs=pred)
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0])
]