Skip to content

Commit

Permalink
Fix CoreML P6 inference (ultralytics#6700)
Browse files Browse the repository at this point in the history
* Fix CoreML P6 inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
2 people authored and eladco committed Mar 10, 2022
1 parent 13ea477 commit 385c799
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -428,7 +428,8 @@ def forward(self, im, augment=False, visualize=False, val=False):
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
else:
y = y[sorted(y)[-1]] # last output
k = 'var_' + str(sorted(int(k.replace('var_', '')) for k in y)[-1]) # output key
y = y[k] # output
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
if self.saved_model: # SavedModel
Expand Down

0 comments on commit 385c799

Please sign in to comment.