Commit
•
56697c4
1
Parent(s):
3670d91
Fix CoreML P6 inference (#6700)
Browse files* 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>
- models/common.py +2 -1
models/common.py
CHANGED
@@ -428,7 +428,8 @@ class DetectMultiBackend(nn.Module):
|
|
428 |
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
|
429 |
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
|
430 |
else:
|
431 |
-
|
|
|
432 |
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
433 |
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
434 |
if self.saved_model: # SavedModel
|
|
|
428 |
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
|
429 |
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
|
430 |
else:
|
431 |
+
k = 'var_' + str(sorted(int(k.replace('var_', '')) for k in y)[-1]) # output key
|
432 |
+
y = y[k] # output
|
433 |
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
434 |
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
435 |
if self.saved_model: # SavedModel
|