Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Fix TorchScript JSON string key bug (#9015)

Browse Source 
* Fix TorchScript JSON string key bug

Resolves https://github.com/ultralytics/yolov5/issues/9011

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

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

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

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocherandpre-commit-ci[bot] authored and GitHub committed 2022-08-18 11:50:24 +02:00
1 parent 5c854fab5e
commit 529aafd737
1 file changed
+4 -2
Download Patch File Download Diff File Expand all files Collapse all files
models/common.py
+4 -2
View File
@@ -337,8 +337,10 @@ class DetectMultiBackend(nn.Module):
extra_files = {'config.txt': ''} # model metadata extra_files = {'config.txt': ''} # model metadata
model = torch.jit.load(w, _extra_files=extra_files) model = torch.jit.load(w, _extra_files=extra_files)
model.half() if fp16 else model.float() model.half() if fp16 else model.float()
if extra_files['config.txt']: if extra_files['config.txt']: # load metadata dict
d = json.loads(extra_files['config.txt']) # extra_files dict d = json.loads(extra_files['config.txt'],
object_hook=lambda d: {int(k) if k.isdigit() else k: v
for k, v in d.items()})
stride, names = int(d['stride']), d['names'] stride, names = int(d['stride']), d['names']
elif dnn: # ONNX OpenCV DNN elif dnn: # ONNX OpenCV DNN
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...') LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')