diff --git a/machine-learning/immich_ml/models/constants.py b/machine-learning/immich_ml/models/constants.py index 597e79222d..0dbc9a2479 100644 --- a/machine-learning/immich_ml/models/constants.py +++ b/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", diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index fdce421fdb..f87ae4b0c4 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/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, diff --git a/machine-learning/immich_ml/models/ocr/recognition.py b/machine-learning/immich_ml/models/ocr/recognition.py index 969f467537..f5669c0a75 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/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, diff --git a/machine-learning/immich_ml/models/ocr/schemas.py b/machine-learning/immich_ml/models/ocr/schemas.py index a0237f1432..78e8619a0b 100644 --- a/machine-learning/immich_ml/models/ocr/schemas.py +++ b/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 ModelType as RapidModelType +from rapidocr.utils.typings import EngineType, LangRec 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] diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index 5d0737316b..be574c6397 100644 --- a/machine-learning/test_main.py +++ b/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 rapidocr.utils.typings import ModelType as RapidModelType -from rapidocr.utils.typings import OCRVersion +from immich_ml.models.ocr.schemas import OcrOptions 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"