HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.170 apply is_list fixes for torch.Tensor inputs (#4713)

Browse source 
Co-authored-by: Gezhi Zhang <765724965@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-09-04 02:22:05 +02:00 • committed by GitHub
parent a1c1d6b483
commit aa9133bb88
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 74 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

8
ultralytics/models/yolo/classify/predict.py
View file

@ -4,7 +4,7 @@ import torch
from ultralytics.engine.predictor import BasePredictor
from ultralytics.engine.results import Results
from ultralytics.utils import DEFAULT_CFG
from ultralytics.utils import DEFAULT_CFG, ops
class ClassificationPredictor(BasePredictor):
@ -38,10 +38,12 @@ class ClassificationPredictor(BasePredictor):
def postprocess(self, preds, img, orig_imgs):
"""Post-processes predictions to return Results objects."""
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, pred in enumerate(preds):
orig_img = orig_imgs[i] if is_list else orig_imgs
orig_img = orig_imgs[i]
img_path = self.batch[0][i]
results.append(Results(orig_img, path=img_path, names=self.model.names, probs=pred))
return results