ultralytics 8.0.170 apply is_list fixes for torch.Tensor inputs (#4713)
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:
parent
a1c1d6b483
commit
aa9133bb88
15 changed files with 74 additions and 36 deletions
|
|
@ -37,10 +37,12 @@ class PosePredictor(DetectionPredictor):
|
|||
classes=self.args.classes,
|
||||
nc=len(self.model.names))
|
||||
|
||||
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]
|
||||
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round()
|
||||
pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
|
||||
pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue