diff --git a/machine-learning/immich_ml/models/constants.py b/machine-learning/immich_ml/models/constants.py index 0815410495..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 @@ -86,8 +89,21 @@ _PADDLE_MODELS = { "KOREAN__PP-OCRv5_mobile", "LATIN__PP-OCRv5_mobile", "TH__PP-OCRv5_mobile", + "PP-OCRv6_tiny", + "PP-OCRv6_small", + "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 48dd87854b..6419f44abb 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/machine-learning/immich_ml/models/ocr/detection.py @@ -7,11 +7,11 @@ 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.models.constants import PADDLE_MODEL_SPECS from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType from immich_ml.sessions.ort import OrtSession @@ -41,13 +41,14 @@ class TextDetector(InferenceModel): ) def _download(self) -> None: + ocr_version, model_type = PADDLE_MODEL_SPECS[self.model_name.split("__")[-1]] 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 ea7ff565d1..4e6d936e4f 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/machine-learning/immich_ml/models/ocr/recognition.py @@ -7,12 +7,12 @@ 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 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 @@ -37,13 +37,14 @@ class TextRecognizer(InferenceModel): super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX) def _download(self) -> None: + ocr_version, model_type = PADDLE_MODEL_SPECS[self.model_name.split("__")[-1]] 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/pyproject.toml b/machine-learning/pyproject.toml index e95ef7afef..c4667c50b7 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/uv.lock b/machine-learning/uv.lock index fa92147445..867df284b5 100644 --- a/machine-learning/uv.lock +++ b/machine-learning/uv.lock @@ -875,7 +875,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" }, @@ -2125,7 +2125,7 @@ wheels = [ [[package]] name = "rapidocr" -version = "3.8.1" +version = "3.9.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorlog" }, @@ -2141,7 +2141,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 7f3c085f84..50b8a955c6 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' },