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

refactor(ml): replace OCR version and type resolver with PADDLE_MODEL_SPECS dict

Browse source
This commit is contained in:
ducvuongpham 2026-07-09 08:06:41 +09:00
parent 3f300ce5c5
commit 73cd5bf699
No known key found for this signature in database
5 changed files with 21 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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

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

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

20
machine-learning/test_main.py
View file

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