Skip to content

Commit

Permalink
Add label ids to model xml in release 1.5 (#2591)
Browse files Browse the repository at this point in the history
Add label ids to model xml
  • Loading branch information
ashwinvaidya17 authored Nov 1, 2023
1 parent 23a4644 commit a5d12b7
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 a5d12b7

Please sign in to comment.