feat(train): add simple mlflow output

guillaume-chervet · May 24, 2023 · 2f90295 · 2f90295
1 parent 13ba04b
commit 2f90295
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/train/evaluate/azureml_step.py b/train/evaluate/azureml_step.py
@@ -18,8 +18,13 @@ def evaluate_step(
     integration_output: Output(type=URI_FOLDER),
 ):
     from evaluate import evaluate
+    import mlflow
 
-    evaluate(logging, Path(model_input), Path(images_input),
+    statistics = evaluate(logging, Path(model_input), Path(images_input),
              Path(model_output), Path(integration_output))
 
+    mlflow.log_metric("ok", statistics["ok"])
+    mlflow.log_metric("ko", statistics["ko"])
+    mlflow.log_metric("total", statistics["total"])
+
 
diff --git a/train/evaluate/evaluate.py b/train/evaluate/evaluate.py
@@ -69,3 +69,4 @@ def evaluate(logging, input_model_directory: Path, input_images_directory: Path,
         truth_filename = Path(filename).stem + "_" + Path(filename).suffix.lower().replace(".", '') + ".json"
         with open(ground_truth_directory / truth_filename, 'w') as file_stream:
             json.dump(data, file_stream, indent=4)
+    return statistics