diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index d50039c54ed9..daa7189afdf6 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -56,6 +56,7 @@ from copy import deepcopy from pathlib import Path +import numpy as np import torch from ultralytics.nn.autobackend import check_class_names @@ -362,6 +363,16 @@ def export_openvino(self, prefix=colorstr('OpenVINO:')): model_name=self.pretty_name, framework='onnx', compress_to_fp16=self.args.half) # export + ov_model.set_rt_info("YoloV8", ["model_info", "model_type"]) + ov_model.set_rt_info("fit_to_window_letterbox", ["model_info", "resize_type"]) + ov_model.set_rt_info(True, ["model_info", "reverse_input_channels"]) + ov_model.set_rt_info(114, ["model_info", "pad_value"]) + ov_model.set_rt_info([255.0], ["model_info", "scale_values"]) + ov_model.set_rt_info(self.args.iou, ["model_info", "iou_threshold"]) + labels = [] + for i in range(len(self.model.names)): + labels.append(self.model.names[i].replace(" ", "_")) + ov_model.set_rt_info(labels, ["model_info", "labels"]) ov.serialize(ov_model, f_ov) # save yaml_save(Path(f) / 'metadata.yaml', self.metadata) # add metadata.yaml return f, None