refactor(ml): replace OCR version and type resolver with PADDLE_MODEL_SPECS dict

machine-learning/immich_ml/models/constants.py

@@ -1,3 +1,6 @@
+from rapidocr.utils.typings import ModelType as RapidModelType
+from rapidocr.utils.typings import OCRVersion
+
 from immich_ml.config import clean_name
 from immich_ml.schemas import ModelSource
 
@@ -91,6 +94,16 @@ _PADDLE_MODELS = {
     "PP-OCRv6_medium",
 }
 
+
+PADDLE_MODEL_SPECS: dict[str, tuple[OCRVersion, RapidModelType]] = {
+    "PP-OCRv5_server": (OCRVersion.PPOCRV5, RapidModelType.SERVER),
+    "PP-OCRv5_mobile": (OCRVersion.PPOCRV5, RapidModelType.MOBILE),
+    "PP-OCRv6_tiny": (OCRVersion.PPOCRV6, RapidModelType.TINY),
+    "PP-OCRv6_small": (OCRVersion.PPOCRV6, RapidModelType.SMALL),
+    "PP-OCRv6_medium": (OCRVersion.PPOCRV6, RapidModelType.MEDIUM),
+}
+
+
 SUPPORTED_PROVIDERS = [
     "CUDAExecutionProvider",
     "MIGraphXExecutionProvider",

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

@@ -11,10 +11,11 @@ from rapidocr.utils.typings import EngineType, LangDet, TaskType
 
 from immich_ml.config import log
 from immich_ml.models.base import InferenceModel
+from immich_ml.models.constants import PADDLE_MODEL_SPECS
 from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
 from immich_ml.sessions.ort import OrtSession
 
-from .schemas import TextDetectionOutput, resolve_ocr_version_and_type
+from .schemas import TextDetectionOutput
 
 
 class TextDetector(InferenceModel):
@@ -40,7 +41,7 @@ class TextDetector(InferenceModel):
         )
 
     def _download(self) -> None:
-        ocr_version, model_type = resolve_ocr_version_and_type(self.model_name)
+        ocr_version, model_type = PADDLE_MODEL_SPECS[self.model_name.split("__")[-1]]
         model_info = InferSession.get_model_url(
             FileInfo(
                 engine_type=EngineType.ONNXRUNTIME,

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

@@ -12,11 +12,12 @@ from rapidocr.utils.vis_res import VisRes
 
 from immich_ml.config import log, settings
 from immich_ml.models.base import InferenceModel
+from immich_ml.models.constants import PADDLE_MODEL_SPECS
 from immich_ml.models.transforms import pil_to_cv2
 from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
 from immich_ml.sessions.ort import OrtSession
 
-from .schemas import OcrOptions, TextDetectionOutput, TextRecognitionOutput, resolve_ocr_version_and_type
+from .schemas import OcrOptions, TextDetectionOutput, TextRecognitionOutput
 
 
 class TextRecognizer(InferenceModel):
@@ -36,7 +37,7 @@ class TextRecognizer(InferenceModel):
         super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
 
     def _download(self) -> None:
-        ocr_version, model_type = resolve_ocr_version_and_type(self.model_name)
+        ocr_version, model_type = PADDLE_MODEL_SPECS[self.model_name.split("__")[-1]]
         model_info = InferSession.get_model_url(
             FileInfo(
                 engine_type=EngineType.ONNXRUNTIME,

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

@@ -2,21 +2,10 @@ from typing import Any, Iterable
 
 import numpy as np
 import numpy.typing as npt
-from rapidocr.utils.typings import EngineType, LangRec, OCRVersion
+from rapidocr.utils.typings import EngineType, LangRec
-from rapidocr.utils.typings import ModelType as RapidModelType
 from typing_extensions import TypedDict
 
 
-def resolve_ocr_version_and_type(model_name: str) -> tuple[OCRVersion, RapidModelType]:
-    if "PP-OCRv6" in model_name:
-        if "tiny" in model_name:
-            return OCRVersion.PPOCRV6, RapidModelType.TINY
-        if "medium" in model_name:
-            return OCRVersion.PPOCRV6, RapidModelType.MEDIUM
-        return OCRVersion.PPOCRV6, RapidModelType.SMALL
-    return OCRVersion.PPOCRV5, RapidModelType.MOBILE if "mobile" in model_name else RapidModelType.SERVER
-
-
 class TextDetectionOutput(TypedDict):
     boxes: npt.NDArray[np.float32]
     scores: npt.NDArray[np.float32]

machine-learning/test_main.py

@@ -28,9 +28,7 @@ from immich_ml.models.facial_recognition.detection import FaceDetector
 from immich_ml.models.facial_recognition.recognition import FaceRecognizer
 from immich_ml.models.ocr.detection import TextDetector
 from immich_ml.models.ocr.recognition import TextRecognizer
-from immich_ml.models.ocr.schemas import OcrOptions, resolve_ocr_version_and_type
+from immich_ml.models.ocr.schemas import OcrOptions
-from rapidocr.utils.typings import ModelType as RapidModelType
-from rapidocr.utils.typings import OCRVersion
 from immich_ml.schemas import ModelFormat, ModelPrecision, ModelTask, ModelType
 from immich_ml.sessions.ann import AnnSession
 from immich_ml.sessions.ort import OrtSession
@@ -1021,22 +1019,6 @@ class TestFaceRecognition:
 
 
 class TestOcr:
-    @pytest.mark.parametrize(
-        ("model_name", "expected"),
-        [
-            ("PP-OCRv6_tiny", (OCRVersion.PPOCRV6, RapidModelType.TINY)),
-            ("PP-OCRv6_small", (OCRVersion.PPOCRV6, RapidModelType.SMALL)),
-            ("PP-OCRv6_medium", (OCRVersion.PPOCRV6, RapidModelType.MEDIUM)),
-            ("PP-OCRv5_mobile", (OCRVersion.PPOCRV5, RapidModelType.MOBILE)),
-            ("PP-OCRv5_server", (OCRVersion.PPOCRV5, RapidModelType.SERVER)),
-            ("LATIN__PP-OCRv5_mobile", (OCRVersion.PPOCRV5, RapidModelType.MOBILE)),
-        ],
-    )
-    def test_resolve_ocr_version_and_type(
-        self, model_name: str, expected: tuple[OCRVersion, RapidModelType]
-    ) -> None:
-        assert resolve_ocr_version_and_type(model_name) == expected
-
     def test_set_det_min_score(self, path: mock.Mock) -> None:
         path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"