fix(ml): ocr inputs not resized correctly (#23541)

* fix resizing, use pillow

* unused import

* linting

* lanczos

* optimizations

fused operations

unused import

machine-learning/immich_ml/models/ocr/detection.py

@@ -1,8 +1,10 @@
 from typing import Any
 
+import cv2
 import numpy as np
+from numpy.typing import NDArray
 from PIL import Image
-from rapidocr.ch_ppocr_det import TextDetector as RapidTextDetector
+from rapidocr.ch_ppocr_det.utils import DBPostProcess
 from rapidocr.inference_engine.base import FileInfo, InferSession
 from rapidocr.utils import DownloadFile, DownloadFileInput
 from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType
@@ -10,11 +12,10 @@ from rapidocr.utils.typings import ModelType as RapidModelType
 
 from immich_ml.config import log
 from immich_ml.models.base import InferenceModel
-from immich_ml.models.transforms import decode_cv2
 from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
 from immich_ml.sessions.ort import OrtSession
 
-from .schemas import OcrOptions, TextDetectionOutput
+from .schemas import TextDetectionOutput
 
 
 class TextDetector(InferenceModel):
@@ -24,13 +25,20 @@ class TextDetector(InferenceModel):
     def __init__(self, model_name: str, **model_kwargs: Any) -> None:
         super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
         self.max_resolution = 736
-        self.min_score = 0.5
-        self.score_mode = "fast"
+        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)
         self._empty: TextDetectionOutput = {
-            "image": np.empty(0, dtype=np.float32),
             "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", 0.5),
+            max_candidates=1000,
+            unclip_ratio=1.6,
+            use_dilation=True,
+            score_mode="fast",
+        )
 
     def _download(self) -> None:
         model_info = InferSession.get_model_url(
@@ -52,35 +60,65 @@ class TextDetector(InferenceModel):
 
     def _load(self) -> ModelSession:
         # TODO: support other runtime sessions
-        session = OrtSession(self.model_path)
-        self.model = RapidTextDetector(
-            OcrOptions(
-                session=session.session,
-                limit_side_len=self.max_resolution,
-                limit_type="min",
-                box_thresh=self.min_score,
-                score_mode=self.score_mode,
-            )
-        )
-        return session
+        return OrtSession(self.model_path)
 
-    def _predict(self, inputs: bytes | Image.Image) -> TextDetectionOutput:
-        results = self.model(decode_cv2(inputs))
-        if results.boxes is None or results.scores is None or results.img is None:
+    # partly adapted from RapidOCR
+    def _predict(self, inputs: Image.Image) -> TextDetectionOutput:
+        w, h = inputs.size
+        if w < 32 or h < 32:
+            return self._empty
+        out = self.session.run(None, {"x": self._transform(inputs)})[0]
+        boxes, scores = self.postprocess(out, (h, w))
+        if len(boxes) == 0:
             return self._empty
         return {
-            "image": results.img,
-            "boxes": np.array(results.boxes, dtype=np.float32),
-            "scores": np.array(results.scores, dtype=np.float32),
+            "boxes": self.sorted_boxes(boxes),
+            "scores": np.array(scores, dtype=np.float32),
         }
 
+    # 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
+
+        resize_h = int(img.height * ratio)
+        resize_w = int(img.width * ratio)
+
+        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)
+
+        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
-            self.model.limit_side_len = max_resolution
         if (min_score := kwargs.get("minScore")) is not None:
-            self.min_score = min_score
-            self.model.postprocess_op.box_thresh = min_score
+            self.postprocess.box_thresh = min_score
         if (score_mode := kwargs.get("scoreMode")) is not None:
-            self.score_mode = score_mode
-            self.model.postprocess_op.score_mode = score_mode
+            self.postprocess.score_mode = score_mode