Skip to content

Commit

Permalink
OpenVINO metadata fix2 (#7954)
Browse files Browse the repository at this point in the history
* Bug Fixed: OpenVINO metadata

* Bug Fixed: OpenVINO metadata

* Update export.py

* Update export.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
xylieong and glenn-jocher committed May 24, 2022
1 parent 0dd66e2 commit 27911dc
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion export.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def export_openvino(model, file, half, prefix=colorstr('OpenVINO:')):

cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
subprocess.check_output(cmd.split()) # export
with open(Path(f) / file.with_suffix('.yaml'), 'w') as g:
with open(Path(f) / file.with_suffix('.yaml').name, 'w') as g:
yaml.dump({'stride': int(max(model.stride)), 'names': model.names}, g) # add metadata.yaml

LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
Expand Down
2 changes: 1 addition & 1 deletion models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -368,7 +368,7 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False,
network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin'))
executable_network = ie.compile_model(model=network, device_name="CPU")
output_layer = next(iter(executable_network.outputs))
meta = w.with_suffix('.yaml')
meta = Path(w).with_suffix('.yaml')
if meta.exists():
stride, names = self._load_metadata(meta) # load metadata
elif engine: # TensorRT
Expand Down

0 comments on commit 27911dc

Please sign in to comment.