update tests

machine-learning/conftest.py

@@ -1,11 +1,10 @@
 import json
-from typing import Any, Iterator
+from types import SimpleNamespace
+from typing import Any, Callable, Iterator
 from unittest import mock
 
-import numpy as np
 import pytest
 from fastapi.testclient import TestClient
-from numpy.typing import NDArray
 from PIL import Image
 
 from immich_ml.config import log
@@ -17,11 +16,6 @@ def pil_image() -> Image.Image:
     return Image.new("RGB", (600, 800))
 
 
-@pytest.fixture
-def cv_image(pil_image: Image.Image) -> NDArray[np.float32]:
-    return np.asarray(pil_image)[:, :, ::-1]  # PIL uses RGB while cv2 uses BGR
-
-
 @pytest.fixture
 def mock_get_model() -> Iterator[mock.Mock]:
     with mock.patch("immich_ml.models.cache.from_model_type", autospec=True) as mocked:
@@ -129,6 +123,18 @@ def ort_session() -> Iterator[mock.Mock]:
         yield mocked
 
 
+@pytest.fixture(scope="function")
+def stub_session() -> Callable[..., mock.Mock]:
+    def _make(shape: tuple[Any, ...], outputs: Any = None, name: str = "input.1") -> mock.Mock:
+        session = mock.Mock()
+        session.get_inputs.return_value = [SimpleNamespace(name=name, shape=shape)]
+        if outputs is not None:
+            session.run.return_value = outputs
+        return session
+
+    return _make
+
+
 @pytest.fixture(scope="function")
 def ann_session() -> Iterator[mock.Mock]:
     with mock.patch("immich_ml.sessions.ann.Ann") as mocked:

machine-learning/immich_ml/models/facial_recognition/_ops.py

@@ -52,7 +52,7 @@ def decode_scrfd(
 
 def nms(boxes: NDArray[np.float32], scores: NDArray[np.float32], threshold: float = 0.4) -> NDArray[np.intp]:
     wh = np.column_stack([boxes[:, 0], boxes[:, 1], boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]])
-    keep = cv2.dnn.NMSBoxes(wh.tolist(), scores.tolist(), 0.0, threshold) # NMSBoxes treats the inputs as Sequences
+    keep = cv2.dnn.NMSBoxes(wh.tolist(), scores.tolist(), 0.0, threshold)  # NMSBoxes treats the inputs as Sequences
     return np.asarray(keep, dtype=np.intp).reshape(-1)
 
 

machine-learning/test_main.py

@@ -7,7 +7,6 @@ from types import SimpleNamespace
 from typing import Any, Callable
 from unittest import mock
 
-import cv2
 import numpy as np
 import onnxruntime as ort
 import orjson
@@ -781,89 +780,189 @@ class TestCLIP:
         assert np.allclose(tokens["attention_mask"], np.array([mock_attention_mask], dtype=np.int32), atol=0)
 
 
+def make_scrfd_heads(detections: list[tuple[int, int, float]]) -> list[np.ndarray]:
+    """Build the 9 head tensors a SCRFD keypoint model emits at 640x640.
+
+    `detections` is a list of (cell_x, cell_y, score) placed on the stride-8 level.
+    Distances and keypoint offsets are fixed so the decoded geometry is known by
+    hand rather than derived from the code under test:
+        box  = [cx - 1*8, cy - 2*8, cx + 3*8, cy + 4*8]
+        kps  = [cx, cy] + [0, 1, 2, ... 9] * 8, reshaped to 5 points
+    """
+    heads: list[np.ndarray] = []
+    counts = [(640 // stride) ** 2 * 2 for stride in (8, 16, 32)]
+    for channels in (1, 4, 10):
+        for n in counts:
+            heads.append(np.zeros((n, channels), dtype=np.float32))
+    for cell_x, cell_y, score in detections:
+        i = 2 * (cell_y * 80 + cell_x)  # anchor-major, 2 anchors per cell
+        heads[0][i] = score
+        heads[3][i] = [1, 2, 3, 4]
+        heads[6][i] = np.arange(10)
+    return heads
+
+
+def expected_box(cell_x: int, cell_y: int) -> list[float]:
+    cx, cy = cell_x * 8, cell_y * 8
+    return [cx - 8, cy - 16, cx + 24, cy + 32]
+
+
+def expected_landmarks(cell_x: int, cell_y: int) -> np.ndarray:
+    cx, cy = cell_x * 8, cell_y * 8
+    return (np.tile([cx, cy], 5) + np.arange(10) * 8).reshape(5, 2).astype(np.float32)
+
+
+@pytest.fixture
+def batch_axis(mocker: MockerFixture) -> SimpleNamespace:
+    """Patches for the `_add_batch_axis` path: the onnx module it rewrites the model
+    through, and the download it would otherwise trigger on load."""
+    mocker.patch("immich_ml.models.base.InferenceModel.download")
+    return SimpleNamespace(
+        onnx=mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True),
+        update_dims=mocker.patch(
+            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
+        ),
+    )
+
+
 class TestFaceRecognition:
-    def test_set_min_score(self, snapshot_download: mock.Mock, ort_session: mock.Mock, path: mock.Mock) -> None:
-        path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
-
-        face_detector = FaceDetector("buffalo_s", min_score=0.5, cache_dir="test_cache")
-        face_detector.load()
-
-        assert face_detector.min_score == 0.5
-        assert face_detector.model.det_thresh == 0.5
-
-    def test_detection(self, cv_image: cv2.Mat, mocker: MockerFixture) -> None:
+    def test_detection(self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture) -> None:
         mocker.patch.object(FaceDetector, "load")
-        face_detector = FaceDetector("buffalo_s", min_score=0.0, cache_dir="test_cache")
+        face_detector = FaceDetector("buffalo_s", cache_dir="test_cache")
 
-        det_model = mock.Mock()
-        num_faces = 2
-        bbox = np.random.rand(num_faces, 4).astype(np.float32)
-        scores = np.array([[0.67]] * num_faces).astype(np.float32)
-        kpss = np.random.rand(num_faces, 5, 2).astype(np.float32)
-        det_model.detect.return_value = (np.concatenate([bbox, scores], axis=-1), kpss)
-        face_detector.model = det_model
+        session = stub_session((1, 3, 640, 640), outputs=make_scrfd_heads([(10, 10, 0.9), (50, 50, 0.8)]))
+        face_detector.session = session
 
-        faces = face_detector.predict(cv_image)
+        faces = face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.7)
 
         assert isinstance(faces, dict)
-        assert isinstance(faces.get("boxes", None), np.ndarray)
-        assert isinstance(faces.get("landmarks", None), np.ndarray)
-        assert isinstance(faces.get("scores", None), np.ndarray)
-        assert np.equal(faces["boxes"], bbox.round()).all()
-        assert np.equal(faces["landmarks"], kpss).all()
-        assert np.equal(faces["scores"], scores).all()
-        det_model.detect.assert_called_once()
+        assert set(faces) == {"boxes", "scores", "landmarks"}
+        # NMS returns highest score first
+        assert faces["boxes"].tolist() == [expected_box(10, 10), expected_box(50, 50)]
+        assert np.allclose(faces["scores"], [0.9, 0.8])
+        assert faces["landmarks"].shape == (2, 5, 2)
+        assert np.allclose(faces["landmarks"][0], expected_landmarks(10, 10))
+        assert np.allclose(faces["landmarks"][1], expected_landmarks(50, 50))
 
-    def test_recognition(self, cv_image: cv2.Mat, mocker: MockerFixture) -> None:
+    def test_detection_applies_min_score_per_request(
+        self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture
+    ) -> None:
+        mocker.patch.object(FaceDetector, "load")
+        face_detector = FaceDetector("buffalo_s", cache_dir="test_cache")
+
+        session = stub_session((1, 3, 640, 640), outputs=make_scrfd_heads([(10, 10, 0.9), (50, 50, 0.5)]))
+        face_detector.session = session
+
+        # the threshold is a request parameter, so the same loaded model must honour both
+        assert face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.7)["boxes"].shape[0] == 1
+        assert face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.4)["boxes"].shape[0] == 2
+
+    def test_detection_scales_boxes_back_to_the_original_image(
+        self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture
+    ) -> None:
+        mocker.patch.object(FaceDetector, "load")
+        face_detector = FaceDetector("buffalo_s", cache_dir="test_cache")
+
+        session = stub_session((1, 3, 640, 640), outputs=make_scrfd_heads([(10, 10, 0.9)]))
+        face_detector.session = session
+
+        # a 320x320 image is letterboxed up to 640, so coordinates come back halved
+        faces = face_detector.predict(Image.new("RGB", (320, 320)), minScore=0.7)
+
+        assert faces["boxes"].tolist() == [[v / 2 for v in expected_box(10, 10)]]
+        assert np.allclose(faces["landmarks"][0], expected_landmarks(10, 10) / 2)
+
+    def test_recognition(self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture) -> None:
         mocker.patch.object(FaceRecognizer, "load")
         mocker.patch(
             "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
             return_value=["CPUExecutionProvider"],
         )
-        face_recognizer = FaceRecognizer("buffalo_s", min_score=0.0, cache_dir="test_cache")
+        face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache")
 
+        # a uniform grey image whose crops land wholly inside it, so every sampled
+        # pixel is 128 and the normalised value the session receives is exact
+        image = Image.new("RGB", (600, 800), (128, 128, 128))
         num_faces = 2
+        arcface_dst = np.array(
+            [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]],
+            dtype=np.float32,
+        )
+        kpss = np.stack([arcface_dst * 2 + [200, 300], arcface_dst * 2 + [220, 320]]).astype(np.float32)
         bbox = np.random.rand(num_faces, 4).astype(np.float32)
         scores = np.array([0.67] * num_faces).astype(np.float32)
-        kpss = np.random.rand(num_faces, 5, 2).astype(np.float32)
-        faces = {"boxes": bbox, "landmarks": kpss, "scores": scores}
+        embeddings = np.random.rand(num_faces, 512).astype(np.float32)
 
-        rec_model = mock.Mock()
-        embedding = np.random.rand(num_faces, 512).astype(np.float32)
-        rec_model.get_feat.return_value = embedding
-        face_recognizer.model = rec_model
+        session = stub_session(("batch", 3, 112, 112), outputs=[embeddings])
+        face_recognizer.session = session
 
-        faces = face_recognizer.predict(cv_image, faces)
+        faces = face_recognizer.predict(image, {"boxes": bbox, "landmarks": kpss, "scores": scores})
 
         assert isinstance(faces, list)
         assert len(faces) == num_faces
         for face in faces:
-            assert isinstance(face.get("boundingBox"), dict)
             assert set(face["boundingBox"]) == {"x1", "y1", "x2", "y2"}
             assert all(isinstance(val, np.float32) for val in face["boundingBox"].values())
-            embedding_str = face.get("embedding")
-            assert isinstance(embedding_str, str)
-            embedding = orjson.loads(embedding_str)
+            embedding = orjson.loads(face["embedding"])
             assert isinstance(embedding, list)
             assert len(embedding) == 512
             assert isinstance(face.get("score", None), np.float32)
 
-        rec_model.get_feat.assert_called_once()
-        call_args = rec_model.get_feat.call_args_list[0].args
-        assert len(call_args) == 1
-        assert isinstance(call_args[0], list)
-        assert isinstance(call_args[0][0], np.ndarray)
-        assert call_args[0][0].shape == (112, 112, 3)
+        session.run.assert_called_once()
+        crops = session.run.call_args.args[1]["input.1"]
+        assert crops.shape == (num_faces, 3, 112, 112)
+        assert crops.dtype == np.float32
+        # mean/std 127.5, not raw 0-255. atol is loose enough for the float32 cancellation
+        # in normalize's scale-then-subtract, but still rejects a wrong mean or std
+        assert np.allclose(crops, (128 - 127.5) / 127.5, atol=1e-6)
+
+    def test_recognition_returns_early_without_faces(self, pil_image: Image.Image, mocker: MockerFixture) -> None:
+        mocker.patch.object(FaceRecognizer, "load")
+        mocker.patch(
+            "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
+            return_value=["CPUExecutionProvider"],
+        )
+        face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache")
+        session = mock.Mock()
+        face_recognizer.session = session
+
+        empty = {
+            "boxes": np.empty((0, 4), dtype=np.float32),
+            "landmarks": np.empty((0, 5, 2), dtype=np.float32),
+            "scores": np.empty(0, dtype=np.float32),
+        }
+
+        assert face_recognizer.predict(pil_image, empty) == []
+        session.run.assert_not_called()
+
+    def test_recognition_batches_when_batch_size_is_set(
+        self, pil_image: Image.Image, stub_session: Callable[..., mock.Mock], mocker: MockerFixture
+    ) -> None:
+        mocker.patch.object(FaceRecognizer, "load")
+        mocker.patch(
+            "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
+            return_value=["CPUExecutionProvider"],
+        )
+        face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache")
+        face_recognizer.batch_size = 2
+
+        num_faces = 5
+        session = stub_session((1, 3, 112, 112))
+        session.run.side_effect = lambda _, feed: [np.zeros((feed["input.1"].shape[0], 512), dtype=np.float32)]
+        face_recognizer.session = session
+
+        faces = {
+            "boxes": np.random.rand(num_faces, 4).astype(np.float32),
+            "landmarks": (np.random.rand(num_faces, 5, 2) * 100).astype(np.float32),
+            "scores": np.array([0.67] * num_faces, dtype=np.float32),
+        }
+        assert len(face_recognizer.predict(pil_image, faces)) == num_faces
+        assert session.run.call_count == 3  # 2 + 2 + 1
+        assert [c.args[1]["input.1"].shape[0] for c in session.run.call_args_list] == [2, 2, 1]
 
     def test_recognition_adds_batch_axis_for_ort(
-        self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
+        self, batch_axis: SimpleNamespace, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
     ) -> None:
-        onnx = mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True)
-        update_dims = mocker.patch(
-            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
-        )
-        mocker.patch("immich_ml.models.base.InferenceModel.download")
-        mocker.patch("immich_ml.models.facial_recognition.recognition.ArcFaceONNX")
         mocker.patch(
             "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
             return_value=["CPUExecutionProvider"],
@@ -884,24 +983,20 @@ class TestFaceRecognition:
         output_dims.type.tensor_type.shape.dim = [SimpleNamespace(dim_value=size) for size in [1, 800]]
         proto.graph.output = [output_dims]
 
-        onnx.load.return_value = proto
+        batch_axis.onnx.load.return_value = proto
 
         face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path)
         face_recognizer.load()
 
         assert face_recognizer.batch_size is None
-        update_dims.assert_called_once_with(proto, {"input.1": ["batch", 3, 224, 224]}, {"output.1": ["batch", 800]})
-        onnx.save.assert_called_once_with(update_dims.return_value, face_recognizer.model_path)
+        batch_axis.update_dims.assert_called_once_with(
+            proto, {"input.1": ["batch", 3, 224, 224]}, {"output.1": ["batch", 800]}
+        )
+        batch_axis.onnx.save.assert_called_once_with(batch_axis.update_dims.return_value, face_recognizer.model_path)
 
     def test_recognition_does_not_add_batch_axis_if_exists(
-        self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
+        self, batch_axis: SimpleNamespace, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
     ) -> None:
-        onnx = mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True)
-        update_dims = mocker.patch(
-            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
-        )
-        mocker.patch("immich_ml.models.base.InferenceModel.download")
-        mocker.patch("immich_ml.models.facial_recognition.recognition.ArcFaceONNX")
         mocker.patch(
             "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
             return_value=["CPUExecutionProvider"],
@@ -916,88 +1011,53 @@ class TestFaceRecognition:
         face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path)
         face_recognizer.load()
 
+        # batching is available here, so the axis is only skipped because it already exists
         assert face_recognizer.batch_size is None
-        update_dims.assert_not_called()
-        onnx.load.assert_not_called()
-        onnx.save.assert_not_called()
+        batch_axis.update_dims.assert_not_called()
+        batch_axis.onnx.load.assert_not_called()
+        batch_axis.onnx.save.assert_not_called()
 
-    def test_recognition_does_not_add_batch_axis_for_armnn(
-        self, ann_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
+    @pytest.mark.parametrize(
+        ("session_fixture", "suffix", "model_kwargs", "available_providers"),
+        [
+            pytest.param("ann_session", ".armnn", {"model_format": ModelFormat.ARMNN}, None, id="armnn"),
+            pytest.param(
+                "ort_session", ".onnx", {}, ["OpenVINOExecutionProvider", "CPUExecutionProvider"], id="openvino"
+            ),
+            pytest.param(
+                "ort_session", ".onnx", {}, ["MIGraphXExecutionProvider", "CPUExecutionProvider"], id="migraphx"
+            ),
+        ],
+    )
+    def test_recognition_does_not_add_batch_axis_when_batching_is_unsupported(
+        self,
+        request: pytest.FixtureRequest,
+        batch_axis: SimpleNamespace,
+        path: mock.Mock,
+        ort_pybind: mock.Mock,
+        mocker: MockerFixture,
+        session_fixture: str,
+        suffix: str,
+        model_kwargs: dict[str, Any],
+        available_providers: list[str] | None,
     ) -> None:
-        onnx = mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True)
-        update_dims = mocker.patch(
-            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
-        )
-        mocker.patch("immich_ml.models.base.InferenceModel.download")
-        mocker.patch("immich_ml.models.facial_recognition.recognition.ArcFaceONNX")
-        path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".armnn"
-
-        inputs = [SimpleNamespace(name="input.1", shape=("batch", 3, 224, 224))]
-        outputs = [SimpleNamespace(name="output.1", shape=("batch", 800))]
-        ann_session.return_value.get_inputs.return_value = inputs
-        ann_session.return_value.get_outputs.return_value = outputs
-
-        face_recognizer = FaceRecognizer("buffalo_s", model_format=ModelFormat.ARMNN, cache_dir=path)
-        face_recognizer.load()
-
-        assert face_recognizer.batch_size == 1
-        update_dims.assert_not_called()
-        onnx.load.assert_not_called()
-        onnx.save.assert_not_called()
-
-    def test_recognition_does_not_add_batch_axis_for_openvino(
-        self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
-    ) -> None:
-        onnx = mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True)
-        update_dims = mocker.patch(
-            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
-        )
-        mocker.patch("immich_ml.models.base.InferenceModel.download")
-        mocker.patch("immich_ml.models.facial_recognition.recognition.ArcFaceONNX")
-        path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
-
-        inputs = [SimpleNamespace(name="input.1", shape=("batch", 3, 224, 224))]
-        outputs = [SimpleNamespace(name="output.1", shape=("batch", 800))]
-        ort_session.return_value.get_inputs.return_value = inputs
-        ort_session.return_value.get_outputs.return_value = outputs
-
-        face_recognizer = FaceRecognizer(
-            "buffalo_s", model_format=ModelFormat.ARMNN, cache_dir=path, providers=["OpenVINOExecutionProvider"]
-        )
-        face_recognizer.load()
-
-        assert face_recognizer.batch_size == 1
-        update_dims.assert_not_called()
-        onnx.load.assert_not_called()
-        onnx.save.assert_not_called()
-
-    def test_recognition_does_not_add_batch_axis_for_migraphx(
-        self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
-    ) -> None:
-        onnx = mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True)
-        update_dims = mocker.patch(
-            "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True
-        )
-        mocker.patch("immich_ml.models.base.InferenceModel.download")
-        mocker.patch("immich_ml.models.facial_recognition.recognition.ArcFaceONNX")
+        session = request.getfixturevalue(session_fixture)
+        ort_pybind.get_available_openvino_device_ids.return_value = ["CPU"]
         mocker.patch(
             "immich_ml.models.facial_recognition.recognition.ort.get_available_providers",
-            return_value=["MIGraphXExecutionProvider", "CPUExecutionProvider"],
+            return_value=available_providers,
         )
-        path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
+        path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = suffix
+        session.return_value.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 224, 224))]
+        session.return_value.get_outputs.return_value = [SimpleNamespace(name="output.1", shape=(1, 800))]
 
-        inputs = [SimpleNamespace(name="input.1", shape=(1, 3, 224, 224))]
-        outputs = [SimpleNamespace(name="output.1", shape=(1, 800))]
-        ort_session.return_value.get_inputs.return_value = inputs
-        ort_session.return_value.get_outputs.return_value = outputs
-
-        face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path)
+        face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path, **model_kwargs)
         face_recognizer.load()
 
         assert face_recognizer.batch_size == 1
-        update_dims.assert_not_called()
-        onnx.load.assert_not_called()
-        onnx.save.assert_not_called()
+        batch_axis.update_dims.assert_not_called()
+        batch_axis.onnx.load.assert_not_called()
+        batch_axis.onnx.save.assert_not_called()
 
     def test_set_custom_max_batch_size(self, mocker: MockerFixture) -> None:
         mocker.patch.object(settings, "max_batch_size", MaxBatchSize(facial_recognition=2))
@@ -1045,10 +1105,10 @@ class TestOcr:
 
         rapid_recognizer.assert_called_once_with(
             OcrOptions(
-              session=ort_session.return_value,
-              rec_batch_num=6,
-              rec_img_shape=(3, 48, 320),
-              model_root_dir=text_recognizer.cache_dir,
+                session=ort_session.return_value,
+                rec_batch_num=6,
+                rec_img_shape=(3, 48, 320),
+                model_root_dir=text_recognizer.cache_dir,
             )
         )
 
@@ -1063,10 +1123,10 @@ class TestOcr:
 
         rapid_recognizer.assert_called_once_with(
             OcrOptions(
-              session=ort_session.return_value,
-              rec_batch_num=4,
-              rec_img_shape=(3, 48, 320),
-              model_root_dir=text_recognizer.cache_dir,
+                session=ort_session.return_value,
+                rec_batch_num=4,
+                rec_img_shape=(3, 48, 320),
+                model_root_dir=text_recognizer.cache_dir,
             )
         )
 
@@ -1083,10 +1143,10 @@ class TestOcr:
 
         rapid_recognizer.assert_called_once_with(
             OcrOptions(
-              session=ort_session.return_value,
-              rec_batch_num=6,
-              rec_img_shape=(3, 48, 320),
-              model_root_dir=text_recognizer.cache_dir,
+                session=ort_session.return_value,
+                rec_batch_num=6,
+                rec_img_shape=(3, 48, 320),
+                model_root_dir=text_recognizer.cache_dir,
             )
         )