diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index de4de264f8..5c93210f2f 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/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) diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index aa0aa0af43..1a4fbc1a0a 100644 --- a/machine-learning/test_main.py +++ b/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: