diff --git a/machine-learning/immich_ml/models/constants.py b/machine-learning/immich_ml/models/constants.py index 0815410495..597e79222d 100644 --- a/machine-learning/immich_ml/models/constants.py +++ b/machine-learning/immich_ml/models/constants.py @@ -86,6 +86,9 @@ _PADDLE_MODELS = { "KOREAN__PP-OCRv5_mobile", "LATIN__PP-OCRv5_mobile", "TH__PP-OCRv5_mobile", + "PP-OCRv6_tiny", + "PP-OCRv6_small", + "PP-OCRv6_medium", } SUPPORTED_PROVIDERS = [ diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index 0a2cb8ad91..fdce421fdb 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/machine-learning/immich_ml/models/ocr/detection.py @@ -7,15 +7,14 @@ from PIL import Image from rapidocr.ch_ppocr_det.utils import DBPostProcess from rapidocr.inference_engine.base import FileInfo, InferSession from rapidocr.utils.download_file import DownloadFile, DownloadFileInput -from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType -from rapidocr.utils.typings import ModelType as RapidModelType +from rapidocr.utils.typings import EngineType, LangDet, TaskType from immich_ml.config import log from immich_ml.models.base import InferenceModel from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType from immich_ml.sessions.ort import OrtSession -from .schemas import TextDetectionOutput +from .schemas import TextDetectionOutput, resolve_ocr_version_and_type class TextDetector(InferenceModel): @@ -41,13 +40,14 @@ class TextDetector(InferenceModel): ) def _download(self) -> None: + ocr_version, model_type = resolve_ocr_version_and_type(self.model_name) model_info = InferSession.get_model_url( FileInfo( engine_type=EngineType.ONNXRUNTIME, - ocr_version=OCRVersion.PPOCRV5, + ocr_version=ocr_version, task_type=TaskType.DET, lang_type=LangDet.CH, - model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, + model_type=model_type, ) ) download_params = DownloadFileInput( diff --git a/machine-learning/immich_ml/models/ocr/recognition.py b/machine-learning/immich_ml/models/ocr/recognition.py index 94f40c9285..969f467537 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/machine-learning/immich_ml/models/ocr/recognition.py @@ -7,8 +7,7 @@ from rapidocr.ch_ppocr_rec import TextRecInput from rapidocr.ch_ppocr_rec import TextRecognizer as RapidTextRecognizer from rapidocr.inference_engine.base import FileInfo, InferSession from rapidocr.utils.download_file import DownloadFile, DownloadFileInput -from rapidocr.utils.typings import EngineType, LangRec, OCRVersion, TaskType -from rapidocr.utils.typings import ModelType as RapidModelType +from rapidocr.utils.typings import EngineType, LangRec, TaskType from rapidocr.utils.vis_res import VisRes from immich_ml.config import log, settings @@ -17,7 +16,7 @@ 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 +from .schemas import OcrOptions, TextDetectionOutput, TextRecognitionOutput, resolve_ocr_version_and_type class TextRecognizer(InferenceModel): @@ -37,13 +36,14 @@ 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) model_info = InferSession.get_model_url( FileInfo( engine_type=EngineType.ONNXRUNTIME, - ocr_version=OCRVersion.PPOCRV5, + ocr_version=ocr_version, task_type=TaskType.REC, lang_type=self.language, - model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, + model_type=model_type, ) ) download_params = DownloadFileInput( diff --git a/machine-learning/immich_ml/models/ocr/schemas.py b/machine-learning/immich_ml/models/ocr/schemas.py index 78e8619a0b..a0237f1432 100644 --- a/machine-learning/immich_ml/models/ocr/schemas.py +++ b/machine-learning/immich_ml/models/ocr/schemas.py @@ -2,10 +2,21 @@ from typing import Any, Iterable import numpy as np import numpy.typing as npt -from rapidocr.utils.typings import EngineType, LangRec +from rapidocr.utils.typings import EngineType, LangRec, OCRVersion +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] diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml index fe22b91a14..ed7a007c65 100644 --- a/machine-learning/pyproject.toml +++ b/machine-learning/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "rich>=13.4.2", "tokenizers>=0.15.0,<1.0", "uvicorn[standard]>=0.22.0,<1.0", - "rapidocr>=3.1.0", + "rapidocr>=3.9.1", ] [dependency-groups] diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index be574c6397..5d0737316b 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -28,7 +28,9 @@ 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 +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.schemas import ModelFormat, ModelPrecision, ModelTask, ModelType from immich_ml.sessions.ann import AnnSession from immich_ml.sessions.ort import OrtSession @@ -1019,6 +1021,22 @@ 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" diff --git a/machine-learning/uv.lock b/machine-learning/uv.lock index 8e081c12b4..24a78af64f 100644 --- a/machine-learning/uv.lock +++ b/machine-learning/uv.lock @@ -1076,7 +1076,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.0.0,<3" }, { name = "pydantic-settings", specifier = ">=2.5.2,<3" }, { name = "python-multipart", specifier = ">=0.0.6,<1.0" }, - { name = "rapidocr", specifier = ">=3.1.0" }, + { name = "rapidocr", specifier = ">=3.9.1" }, { name = "rich", specifier = ">=13.4.2" }, { name = "rknn-toolkit-lite2", marker = "extra == 'rknn'", specifier = ">=2.3.0,<3" }, { name = "tokenizers", specifier = ">=0.15.0,<1.0" }, @@ -2528,7 +2528,7 @@ wheels = [ [[package]] name = "rapidocr" -version = "3.8.1" +version = "3.9.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorlog" }, @@ -2544,7 +2544,7 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/4a/fa521d947f0fc7bb304bf11bec4cb66266bd81494588b4cb48dc01001719/rapidocr-3.8.1-py3-none-any.whl", hash = "sha256:650044b1fbce9e6bae5cae462dcf8be754cde11e2f23fc51f65dcc08deae2c46", size = 15080319, upload-time = "2026-04-11T07:13:22.56Z" }, + { url = "https://files.pythonhosted.org/packages/a0/23/e8d7251c53137b5b66d89e85cb0bc3e6bcfaec9527f29b477ec04389c8b2/rapidocr-3.9.1-py3-none-any.whl", hash = "sha256:600885e4e94e0b427abad394fccb0ec1d3c9118a215ca435bf7680aeae0e292b", size = 27274755, upload-time = "2026-07-02T13:37:00.827Z" }, ] [[package]] diff --git a/web/src/routes/admin/system-settings/MachineLearningSettings.svelte b/web/src/routes/admin/system-settings/MachineLearningSettings.svelte index 92b661fe8e..127871d734 100644 --- a/web/src/routes/admin/system-settings/MachineLearningSettings.svelte +++ b/web/src/routes/admin/system-settings/MachineLearningSettings.svelte @@ -275,6 +275,9 @@ name="ocr-model" bind:value={configToEdit.machineLearning.ocr.modelName} options={[ + { text: 'PP-OCRv6_medium (multilingual)', value: 'PP-OCRv6_medium' }, + { text: 'PP-OCRv6_small (multilingual)', value: 'PP-OCRv6_small' }, + { text: 'PP-OCRv6_tiny (multilingual, excludes Japanese)', value: 'PP-OCRv6_tiny' }, { text: 'PP-OCRv5_server (Chinese, Japanese and English)', value: 'PP-OCRv5_server' }, { text: 'PP-OCRv5_mobile (Chinese, Japanese and English)', value: 'PP-OCRv5_mobile' }, { text: 'PP-OCRv5_mobile (English-only)', value: 'EN__PP-OCRv5_mobile' },