From 551f639522f6aeb3ebb9ce0aeac92eab9b2be2fa Mon Sep 17 00:00:00 2001
From: "Tybulewicz, Tomasz" <tomasz.tybulewicz@intel.com>
Date: Wed, 1 Oct 2025 14:03:12 +0200
Subject: [PATCH] Option to save test results as images

---
 tests/accuracy/conftest.py      |  6 ++++++
 tests/accuracy/test_accuracy.py | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/tests/accuracy/conftest.py b/tests/accuracy/conftest.py
index edf47a9f..b7fc4ed4 100644
--- a/tests/accuracy/conftest.py
+++ b/tests/accuracy/conftest.py
@@ -28,6 +28,12 @@ def pytest_addoption(parser):
         default=False,
         help="whether to dump results into json file",
     )
+    parser.addoption(
+        "--results-dir",
+        action="store",
+        default="",
+        help="directory to store inference result",
+    )
 
 
 def pytest_configure(config):
diff --git a/tests/accuracy/test_accuracy.py b/tests/accuracy/test_accuracy.py
index eae25c65..4cc748f6 100644
--- a/tests/accuracy/test_accuracy.py
+++ b/tests/accuracy/test_accuracy.py
@@ -50,6 +50,7 @@
     InstanceSegmentationTiler,
     SemanticSegmentationTiler,
 )
+from model_api.visualizer import Visualizer
 
 # Mapping of model type strings to actual classes for security
 MODEL_TYPE_MAPPING = {
@@ -111,6 +112,11 @@ def data(pytestconfig):
     return pytestconfig.getoption("data")
 
 
+@pytest.fixture(scope="session")
+def results_dir(pytestconfig):
+    return pytestconfig.getoption("results_dir")
+
+
 @pytest.fixture(scope="session")
 def device(pytestconfig):
     return pytestconfig.getoption("device")
@@ -139,7 +145,7 @@ def pytest_generate_tests(metafunc):
         metafunc.parametrize("model_data", config_data)
 
 
-def test_image_models(data, device, dump, result, model_data):  # noqa: C901
+def test_image_models(data, device, dump, result, model_data, results_dir):  # noqa: C901
     name = model_data["name"]
     if name.endswith((".xml", ".onnx")):
         name = f"{data}/{name}"
@@ -235,6 +241,9 @@ def test_image_models(data, device, dump, result, model_data):  # noqa: C901
                     )
             else:
                 outputs = model(image)
+
+            store_outputs(name, image, device, outputs, results_dir)
+
             if isinstance(outputs, ClassificationResult) or type(outputs) is DetectionResult:
                 assert len(test_data["reference"]) == 1
                 output_str = str(outputs)
@@ -295,3 +304,23 @@ def test_image_models(data, device, dump, result, model_data):  # noqa: C901
 
     if dump:
         result[-1]["test_data"] = inference_results
+
+
+def store_outputs(name, image, device, result, results_dir: str) -> None:
+    if not results_dir:
+        return
+
+    Path(results_dir).mkdir(exist_ok=True, parents=True)
+
+    iteration = 1
+    while True:
+        path = Path(results_dir) / f"{Path(name).stem}_{iteration}_{device}.png"
+        if not path.exists():
+            break
+        iteration += 1
+
+    visualizer = Visualizer()
+    try:
+        visualizer.save(image, result, path)
+    except (TypeError, ValueError) as e:
+        print(f"Cannot save the output visualization for {name}. Error: {e}")