Add post-process inference test

ankandrew · Apr 5, 2024 · 45755df · 45755df
1 parent 9005296
commit 45755df
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/test/fast_lp_ocr/inference/__init__.py b/test/fast_lp_ocr/inference/__init__.py
diff --git a/test/fast_lp_ocr/inference/test_process.py b/test/fast_lp_ocr/inference/test_process.py
@@ -0,0 +1,45 @@
+"""
+Tests for inference process module.
+"""
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+
+from fast_plate_ocr.inference.process import postprocess_output
+
+
+@pytest.mark.parametrize(
+    "model_output, max_plate_slots, model_alphabet, expected_plates",
+    [
+        (
+            np.array(
+                [
+                    [[0.5, 0.4, 0.1], [0.2, 0.6, 0.2], [0.1, 0.4, 0.5]],
+                    [[0.1, 0.1, 0.8], [0.2, 0.2, 0.6], [0.1, 0.4, 0.5]],
+                ],
+                dtype=np.float32,
+            ),
+            3,
+            "ABC",
+            ["ABC", "CCC"],
+        ),
+        (
+            np.array(
+                [[[0.1, 0.4, 0.5], [0.6, 0.2, 0.2], [0.1, 0.5, 0.4]]],
+                dtype=np.float32,
+            ),
+            3,
+            "ABC",
+            ["CAB"],
+        ),
+    ],
+)
+def test_postprocess_output(
+    model_output: npt.NDArray,
+    max_plate_slots: int,
+    model_alphabet: str,
+    expected_plates: list[str],
+) -> None:
+    actual_plate = postprocess_output(model_output, max_plate_slots, model_alphabet)
+    assert actual_plate == expected_plates