Skip to content

Commit

Permalink
Add dumping onnx model to model entity
Browse files Browse the repository at this point in the history
  • Loading branch information
sovrasov committed Apr 6, 2023
1 parent aed5f0f commit 4111f06
Show file tree
Hide file tree
Showing 8 changed files with 33 additions and 7 deletions.
3 changes: 3 additions & 0 deletions otx/algorithms/action/tasks/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -362,10 +362,13 @@ def export(
exporter.export()
bin_file = [f for f in os.listdir(self._output_path) if f.endswith(".bin")][0]
xml_file = [f for f in os.listdir(self._output_path) if f.endswith(".xml")][0]
onnx_file = [f for f in os.listdir(self._output_path) if f.endswith(".onnx")][0]
with open(os.path.join(self._output_path, bin_file), "rb") as f:
output_model.set_data("openvino.bin", f.read())
with open(os.path.join(self._output_path, xml_file), "rb") as f:
output_model.set_data("openvino.xml", f.read())
with open(onnx_file, "rb") as file:
output_model.set_data("model.onnx", file.read())
output_model.set_data(
"confidence_threshold", np.array([self.confidence_threshold], dtype=np.float32).tobytes()
)
Expand Down
2 changes: 2 additions & 0 deletions otx/algorithms/anomaly/tasks/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,8 @@ def export(
output_model.set_data("openvino.bin", file.read())
with open(xml_file, "rb") as file:
output_model.set_data("openvino.xml", file.read())
with open(onnx_path, "rb") as file:
output_model.set_data("model.onnx", file.read())

output_model.precision = self.precision
output_model.optimization_methods = self.optimization_methods
Expand Down
7 changes: 5 additions & 2 deletions otx/algorithms/classification/tasks/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,19 +249,22 @@ def export(

bin_file = outputs.get("bin")
xml_file = outputs.get("xml")
onnx_file = outputs.get("onnx")

inference_config = get_cls_inferencer_configuration(self._task_environment.label_schema)
deploy_cfg = get_cls_deploy_config(self._task_environment.label_schema, inference_config)
ir_extra_data = get_cls_model_api_configuration(self._task_environment.label_schema, inference_config)
ir_extra_data[("otx_config",)] = json.dumps(deploy_cfg, ensure_ascii=False)
embed_ir_model_data(xml_file, ir_extra_data)

if xml_file is None or bin_file is None:
raise RuntimeError("invalid status of exporting. bin and xml should not be None")
if xml_file is None or bin_file is None or onnx_file is None:
raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
with open(bin_file, "rb") as f:
output_model.set_data("openvino.bin", f.read())
with open(xml_file, "rb") as f:
output_model.set_data("openvino.xml", f.read())
with open(onnx_file, "rb") as f:
output_model.set_data("model.onnx", f.read())
output_model.precision = self._precision
output_model.has_xai = dump_features
output_model.set_data(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs): # noqa: C901
"outputs": {
"bin": os.path.join(cfg.work_dir, f"{model_name}.bin"),
"xml": os.path.join(cfg.work_dir, f"{model_name}.xml"),
"onnx": os.path.join(cfg.work_dir, f"{model_name}.onnx"),
"partitioned": [
{
"bin": os.path.join(cfg.work_dir, name.replace(".xml", ".bin")),
Expand Down
7 changes: 5 additions & 2 deletions otx/algorithms/detection/adapters/mmdet/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -513,18 +513,21 @@ def export(

bin_file = outputs.get("bin")
xml_file = outputs.get("xml")
onnx_file = outputs.get("onnx")

ir_extra_data = get_det_model_api_configuration(
self._task_environment.label_schema, self._task_type, self.confidence_threshold
)
embed_ir_model_data(xml_file, ir_extra_data)

if xml_file is None or bin_file is None:
raise RuntimeError("invalid status of exporting. bin and xml should not be None")
if xml_file is None or bin_file is None or onnx_file is None:
raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
with open(bin_file, "rb") as f:
output_model.set_data("openvino.bin", f.read())
with open(xml_file, "rb") as f:
output_model.set_data("openvino.xml", f.read())
with open(onnx_file, "rb") as f:
output_model.set_data("model.onnx", f.read())
output_model.set_data(
"confidence_threshold",
np.array([self.confidence_threshold], dtype=np.float32).tobytes(),
Expand Down
8 changes: 6 additions & 2 deletions otx/algorithms/segmentation/tasks/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,12 +175,16 @@ def export(

bin_file = outputs.get("bin")
xml_file = outputs.get("xml")
if xml_file is None or bin_file is None:
raise RuntimeError("invalid status of exporting. bin and xml should not be None")
onnx_file = outputs.get("onnx")

if xml_file is None or bin_file is None or onnx_file is None:
raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
with open(bin_file, "rb") as f:
output_model.set_data("openvino.bin", f.read())
with open(xml_file, "rb") as f:
output_model.set_data("openvino.xml", f.read())
with open(onnx_file, "rb") as f:
output_model.set_data("model.onnx", f.read())
output_model.precision = self._precision
output_model.optimization_methods = self._optimization_methods
output_model.has_xai = dump_features
Expand Down
1 change: 1 addition & 0 deletions tests/test_suite/run_test_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ def otx_export_testing(template, root, dump_features=False, half_precision=False
path_to_xml = os.path.join(save_path, "openvino.xml")
assert os.path.exists(path_to_xml)
assert os.path.exists(os.path.join(save_path, "openvino.bin"))
assert os.path.exists(os.path.join(save_path, "model.onnx"))
assert os.path.exists(os.path.join(save_path, "label_schema.json"))

if dump_features:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,16 +131,25 @@ def test_export_with_model_files(self, mocker, precision: ModelPrecision, dump_f
f.write(b"foo")
with open(f"{self.output_path}/model.bin", "wb") as f:
f.write(b"bar")
with open(f"{self.output_path}/model.onnx", "wb") as f:
f.write(b"foo")

mocker.patch("otx.algorithms.classification.tasks.inference.embed_ir_model_data", return_value=None)
fake_output = {"outputs": {"bin": f"{self.output_path}/model.xml", "xml": f"{self.output_path}/model.bin"}}
fake_output = {
"outputs": {
"bin": f"{self.output_path}/model.xml",
"xml": f"{self.output_path}/model.bin",
"onnx": f"{self.output_path}/model.onnx",
}
}
mock_run_task = mocker.patch.object(BaseTask, "_run_task", return_value=fake_output)
self.cls_inference_task.export(ExportType.OPENVINO, self.model, precision, dump_features)

mock_run_task.assert_called_once()
assert self.model.has_xai == dump_features
assert self.model.get_data("openvino.bin")
assert self.model.get_data("openvino.xml")
assert self.model.get_data("model.onnx")

@e2e_pytest_unit
def test_unload(self, mocker):
Expand Down

0 comments on commit 4111f06

Please sign in to comment.