helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2026-08-15 13:03:57 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(ml): add PP-OCRv6 OCR models

Browse source
This commit is contained in:
ducvuongpham 2026-07-08 22:23:50 +09:00
parent 439ace9dba
commit 3f300ce5c5
No known key found for this signature in database
8 changed files with 51 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

3
machine-learning/immich_ml/models/constants.py
View file

@ -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 = [

10
machine-learning/immich_ml/models/ocr/detection.py
View file

@ -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(

10
machine-learning/immich_ml/models/ocr/recognition.py
View file

@ -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(

13
machine-learning/immich_ml/models/ocr/schemas.py
View file

@ -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]

2
machine-learning/pyproject.toml
View file

@ -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]

20
machine-learning/test_main.py
View file

@ -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"

6
machine-learning/uv.lock generated
View file

@ -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]]

3
web/src/routes/admin/system-settings/MachineLearningSettings.svelte
View file

@ -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' },