fix(ml): Round OCR max resolution up to a multiple of 32 (#32263)

fix(ml): round ocr max resolution up to a multiple of 32

machine-learning/immich_ml/models/ocr/detection.py

@@ -1,3 +1,4 @@
+import math
 from collections.abc import Hashable, Sequence
 from typing import Any, Self
 
@@ -23,9 +24,10 @@ class TextDetector(TextModel[TextDetectionOptions]):
         self, model_name: str, max_resolution: int = TextDetectionOptions.max_resolution, **model_kwargs: Any
     ) -> None:
         super().__init__(model_name, **model_kwargs)
+        size = self._canvas_size(max_resolution)
         # RKNPU ships a binary per short side, which the label picks
-        canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(max_resolution))
-        self.shape_policy = ShapePolicy(dims=canvases, label=f"res{max_resolution}")
+        canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(size))
+        self.shape_policy = ShapePolicy(dims=canvases, label=f"res{size}")
         self.scale = np.float32(1.0 / 127.5)  # (x/255 - 0.5) / 0.5
         self._empty: TextDetectionOutput = {
             "boxes": np.empty(0, dtype=np.float32),
@@ -35,7 +37,7 @@ class TextDetector(TextModel[TextDetectionOptions]):
 
     @classmethod
     def graph(cls, options: TextDetectionOptions) -> Hashable:
-        return options.max_resolution
+        return cls._canvas_size(options.max_resolution)
 
     @classmethod
     def create(cls, model_name: str, options: TextDetectionOptions) -> Self:
@@ -97,6 +99,11 @@ class TextDetector(TextModel[TextDetectionOptions]):
             default=(height, width),
         )
 
+    @staticmethod
+    def _canvas_size(max_resolution: int) -> int:
+        # the graph only takes sides divisible by 32; up, so the short side _transform rounds to always fits
+        return math.ceil(max_resolution / 32) * 32
+
     @staticmethod
     def _round32(value: float) -> int:
         return int(round(value / 32) * 32)

machine-learning/test_main.py

@@ -1742,6 +1742,22 @@ class TestOcr:
         # the legacy path swaps to BGR on the way in; the fused graph does its own
         assert fed[0, 0, 0].tolist() == [10, 20, 30]
 
+    @pytest.mark.parametrize(("max_resolution", "label"), [(736, "res736"), (1000, "res1024"), (1080, "res1088")])
+    def test_det_rounds_its_canvases_up_to_a_multiple_of_32(
+        self, path: mock.Mock, max_resolution: int, label: str
+    ) -> None:
+        text_detector = TextDetector("PP-OCRv5_mobile", max_resolution=max_resolution, cache_dir=path)
+
+        dims = [dim for shape in text_detector.shape_policy.dims for dim in (shape.height, shape.width)]
+        assert dims and all(dim is not None and dim % 32 == 0 for dim in dims)  # the graph cannot take any other size
+        assert text_detector.shape_policy.label == label  # up, so 1080 finds the res1088 binary RKNPU ships
+
+    def test_det_resolutions_in_one_32_band_share_a_graph(self) -> None:
+        graph = [TextDetector.graph(TextDetectionOptions(max_resolution=size)) for size in (1080, 1088, 1089)]
+
+        assert graph[0] == graph[1]
+        assert graph[1] != graph[2]  # the next band compiles canvases of its own
+
     def test_rec_runs_the_batch_at_a_compiled_width(
         self, path: mock.Mock, mocker: MockerFixture, stub_session: Callable[..., mock.Mock]
     ) -> None: