mirror of
https://github.com/immich-app/immich
synced 2026-08-29 13:15:45 +00:00
176 lines
7.5 KiB
Python
176 lines
7.5 KiB
Python
import math
|
|
from functools import cached_property
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
from PIL import Image
|
|
from rapidocr import LangRec
|
|
from rapidocr.inference_engine.base import FileInfo, InferSession
|
|
from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
|
|
from rapidocr.utils.typings import EngineType, OCRVersion, TaskType
|
|
from rapidocr.utils.typings import ModelType as RapidModelType
|
|
|
|
from immich_ml.config import log, settings
|
|
from immich_ml.models.base import InferenceModel
|
|
from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
|
|
from immich_ml.sessions.ort import OrtSession
|
|
|
|
from .ctc import CtcDecoder, greedy
|
|
from .schemas import TextDetectionOutput, TextRecognitionOutput
|
|
|
|
REC_HEIGHT = 48
|
|
REC_BASE_RATIO = 320 / REC_HEIGHT # PP-OCR's rec_image_shape floor
|
|
SCALE = np.float32(1.0 / 127.5)
|
|
|
|
|
|
class TextRecognizer(InferenceModel):
|
|
depends = [(ModelType.DETECTION, ModelTask.OCR)]
|
|
identity = (ModelType.RECOGNITION, ModelTask.OCR)
|
|
|
|
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
|
self.language = LangRec[model_name.split("__")[0]] if "__" in model_name else LangRec.CH
|
|
self._empty: TextRecognitionOutput = {
|
|
"box": np.empty(0, dtype=np.float32),
|
|
"boxScore": np.empty(0, dtype=np.float32),
|
|
"text": [],
|
|
"textScore": np.empty(0, dtype=np.float32),
|
|
}
|
|
super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
|
|
max_batch_size = settings.max_batch_size and settings.max_batch_size.ocr
|
|
self.batch_size = max_batch_size if max_batch_size else 6
|
|
|
|
def _download(self) -> None:
|
|
model_info = InferSession.get_model_url(
|
|
FileInfo(
|
|
engine_type=EngineType.ONNXRUNTIME,
|
|
ocr_version=OCRVersion.PPOCRV5,
|
|
task_type=TaskType.REC,
|
|
lang_type=self.language,
|
|
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 runtimes
|
|
session = OrtSession(self.model_path)
|
|
self.decoder = (
|
|
CtcDecoder(character.splitlines())
|
|
if (character := session.get_metadata().get("character")) is not None
|
|
else CtcDecoder.from_file(self.model_dir / "charset.txt")
|
|
)
|
|
return session
|
|
|
|
@cached_property
|
|
def raw_input(self) -> bool:
|
|
return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally
|
|
|
|
def _predict(self, img: Image.Image, texts: TextDetectionOutput, minScore: float = 0.9) -> TextRecognitionOutput:
|
|
boxes, box_scores = texts["boxes"], texts["scores"]
|
|
if boxes.shape[0] == 0:
|
|
return self._empty
|
|
|
|
widths, heights, coeffs = self._crop_geometry(boxes)
|
|
upright = heights * 2 >= widths * 3 # PP-OCR stands a crop up when it is 1.5x taller than it is wide
|
|
ratios = np.where(upright, heights / widths, widths / heights)
|
|
order = np.argsort(ratios) # group similar widths so each batch pads the least it can
|
|
|
|
text_list: list[str] = [""] * len(order)
|
|
score_list = np.zeros(len(order), dtype=np.float32)
|
|
input_name = self.session.get_inputs()[0].name
|
|
for start in range(0, len(order), self.batch_size):
|
|
chunk = order[start : start + self.batch_size]
|
|
width = int(REC_HEIGHT * max(REC_BASE_RATIO, ratios[chunk[-1]]))
|
|
images: NDArray[Any] = (
|
|
np.full((len(chunk), REC_HEIGHT, width, 3), 127, dtype=np.uint8)
|
|
if self.raw_input
|
|
else np.zeros((len(chunk), 3, REC_HEIGHT, width), dtype=np.float32)
|
|
)
|
|
for i, index in enumerate(chunk):
|
|
crop = self._crop(img, coeffs[index], widths[index], heights[index], upright[index])
|
|
resized_w = max(1, min(width, math.ceil(REC_HEIGHT * ratios[index])))
|
|
if self.raw_input:
|
|
cv2.resize(crop, (resized_w, REC_HEIGHT), dst=images[i, :, :resized_w])
|
|
continue
|
|
resized = cv2.resize(crop, (resized_w, REC_HEIGHT))
|
|
view = images[i, :, :, :resized_w]
|
|
np.multiply(resized.transpose(2, 0, 1)[::-1], SCALE, out=view) # [::-1] is the RGB -> BGR swap
|
|
view -= 1.0
|
|
|
|
out_indices, out_probs = greedy(self.session.run(None, {input_name: images}))
|
|
chunk_texts, chunk_scores = self.decoder.decode(out_indices, out_probs)
|
|
for index, text, score in zip(chunk, chunk_texts, chunk_scores):
|
|
text_list[index] = text
|
|
score_list[index] = score
|
|
|
|
boxes[:, :, 0] /= img.width
|
|
boxes[:, :, 1] /= img.height
|
|
|
|
valid = score_list > minScore
|
|
valid_list = valid.tolist()
|
|
return {
|
|
"box": boxes.reshape(-1, 8)[valid].reshape(-1),
|
|
"text": [text for text, keep in zip(text_list, valid_list) if keep],
|
|
"boxScore": box_scores[valid],
|
|
"textScore": score_list[valid],
|
|
}
|
|
|
|
def _crop_geometry(
|
|
self, boxes: NDArray[np.float32]
|
|
) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
|
|
"""Upright crop size for each box, and the perspective coefficients that map the box onto it."""
|
|
widths = np.maximum(
|
|
np.linalg.norm(boxes[:, 1] - boxes[:, 0], axis=1), np.linalg.norm(boxes[:, 2] - boxes[:, 3], axis=1)
|
|
).astype(np.int32)
|
|
heights = np.maximum(
|
|
np.linalg.norm(boxes[:, 0] - boxes[:, 3], axis=1), np.linalg.norm(boxes[:, 1] - boxes[:, 2], axis=1)
|
|
).astype(np.int32)
|
|
pts_std = np.zeros((len(boxes), 4, 2), dtype=np.float32)
|
|
pts_std[:, 1:3, 0] = widths[:, None]
|
|
pts_std[:, 2:4, 1] = heights[:, None]
|
|
return widths, heights, self._get_perspective_transform(pts_std, boxes)
|
|
|
|
def _crop(
|
|
self, img: Image.Image, coeffs: NDArray[np.float32], width: int, height: int, upright: bool
|
|
) -> NDArray[np.uint8]:
|
|
crop = img.transform(
|
|
size=(width, height),
|
|
method=Image.Transform.PERSPECTIVE,
|
|
data=coeffs.tolist(),
|
|
resample=Image.Resampling.BICUBIC,
|
|
)
|
|
return np.asarray(crop.rotate(90, expand=True) if upright else crop)
|
|
|
|
def _get_perspective_transform(self, src: NDArray[np.float32], dst: NDArray[np.float32]) -> NDArray[np.float32]:
|
|
N = src.shape[0]
|
|
x, y = src[:, :, 0], src[:, :, 1]
|
|
u, v = dst[:, :, 0], dst[:, :, 1]
|
|
A = np.zeros((N, 8, 8), dtype=np.float32)
|
|
b = np.empty((N, 8, 1), dtype=np.float32)
|
|
|
|
# even rows [x, y, 1, 0, 0, 0, -u*x, -u*y] = u; odd [0, 0, 0, x, y, 1, -v*x, -v*y] = v
|
|
A[:, ::2, 0] = x
|
|
A[:, ::2, 1] = y
|
|
A[:, ::2, 2] = 1
|
|
A[:, ::2, 6] = -u * x
|
|
A[:, ::2, 7] = -u * y
|
|
b[:, ::2, 0] = u
|
|
|
|
A[:, 1::2, 3] = x
|
|
A[:, 1::2, 4] = y
|
|
A[:, 1::2, 5] = 1
|
|
A[:, 1::2, 6] = -v * x
|
|
A[:, 1::2, 7] = -v * y
|
|
b[:, 1::2, 0] = v
|
|
|
|
# h22 is fixed at 1, so this is a determined system rather than a null space
|
|
coeffs: NDArray[np.float32] = np.linalg.solve(A, b)[:, :, 0]
|
|
return coeffs
|