Fix bug in plate accuracy metric and add test

Author: ankandrew

fast_plate_ocr/custom.py

@@ -44,7 +44,7 @@ def plate_acc(y_true, y_pred):
         """
         y_true = ops.reshape(y_true, newshape=(-1, max_plate_slots, vocabulary_size))
         y_pred = ops.reshape(y_pred, newshape=(-1, max_plate_slots, vocabulary_size))
-        et = ops.equal(ops.argmax(y_true), ops.argmax(y_pred))
+        et = ops.equal(ops.argmax(y_true, axis=-1), ops.argmax(y_pred, axis=-1))
         return ops.mean(ops.cast(ops.all(et, axis=-1, keepdims=False), dtype="float32"))
 
     return plate_acc

test/fast_lp_ocr/test_custom.py

@@ -13,7 +13,7 @@
 import pytest
 import torch
 
-from fast_plate_ocr.custom import cat_acc_metric
+from fast_plate_ocr.custom import cat_acc_metric, plate_acc_metric
 
 
 @pytest.mark.parametrize(
@@ -22,6 +22,63 @@
         (torch.tensor([[[1, 0]] * 6]), torch.tensor([[[0.9, 0.1]] * 6]), 1.0),
     ],
 )
-def test_cat_acc(y_true: torch.Tensor, y_pred: torch.Tensor, expected_accuracy: float):
+def test_cat_acc(y_true: torch.Tensor, y_pred: torch.Tensor, expected_accuracy: float) -> None:
     actual_accuracy = cat_acc_metric(2, 1)(y_true, y_pred)
     assert actual_accuracy == expected_accuracy
+
+
+@pytest.mark.parametrize(
+    "y_true, y_pred, expected_accuracy",
+    [
+        (
+            torch.tensor(
+                [
+                    [
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    ],
+                    [
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    ],
+                ]
+            ),
+            torch.tensor(
+                [
+                    [
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    ],
+                    [
+                        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    ],
+                ]
+            ),
+            # First batch slice plate was recognized completely correct but second one wasn't
+            # So 50% of plates were recognized correctly
+            0.5,
+        ),
+    ],
+)
+def test_plate_accuracy(
+    y_true: torch.Tensor, y_pred: torch.Tensor, expected_accuracy: float
+) -> None:
+    actual_accuracy = plate_acc_metric(y_true.shape[1], y_true.shape[2])(y_true, y_pred).item()
+    assert actual_accuracy == expected_accuracy