Skip to content

Commit

Permalink
Add label ids to anomaly OpenVINO model xml (#2590)
Browse files Browse the repository at this point in the history
* Add label ids to model xml

---------
  • Loading branch information
ashwinvaidya17 authored Oct 31, 2023
1 parent 06775f2 commit 8598b05
Showing 1 changed file with 14 additions and 1 deletion.
15 changes: 14 additions & 1 deletion src/otx/algorithms/anomaly/tasks/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,7 +359,20 @@ def _add_metadata_to_ir(self, model_file: str, export_type: ExportType) -> None:

extra_model_data[("model_info", "reverse_input_channels")] = False
extra_model_data[("model_info", "model_type")] = "AnomalyDetection"
extra_model_data[("model_info", "labels")] = "Normal Anomaly"

labels = []
label_ids = []
for label_entity in self.task_environment.label_schema.get_labels(include_empty=False):
label_name = label_entity.name.replace(" ", "_")
# There is a mismatch between labels in OTX and modelAPI
if label_name == "Anomalous":
label_name = "Anomaly"
labels.append(label_name)
label_ids.append(str(label_entity.id_))

extra_model_data[("model_info", "labels")] = " ".join(labels)
extra_model_data[("model_info", "label_ids")] = " ".join(label_ids)

if export_type == ExportType.OPENVINO:
embed_ir_model_data(model_file, extra_model_data)
elif export_type == ExportType.ONNX:
Expand Down

0 comments on commit 8598b05

Please sign in to comment.