refactor(ml): direct ocr processing (#31016)
This commit is contained in:
Mert
@@ -0,0 +1,40 @@
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, Self
|
||||
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
|
||||
def greedy(outputs: Sequence[NDArray[Any]]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
|
||||
"""Models with an in-graph argmax head emit (indices, probs); the rest emit raw logits."""
|
||||
if len(outputs) == 2:
|
||||
return outputs[0], outputs[1]
|
||||
(probs,) = outputs
|
||||
indices = probs.argmax(axis=2)
|
||||
return indices.astype(np.int32), np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0]
|
||||
|
||||
|
||||
class CtcDecoder:
|
||||
def __init__(self, charset: list[str]):
|
||||
self.charset = ["", *charset, " "] # PP-OCR: blank at 0, space last
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, charset_path: Path) -> Self:
|
||||
if not charset_path.is_file():
|
||||
raise FileNotFoundError(f"Recognition charset not found: {charset_path}")
|
||||
return cls(charset_path.read_text(encoding="utf-8").splitlines())
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.charset)
|
||||
|
||||
def decode(self, indices: NDArray[np.int32], probs: NDArray[np.float32]) -> tuple[list[str], NDArray[np.float32]]:
|
||||
keep = np.empty(indices.shape, dtype=bool)
|
||||
keep[:, 0] = True
|
||||
np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:]) # repeats before blanks
|
||||
keep &= indices != 0
|
||||
|
||||
scores: NDArray[np.float32] = np.where(keep, probs, 0).sum(1)
|
||||
scores /= np.maximum(keep.sum(1), 1) # an all-blank row sums to 0, so it stays 0
|
||||
texts = ["".join(map(self.charset.__getitem__, row[k])) for row, k in zip(indices, keep)]
|
||||
return texts, scores
|
||||
@@ -1,10 +1,9 @@
|
||||
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.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
|
||||
@@ -15,6 +14,7 @@ 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
|
||||
|
||||
|
||||
@@ -22,23 +22,14 @@ class TextDetector(InferenceModel):
|
||||
depends = []
|
||||
identity = (ModelType.DETECTION, ModelTask.OCR)
|
||||
|
||||
def __init__(self, model_name: str, min_score: float = 0.5, **model_kwargs: Any) -> None:
|
||||
super().__init__(model_name.split("__")[-1], **model_kwargs, model_format=ModelFormat.ONNX)
|
||||
self.max_resolution = 736
|
||||
self.mean = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
self.std_inv = np.float32(1.0) / (np.array([0.5, 0.5, 0.5], dtype=np.float32) * 255.0)
|
||||
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,
|
||||
box_thresh=model_kwargs.get("minScore", min_score),
|
||||
max_candidates=1000,
|
||||
unclip_ratio=1.6,
|
||||
use_dilation=True,
|
||||
score_mode="fast",
|
||||
)
|
||||
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(
|
||||
@@ -62,65 +53,54 @@ class TextDetector(InferenceModel):
|
||||
# TODO: support other runtime sessions
|
||||
return OrtSession(self.model_path)
|
||||
|
||||
# partly adapted from RapidOCR
|
||||
# TODO: take the options as parameters instead of applying them through `configure`
|
||||
def _predict(self, inputs: Image.Image, **kwargs: Any) -> TextDetectionOutput:
|
||||
w, h = inputs.size
|
||||
if w < 32 or h < 32:
|
||||
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
|
||||
out = self.session.run(None, {"x": self._transform(inputs)})[0]
|
||||
boxes, scores = self.postprocess(out, (h, w))
|
||||
|
||||
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
|
||||
return {
|
||||
"boxes": self.sorted_boxes(boxes),
|
||||
"scores": np.array(scores, dtype=np.float32),
|
||||
}
|
||||
order = self.reading_order(boxes)
|
||||
return {"boxes": boxes[order], "scores": scores[order]}
|
||||
|
||||
# adapted from RapidOCR
|
||||
def _transform(self, img: Image.Image) -> NDArray[np.float32]:
|
||||
if img.height < img.width:
|
||||
ratio = float(self.max_resolution) / img.height
|
||||
else:
|
||||
ratio = float(self.max_resolution) / img.width
|
||||
ratio = min(ratio, 1.0)
|
||||
@cached_property
|
||||
def raw_input(self) -> bool:
|
||||
return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally
|
||||
|
||||
resize_h = int(img.height * ratio)
|
||||
resize_w = int(img.width * ratio)
|
||||
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)
|
||||
|
||||
resize_h = int(round(resize_h / 32) * 32)
|
||||
resize_w = int(round(resize_w / 32) * 32)
|
||||
resized_img = img.resize((int(resize_w), int(resize_h)), resample=Image.Resampling.LANCZOS)
|
||||
@staticmethod
|
||||
def _round32(value: float) -> int:
|
||||
return int(round(value / 32) * 32)
|
||||
|
||||
img_np: NDArray[np.float32] = cv2.cvtColor(np.array(resized_img, dtype=np.float32), cv2.COLOR_RGB2BGR) # type: ignore
|
||||
img_np -= self.mean
|
||||
img_np *= self.std_inv
|
||||
img_np = np.transpose(img_np, (2, 0, 1))
|
||||
return np.expand_dims(img_np, axis=0)
|
||||
|
||||
def sorted_boxes(self, dt_boxes: NDArray[np.float32]) -> NDArray[np.float32]:
|
||||
if len(dt_boxes) == 0:
|
||||
return dt_boxes
|
||||
|
||||
# Sort by y, then identify lines, then sort by (line, x)
|
||||
y_order = np.argsort(dt_boxes[:, 0, 1], kind="stable")
|
||||
sorted_y = dt_boxes[y_order, 0, 1]
|
||||
|
||||
line_ids = np.empty(len(dt_boxes), dtype=np.int32)
|
||||
line_ids[0] = 0
|
||||
np.cumsum(np.abs(np.diff(sorted_y)) >= 10, out=line_ids[1:])
|
||||
|
||||
# Create composite sort key for final ordering
|
||||
# Shift line_ids by large factor, add x for tie-breaking
|
||||
sort_key = line_ids[y_order] * 1e6 + dt_boxes[y_order, 0, 0]
|
||||
final_order = np.argsort(sort_key, kind="stable")
|
||||
sorted_boxes: NDArray[np.float32] = dt_boxes[y_order[final_order]]
|
||||
return sorted_boxes
|
||||
|
||||
def configure(self, **kwargs: Any) -> None:
|
||||
if (max_resolution := kwargs.get("maxResolution")) is not None:
|
||||
self.max_resolution = max_resolution
|
||||
if (min_score := kwargs.get("minScore")) is not None:
|
||||
self.postprocess.box_thresh = min_score
|
||||
if (score_mode := kwargs.get("scoreMode")) is not None:
|
||||
self.postprocess.score_mode = score_mode
|
||||
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
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
"""Vendored from RapidOCR (Apache-2.0, rapidocr/ch_ppocr_det/utils.py @ v3.4).
|
||||
|
||||
`__call__` takes one image's [H, W] probability map rather than the raw [B, 1, H, W], and
|
||||
`unclip` grows the min-area rect directly instead of offsetting a polygon via pyclipper/shapely.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
|
||||
def order_points_clockwise(points: NDArray[np.float32]) -> NDArray[np.float32]:
|
||||
"""Order each group of four points as top-left, top-right, bottom-right, bottom-left."""
|
||||
by_x = np.take_along_axis(points, np.argsort(points[..., 0], axis=-1)[..., None], axis=-2)
|
||||
pairs = by_x.reshape(*points.shape[:-2], 2, 2, 2) # [left, right][member][xy]
|
||||
first, second = pairs[..., 0, :], pairs[..., 1, :]
|
||||
swap = (first[..., 1] > second[..., 1])[..., None] # sorting two points by y is one comparison
|
||||
top, bottom = np.where(swap, second, first), np.where(swap, first, second)
|
||||
return np.concatenate([top, bottom[..., ::-1, :]], axis=-2)
|
||||
|
||||
|
||||
def _span(values: cv2.typing.MatLike, limit: int) -> tuple[int, int]:
|
||||
lo, hi = math.floor(values.min()), math.ceil(values.max())
|
||||
return min(max(lo, 0), limit - 1), min(max(hi, 0), limit - 1)
|
||||
|
||||
|
||||
class DBPostProcess:
|
||||
"""The post process for Differentiable Binarization (DB)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
thresh: float = 0.3,
|
||||
max_candidates: int = 1000,
|
||||
unclip_ratio: float = 2.0,
|
||||
use_dilation: bool = False,
|
||||
):
|
||||
self.thresh = thresh
|
||||
self.max_candidates = max_candidates
|
||||
self.unclip_ratio = unclip_ratio
|
||||
self.min_size = 3
|
||||
self.dilation_kernel = np.ones((2, 2), dtype=np.uint8) if use_dilation else None
|
||||
|
||||
def __call__(
|
||||
self, probs: NDArray[np.float32], ori_shape: tuple[int, int], box_thresh: float, score_mode: str = "fast"
|
||||
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
|
||||
mask: NDArray[np.uint8] = (probs > self.thresh).view(np.uint8) # bool is already 0/1 bytes, so this is free
|
||||
if self.dilation_kernel is not None:
|
||||
cv2.dilate(mask, self.dilation_kernel, dst=mask) # OpenCV buffers rows, so this needs no second image
|
||||
boxes, scores = self.boxes_from_bitmap(probs, mask, ori_shape[1], ori_shape[0], box_thresh, score_mode)
|
||||
return self.filter_det_res(boxes, scores, ori_shape[0], ori_shape[1])
|
||||
|
||||
def boxes_from_bitmap(
|
||||
self,
|
||||
pred: NDArray[np.float32],
|
||||
bitmap: NDArray[np.uint8],
|
||||
dest_width: int,
|
||||
dest_height: int,
|
||||
box_thresh: float,
|
||||
score_mode: str,
|
||||
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
|
||||
height, width = bitmap.shape
|
||||
contours, _ = cv2.findContours(bitmap, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
boxes, scores = [], []
|
||||
for contour in contours[: self.max_candidates]:
|
||||
rect = cv2.minAreaRect(contour)
|
||||
if min(rect[1]) < self.min_size:
|
||||
continue
|
||||
|
||||
score = self.box_score(pred, cv2.boxPoints(rect) if score_mode == "fast" else contour)
|
||||
if box_thresh > score:
|
||||
continue
|
||||
|
||||
box, sside = self.unclip(rect)
|
||||
if sside < self.min_size + 2:
|
||||
continue
|
||||
|
||||
boxes.append(box)
|
||||
scores.append(score)
|
||||
|
||||
if not boxes:
|
||||
return np.empty((0, 4, 2), dtype=np.float32), np.empty(0, dtype=np.float32)
|
||||
scaled = np.stack(boxes)
|
||||
scaled *= [dest_width / width, dest_height / height] # in place, or the list would promote to float64
|
||||
np.clip(np.round(scaled, out=scaled), 0, [dest_width, dest_height], out=scaled)
|
||||
return scaled, np.array(scores, dtype=np.float32)
|
||||
|
||||
@staticmethod
|
||||
def box_score(probs: NDArray[np.float32], poly: cv2.typing.MatLike) -> float:
|
||||
"""Mean probability inside `poly`: the min-area rect for "fast", the contour for "slow"."""
|
||||
h, w = probs.shape[:2]
|
||||
xs, ys = poly.reshape(-1, 2).T
|
||||
xmin, xmax = _span(xs, w)
|
||||
ymin, ymax = _span(ys, h)
|
||||
|
||||
mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
|
||||
points = poly.reshape(1, -1, 2).astype(np.int32, copy=False) # read-only, so aliasing a contour is fine
|
||||
cv2.fillPoly(mask, points, 1, offset=(-xmin, -ymin)) # type: ignore[call-overload]
|
||||
return cv2.mean(probs[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
|
||||
|
||||
def unclip(self, rect: cv2.typing.RotatedRect) -> tuple[NDArray[np.float32], float]:
|
||||
width, height = rect[1]
|
||||
grow = width * height * self.unclip_ratio / (width + height)
|
||||
points = cv2.boxPoints((rect[0], (width + grow, height + grow), rect[2])).astype(np.float32, copy=False)
|
||||
return points, min(width, height) + grow
|
||||
|
||||
def filter_det_res(
|
||||
self, dt_boxes: NDArray[np.float32], scores: NDArray[np.float32], img_height: int, img_width: int
|
||||
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
|
||||
if len(dt_boxes) == 0:
|
||||
return dt_boxes, scores
|
||||
|
||||
boxes = order_points_clockwise(dt_boxes)
|
||||
np.clip(boxes, 0, [img_width - 1, img_height - 1], out=boxes)
|
||||
sides = boxes[:, 1:4:2] - boxes[:, :1] # top edge and left edge
|
||||
np.square(sides, out=sides)
|
||||
keep = sides.sum(2).min(1) >= 16 # coordinates are whole numbers, so int(length) > 3 is length**2 >= 16
|
||||
return boxes[keep], scores[keep]
|
||||
@@ -1,40 +1,45 @@
|
||||
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.ch_ppocr_rec import TextRecInput
|
||||
from rapidocr.ch_ppocr_rec import TextRecognizer as RapidTextRecognizer
|
||||
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, LangRec, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import EngineType, OCRVersion, TaskType
|
||||
from rapidocr.utils.typings import ModelType as RapidModelType
|
||||
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.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 .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, min_score: float = 0.9, **model_kwargs: Any) -> None:
|
||||
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.min_score = model_kwargs.get("minScore", min_score)
|
||||
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),
|
||||
}
|
||||
VisRes.__init__ = lambda self, **kwargs: None # pyright: ignore[reportAttributeAccessIssue]
|
||||
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(
|
||||
@@ -57,100 +62,115 @@ class TextRecognizer(InferenceModel):
|
||||
def _load(self) -> ModelSession:
|
||||
# TODO: support other runtimes
|
||||
session = OrtSession(self.model_path)
|
||||
max_batch_size = settings.max_batch_size and settings.max_batch_size.ocr
|
||||
self.model = RapidTextRecognizer(
|
||||
OcrOptions(
|
||||
session=session.session,
|
||||
rec_batch_num=max_batch_size if max_batch_size else 6,
|
||||
rec_img_shape=(3, 48, 320),
|
||||
lang_type=self.language,
|
||||
model_root_dir=self.cache_dir,
|
||||
)
|
||||
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
|
||||
|
||||
# TODO: take the options as parameters instead of applying them through `configure`
|
||||
def _predict(self, img: Image.Image, texts: TextDetectionOutput, **kwargs: Any) -> TextRecognitionOutput:
|
||||
@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
|
||||
rec = self.model(TextRecInput(img=self.get_crop_img_list(img, boxes)))
|
||||
if rec.txts is None:
|
||||
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
|
||||
|
||||
text_scores = np.array(rec.scores)
|
||||
valid_text_score_idx = text_scores > self.min_score
|
||||
valid_score_idx_list = valid_text_score_idx.tolist()
|
||||
valid = score_list > minScore
|
||||
valid_list = valid.tolist()
|
||||
return {
|
||||
"box": boxes.reshape(-1, 8)[valid_text_score_idx].reshape(-1),
|
||||
"text": [rec.txts[i] for i in range(len(rec.txts)) if valid_score_idx_list[i]],
|
||||
"boxScore": box_scores[valid_text_score_idx],
|
||||
"textScore": text_scores[valid_text_score_idx],
|
||||
"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 get_crop_img_list(self, img: Image.Image, boxes: NDArray[np.float32]) -> list[NDArray[np.uint8]]:
|
||||
img_crop_width = np.maximum(
|
||||
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)
|
||||
img_crop_height = np.maximum(
|
||||
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((img_crop_width.shape[0], 4, 2), dtype=np.float32)
|
||||
pts_std[:, 1:3, 0] = img_crop_width[:, None]
|
||||
pts_std[:, 2:4, 1] = img_crop_height[:, None]
|
||||
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)
|
||||
|
||||
img_crop_sizes = np.stack([img_crop_width, img_crop_height], axis=1)
|
||||
all_coeffs = self._get_perspective_transform(pts_std, boxes)
|
||||
imgs: list[NDArray[np.uint8]] = []
|
||||
for coeffs, dst_size in zip(all_coeffs, img_crop_sizes):
|
||||
dst_img = img.transform(
|
||||
size=tuple(dst_size),
|
||||
method=Image.Transform.PERSPECTIVE,
|
||||
data=tuple(coeffs),
|
||||
resample=Image.Resampling.BICUBIC,
|
||||
)
|
||||
|
||||
dst_width, dst_height = dst_img.size
|
||||
if dst_height * 1.0 / dst_width >= 1.5:
|
||||
dst_img = dst_img.rotate(90, expand=True)
|
||||
imgs.append(pil_to_cv2(dst_img))
|
||||
|
||||
return imgs
|
||||
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, 9), dtype=np.float32)
|
||||
A = np.zeros((N, 8, 8), dtype=np.float32)
|
||||
b = np.empty((N, 8, 1), dtype=np.float32)
|
||||
|
||||
# Fill even rows (0, 2, 4, 6): [x, y, 1, 0, 0, 0, -u*x, -u*y, -u]
|
||||
# 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
|
||||
A[:, ::2, 8] = -u
|
||||
b[:, ::2, 0] = u
|
||||
|
||||
# Fill odd rows (1, 3, 5, 7): [0, 0, 0, x, y, 1, -v*x, -v*y, -v]
|
||||
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
|
||||
A[:, 1::2, 8] = -v
|
||||
b[:, 1::2, 0] = v
|
||||
|
||||
# Solve using SVD for all matrices at once
|
||||
_, _, Vt = np.linalg.svd(A)
|
||||
H = Vt[:, -1, :].reshape(N, 3, 3)
|
||||
H = H / H[:, 2:3, 2:3]
|
||||
|
||||
# Extract the 8 coefficients for each transformation
|
||||
return np.column_stack(
|
||||
[H[:, 0, 0], H[:, 0, 1], H[:, 0, 2], H[:, 1, 0], H[:, 1, 1], H[:, 1, 2], H[:, 2, 0], H[:, 2, 1]]
|
||||
) # pyright: ignore[reportReturnType]
|
||||
|
||||
def configure(self, **kwargs: Any) -> None:
|
||||
self.min_score = kwargs.get("minScore", self.min_score)
|
||||
# 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
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
from typing import Any, Iterable
|
||||
from collections.abc import Iterable
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
from rapidocr.utils.typings import EngineType, LangRec
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
||||
@@ -16,12 +15,3 @@ class TextRecognitionOutput(TypedDict):
|
||||
boxScore: npt.NDArray[np.float32]
|
||||
text: Iterable[str]
|
||||
textScore: npt.NDArray[np.float32]
|
||||
|
||||
|
||||
# RapidOCR expects `engine_type`, `lang_type`, and `font_path` to be attributes
|
||||
class OcrOptions(dict[str, Any]):
|
||||
def __init__(self, lang_type: LangRec | None = None, **options: Any) -> None:
|
||||
super().__init__(**options)
|
||||
self.engine_type = EngineType.ONNXRUNTIME
|
||||
self.lang_type = lang_type
|
||||
self.font_path = None
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from collections.abc import Sequence
|
||||
from enum import Enum
|
||||
from typing import Any, Literal, Protocol, TypeGuard, TypeVar
|
||||
|
||||
@@ -70,13 +71,17 @@ class ModelSession(Protocol):
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, npt.NDArray[np.float32]] | dict[str, npt.NDArray[np.int32]],
|
||||
input_feed: dict[str, npt.NDArray[np.float32]]
|
||||
| dict[str, npt.NDArray[np.int32]]
|
||||
| dict[str, npt.NDArray[np.uint8]],
|
||||
run_options: Any = None,
|
||||
) -> list[npt.NDArray[np.float32]]: ...
|
||||
|
||||
def get_inputs(self) -> list[SessionNode]: ...
|
||||
def get_inputs(self) -> Sequence[SessionNode]: ...
|
||||
|
||||
def get_outputs(self) -> list[SessionNode]: ...
|
||||
def get_outputs(self) -> Sequence[SessionNode]: ...
|
||||
|
||||
def get_metadata(self) -> dict[str, str]: ...
|
||||
|
||||
|
||||
class HasProfiling(Protocol):
|
||||
|
||||
@@ -43,10 +43,13 @@ class AnnSession:
|
||||
shapes = self.ann.output_shapes[self.model]
|
||||
return [AnnNode(f"output.{i + 1}", s) for i, s in enumerate(shapes)]
|
||||
|
||||
def get_metadata(self) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
inputs: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from typing import Any
|
||||
@@ -39,9 +40,7 @@ def _migraphx_mark_compiled_input(key: tuple[str, MigraphxInputSignature]) -> No
|
||||
_migraphx_compiled_inputs.add(key)
|
||||
|
||||
|
||||
def _migraphx_input_signature(
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
|
||||
) -> MigraphxInputSignature:
|
||||
def _migraphx_input_signature(input_feed: dict[str, NDArray[Any]]) -> MigraphxInputSignature:
|
||||
return tuple((name, str(value.dtype), tuple(value.shape)) for name, value in sorted(input_feed.items()))
|
||||
|
||||
|
||||
@@ -66,18 +65,22 @@ class OrtSession:
|
||||
sess_options=self.sess_options,
|
||||
)
|
||||
|
||||
def get_inputs(self) -> list[SessionNode]:
|
||||
inputs: list[SessionNode] = self.session.get_inputs()
|
||||
def get_inputs(self) -> Sequence[SessionNode]:
|
||||
inputs: Sequence[SessionNode] = self.session.get_inputs()
|
||||
return inputs
|
||||
|
||||
def get_outputs(self) -> list[SessionNode]:
|
||||
outputs: list[SessionNode] = self.session.get_outputs()
|
||||
def get_outputs(self) -> Sequence[SessionNode]:
|
||||
outputs: Sequence[SessionNode] = self.session.get_outputs()
|
||||
return outputs
|
||||
|
||||
def get_metadata(self) -> dict[str, str]:
|
||||
metadata: dict[str, str] = self.session.get_modelmeta().custom_metadata_map
|
||||
return metadata
|
||||
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
if "MIGraphXExecutionProvider" in self.providers:
|
||||
|
||||
@@ -54,10 +54,13 @@ class RknnSession:
|
||||
def get_outputs(self) -> list[SessionNode]:
|
||||
return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["output"].items()]
|
||||
|
||||
def get_metadata(self) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
def run(
|
||||
self,
|
||||
output_names: list[str] | None,
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]],
|
||||
input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
input_data: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]
|
||||
|
||||
@@ -27,7 +27,6 @@ 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.schemas import ModelFormat, ModelPrecision, ModelTask, ModelType
|
||||
from immich_ml.sessions.ann import AnnSession
|
||||
from immich_ml.sessions.ort import OrtSession
|
||||
@@ -1079,76 +1078,69 @@ class TestFaceRecognition:
|
||||
|
||||
|
||||
class TestOcr:
|
||||
def test_set_det_min_score(self, path: mock.Mock) -> None:
|
||||
def test_det_min_score_is_per_request(self, path: mock.Mock) -> None:
|
||||
path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
|
||||
text_detector = TextDetector("PP-OCRv5_mobile", cache_dir="test_cache")
|
||||
probs = np.zeros((1, 1, 64, 64), dtype=np.float32)
|
||||
probs[..., 16:32, 8:56] = 0.6
|
||||
text_detector.session = mock.Mock()
|
||||
text_detector.session.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 64, 64))]
|
||||
text_detector.session.run.return_value = [probs]
|
||||
image = Image.new("RGB", (64, 64))
|
||||
|
||||
text_detector = TextDetector("PP-OCRv5_mobile", min_score=0.8, cache_dir="test_cache")
|
||||
assert len(text_detector._predict(image, minScore=0.5)["boxes"]) == 1
|
||||
assert len(text_detector._predict(image, minScore=0.9)["boxes"]) == 0
|
||||
# the default must be unaffected by the request that just ran
|
||||
assert len(text_detector._predict(image)["boxes"]) == 1
|
||||
|
||||
assert text_detector.postprocess.box_thresh == 0.8
|
||||
|
||||
def test_set_rec_min_score(self, path: mock.Mock) -> None:
|
||||
def test_rec_min_score_is_per_request(self, path: mock.Mock, mocker: MockerFixture) -> None:
|
||||
path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
|
||||
text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache")
|
||||
text_recognizer.session = mock.Mock()
|
||||
text_recognizer.session.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 48, 96))]
|
||||
text_recognizer.session.run.return_value = [np.zeros((1, 4, 8), dtype=np.float32)]
|
||||
text_recognizer.decoder = mock.Mock()
|
||||
text_recognizer.decoder.decode.return_value = (["hello"], np.array([0.8], dtype=np.float32))
|
||||
mocker.patch.object(text_recognizer, "_crop", return_value=np.zeros((48, 96, 3), dtype=np.uint8))
|
||||
image = Image.new("RGB", (100, 50))
|
||||
box = np.array([[[0, 0], [96, 0], [96, 48], [0, 48]]], dtype=np.float32)
|
||||
|
||||
text_recognizer = TextRecognizer("PP-OCRv5_mobile", min_score=0.8, cache_dir="test_cache")
|
||||
def texts() -> Any: # _predict normalizes the boxes in place, so each call needs its own
|
||||
return {"boxes": box.copy(), "scores": np.array([0.9], dtype=np.float32)}
|
||||
|
||||
assert text_recognizer.min_score == 0.8
|
||||
assert text_recognizer._predict(image, texts(), minScore=0.7)["text"] == ["hello"]
|
||||
# the default (0.9) rejects a 0.8 score, and must be unaffected by the 0.7 request
|
||||
assert text_recognizer._predict(image, texts())["text"] == []
|
||||
|
||||
def test_set_rec_set_default_max_batch_size(
|
||||
self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
|
||||
) -> None:
|
||||
path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
|
||||
mocker.patch("immich_ml.models.base.InferenceModel.download")
|
||||
rapid_recognizer = mocker.patch("immich_ml.models.ocr.recognition.RapidTextRecognizer")
|
||||
|
||||
text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache")
|
||||
text_recognizer.load()
|
||||
|
||||
rapid_recognizer.assert_called_once_with(
|
||||
OcrOptions(
|
||||
session=ort_session.return_value,
|
||||
rec_batch_num=6,
|
||||
rec_img_shape=(3, 48, 320),
|
||||
model_root_dir=text_recognizer.cache_dir,
|
||||
)
|
||||
)
|
||||
assert text_recognizer.batch_size == 6
|
||||
|
||||
def test_set_custom_max_batch_size(self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture) -> None:
|
||||
path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
|
||||
mocker.patch("immich_ml.models.base.InferenceModel.download")
|
||||
rapid_recognizer = mocker.patch("immich_ml.models.ocr.recognition.RapidTextRecognizer")
|
||||
mocker.patch.object(settings, "max_batch_size", MaxBatchSize(ocr=4))
|
||||
|
||||
text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache")
|
||||
text_recognizer.load()
|
||||
|
||||
rapid_recognizer.assert_called_once_with(
|
||||
OcrOptions(
|
||||
session=ort_session.return_value,
|
||||
rec_batch_num=4,
|
||||
rec_img_shape=(3, 48, 320),
|
||||
model_root_dir=text_recognizer.cache_dir,
|
||||
)
|
||||
)
|
||||
assert text_recognizer.batch_size == 4
|
||||
|
||||
def test_ignore_other_custom_max_batch_size(
|
||||
self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture
|
||||
) -> None:
|
||||
path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx"
|
||||
mocker.patch("immich_ml.models.base.InferenceModel.download")
|
||||
rapid_recognizer = mocker.patch("immich_ml.models.ocr.recognition.RapidTextRecognizer")
|
||||
mocker.patch.object(settings, "max_batch_size", MaxBatchSize(facial_recognition=3))
|
||||
|
||||
text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache")
|
||||
text_recognizer.load()
|
||||
|
||||
rapid_recognizer.assert_called_once_with(
|
||||
OcrOptions(
|
||||
session=ort_session.return_value,
|
||||
rec_batch_num=6,
|
||||
rec_img_shape=(3, 48, 320),
|
||||
model_root_dir=text_recognizer.cache_dir,
|
||||
)
|
||||
)
|
||||
assert text_recognizer.batch_size == 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user