mirror of
https://github.com/immich-app/immich
synced 2026-08-29 13:15:45 +00:00
106 lines
4.5 KiB
Python
106 lines
4.5 KiB
Python
from functools import cached_property
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
from PIL import Image
|
|
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 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 .postprocess import DBPostProcess
|
|
from .schemas import TextDetectionOutput
|
|
|
|
|
|
class TextDetector(InferenceModel):
|
|
depends = []
|
|
identity = (ModelType.DETECTION, ModelTask.OCR)
|
|
|
|
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
|
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
|
self.scale = np.float32(1.0 / 127.5) # (x/255 - 0.5) / 0.5
|
|
self._empty: TextDetectionOutput = {
|
|
"boxes": np.empty(0, dtype=np.float32),
|
|
"scores": np.empty(0, dtype=np.float32),
|
|
}
|
|
self.postprocess = DBPostProcess(thresh=0.3, max_candidates=1000, unclip_ratio=1.6, use_dilation=True)
|
|
|
|
def _download(self) -> None:
|
|
model_info = InferSession.get_model_url(
|
|
FileInfo(
|
|
engine_type=EngineType.ONNXRUNTIME,
|
|
ocr_version=OCRVersion.PPOCRV5,
|
|
task_type=TaskType.DET,
|
|
lang_type=LangDet.CH,
|
|
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
|
)
|
|
)
|
|
download_params = DownloadFileInput(
|
|
file_url=model_info["model_dir"],
|
|
sha256=model_info["SHA256"],
|
|
save_path=self.model_path,
|
|
logger=log,
|
|
)
|
|
DownloadFile.run(download_params)
|
|
|
|
def _load(self) -> ModelSession:
|
|
# TODO: support other runtime sessions
|
|
return OrtSession(self.model_path)
|
|
|
|
def _predict(
|
|
self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast"
|
|
) -> TextDetectionOutput:
|
|
width, height = inputs.size
|
|
if width < 32 or height < 32:
|
|
return self._empty
|
|
|
|
image, content = self._transform(inputs, maxResolution)
|
|
input_name = self.session.get_inputs()[0].name
|
|
probs = self.session.run(None, {input_name: image})[0][0]
|
|
if probs.ndim == 3:
|
|
probs = probs[0]
|
|
boxes, scores = self.postprocess(probs[: content[0], : content[1]], (height, width), minScore, scoreMode)
|
|
if len(boxes) == 0:
|
|
return self._empty
|
|
order = self.reading_order(boxes)
|
|
return {"boxes": boxes[order], "scores": scores[order]}
|
|
|
|
@cached_property
|
|
def raw_input(self) -> bool:
|
|
return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally
|
|
|
|
def _transform(self, img: Image.Image, max_resolution: int) -> tuple[NDArray[Any], tuple[int, int]]:
|
|
ratio = min(max_resolution / min(img.height, img.width), 1.0)
|
|
resize_h = max(self._round32(img.height * ratio), 32)
|
|
resize_w = max(self._round32(img.width * ratio), 32)
|
|
resized = img.resize((resize_w, resize_h), resample=Image.Resampling.LANCZOS)
|
|
if resized.mode != "RGB":
|
|
resized = resized.convert("RGB")
|
|
array = np.asarray(resized)
|
|
if self.raw_input:
|
|
return array[None], (resize_h, resize_w)
|
|
# reverse plane order gets the BGR swap and the CHW transpose for free
|
|
out = np.empty((1, 3, resize_h, resize_w), dtype=np.float32)
|
|
for channel in range(3):
|
|
plane = out[0, 2 - channel]
|
|
np.multiply(array[:, :, channel], self.scale, out=plane)
|
|
plane -= 1.0
|
|
return out, (resize_h, resize_w)
|
|
|
|
@staticmethod
|
|
def _round32(value: float) -> int:
|
|
return int(round(value / 32) * 32)
|
|
|
|
def reading_order(self, boxes: NDArray[np.float32]) -> NDArray[np.intp]:
|
|
"""Indices of `boxes` top to bottom, then left to right within a line."""
|
|
y_order = np.argsort(boxes[:, 0, 1], kind="stable")
|
|
lines = np.zeros(len(boxes), dtype=np.int32) # indexed in y-sorted order, not box order
|
|
np.cumsum(np.diff(boxes[y_order, 0, 1]) >= 10, out=lines[1:]) # ascending, so the diffs are already positive
|
|
order: NDArray[np.intp] = y_order[np.argsort(lines * 1e6 + boxes[y_order, 0, 0], kind="stable")]
|
|
return order
|