diff --git a/machine-learning/immich_ml/models/ocr/ctc.py b/machine-learning/immich_ml/models/ocr/ctc.py new file mode 100644 index 000000000..90959b1b9 --- /dev/null +++ b/machine-learning/immich_ml/models/ocr/ctc.py @@ -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 diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index 48dd87854..258058514 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/machine-learning/immich_ml/models/ocr/detection.py @@ -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 diff --git a/machine-learning/immich_ml/models/ocr/postprocess.py b/machine-learning/immich_ml/models/ocr/postprocess.py new file mode 100644 index 000000000..cdfeefc67 --- /dev/null +++ b/machine-learning/immich_ml/models/ocr/postprocess.py @@ -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] diff --git a/machine-learning/immich_ml/models/ocr/recognition.py b/machine-learning/immich_ml/models/ocr/recognition.py index ea7ff565d..dc445e760 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/machine-learning/immich_ml/models/ocr/recognition.py @@ -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 diff --git a/machine-learning/immich_ml/models/ocr/schemas.py b/machine-learning/immich_ml/models/ocr/schemas.py index 78e8619a0..85ae924f0 100644 --- a/machine-learning/immich_ml/models/ocr/schemas.py +++ b/machine-learning/immich_ml/models/ocr/schemas.py @@ -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 diff --git a/machine-learning/immich_ml/schemas.py b/machine-learning/immich_ml/schemas.py index d5695f5eb..0dcc540d0 100644 --- a/machine-learning/immich_ml/schemas.py +++ b/machine-learning/immich_ml/schemas.py @@ -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): diff --git a/machine-learning/immich_ml/sessions/ann/__init__.py b/machine-learning/immich_ml/sessions/ann/__init__.py index d3279f9a1..c516f9ba9 100644 --- a/machine-learning/immich_ml/sessions/ann/__init__.py +++ b/machine-learning/immich_ml/sessions/ann/__init__.py @@ -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()] diff --git a/machine-learning/immich_ml/sessions/ort.py b/machine-learning/immich_ml/sessions/ort.py index 579cafc1c..95a40782c 100644 --- a/machine-learning/immich_ml/sessions/ort.py +++ b/machine-learning/immich_ml/sessions/ort.py @@ -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: diff --git a/machine-learning/immich_ml/sessions/rknn/__init__.py b/machine-learning/immich_ml/sessions/rknn/__init__.py index 90f2e63ed..721fe5bbd 100644 --- a/machine-learning/immich_ml/sessions/rknn/__init__.py +++ b/machine-learning/immich_ml/sessions/rknn/__init__.py @@ -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()] diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index 567c6ab17..748a8840c 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -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