diff --git a/machine-learning/immich_ml/models/constants.py b/machine-learning/immich_ml/models/constants.py index 2920d8257..d386b0768 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/legacy.py b/machine-learning/immich_ml/models/ocr/legacy.py index db76feb1b..61a74ba66 100644 --- a/machine-learning/immich_ml/models/ocr/legacy.py +++ b/machine-learning/immich_ml/models/ocr/legacy.py @@ -23,13 +23,14 @@ class TextModel[O: Options](InferenceModel[O]): detects = self.model_type == ModelType.DETECTION language = self.model_name.split("__")[0] if "__" in self.model_name else "CH" + version, size = self.model_name.split("__")[-1].rsplit("_", 1) model_info = InferSession.get_model_url( FileInfo( engine_type=EngineType.ONNXRUNTIME, - ocr_version=OCRVersion.PPOCRV5, + ocr_version=OCRVersion(version), task_type=TaskType.DET if detects else TaskType.REC, lang_type=LangDet.CH if detects else LangRec[language], - model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, + model_type=RapidModelType(size), ) ) log.info( diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index 50319d530..5a0a869cc 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -1690,6 +1690,8 @@ class TestOcr: fetched = fetch.call_args.args[0] assert fetched.file_url.endswith("/onnx/PP-OCRv5/rec/en_PP-OCRv5_rec_mobile.onnx") assert fetched.save_path == tmp_path / "recognition/model.onnx" + TextDetector("PP-OCRv6_tiny", cache_dir=tmp_path).download() + assert fetch.call_args.args[0].file_url.endswith("/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx") snapshot_download.assert_not_called() mocker.patch.object(settings, "model_revision", "v2") diff --git a/web/src/routes/admin/system-settings/MachineLearningSettings.svelte b/web/src/routes/admin/system-settings/MachineLearningSettings.svelte index 7f3c085f8..eb4362914 100644 --- a/web/src/routes/admin/system-settings/MachineLearningSettings.svelte +++ b/web/src/routes/admin/system-settings/MachineLearningSettings.svelte @@ -275,6 +275,15 @@ name="ocr-model" bind:value={configToEdit.machineLearning.ocr.modelName} options={[ + { + text: 'PP-OCRv6_medium (Chinese, Japanese, English and other Latin script languages)', + value: 'PP-OCRv6_medium', + }, + { + text: 'PP-OCRv6_small (Chinese, Japanese, English and other Latin script languages)', + value: 'PP-OCRv6_small', + }, + { text: 'PP-OCRv6_tiny (Chinese, English and other Latin script languages)', 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' },