This commit is contained in:
Tada 2026-08-14 15:28:44 -04:00 committed by GitHub
commit bafa42b8f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 33 additions and 12 deletions

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

View file

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

View file

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

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]

View file

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

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