feat: ocr

.github/labeler.yml

@@ -31,7 +31,7 @@ documentation:
 🧠machine-learning:
   - changed-files:
       - any-glob-to-any-file:
-          - machine-learning/app/**
+          - machine-learning/**
 
 changelog:translation:
   - head-branch: ['^chore/translations$']

docker/docker-compose.dev.yml

@@ -122,7 +122,7 @@ services:
     ports:
       - 3003:3003
     volumes:
-      - ../machine-learning:/usr/src/app
+      - ../machine-learning/immich_ml:/usr/src/immich_ml
       - model-cache:/cache
     env_file:
       - .env

e2e/src/api/specs/server.e2e-spec.ts

@@ -113,6 +113,7 @@ describe('/server', () => {
         importFaces: false,
         oauth: false,
         oauthAutoLaunch: false,
+        ocr: false,
         passwordLogin: true,
         search: true,
         sidecar: true,

i18n/en.json

@@ -154,6 +154,18 @@
     "machine_learning_min_detection_score_description": "Minimum confidence score for a face to be detected from 0-1. Lower values will detect more faces but may result in false positives.",
     "machine_learning_min_recognized_faces": "Minimum recognized faces",
     "machine_learning_min_recognized_faces_description": "The minimum number of recognized faces for a person to be created. Increasing this makes Facial Recognition more precise at the cost of increasing the chance that a face is not assigned to a person.",
+    "machine_learning_ocr": "OCR",
+    "machine_learning_ocr_description": "Use machine learning to recognize text in images",
+    "machine_learning_ocr_enabled": "Enable OCR",
+    "machine_learning_ocr_enabled_description": "If disabled, images will not undergo text recognition.",
+    "machine_learning_ocr_max_resolution": "Maximum resolution",
+    "machine_learning_ocr_max_resolution_description": "Previews above this resolution will be resized while preserving aspect ratio. Higher values are more accurate, but take longer to process and use more memory.",
+    "machine_learning_ocr_min_detection_score": "Minimum detection score",
+    "machine_learning_ocr_min_detection_score_description": "Minimum confidence score for text to be detected from 0-1. Lower values will detect more text but may result in false positives.",
+    "machine_learning_ocr_min_recognition_score": "Minimum recognition score",
+    "machine_learning_ocr_min_score_recognition_description": "Minimum confidence score for detected text to be recognized from 0-1. Lower values will recognize more text but may result in false positives.",
+    "machine_learning_ocr_model": "OCR model",
+    "machine_learning_ocr_model_description": "Server models are more accurate than mobile models, but take longer to process and use more memory.",
     "machine_learning_settings": "Machine Learning Settings",
     "machine_learning_settings_description": "Manage machine learning features and settings",
     "machine_learning_smart_search": "Smart Search",
@@ -245,6 +257,7 @@
     "oauth_storage_quota_default_description": "Quota in GiB to be used when no claim is provided.",
     "oauth_timeout": "Request Timeout",
     "oauth_timeout_description": "Timeout for requests in milliseconds",
+    "ocr_job_description": "Use machine learning to recognize text in images",
     "password_enable_description": "Login with email and password",
     "password_settings": "Password Login",
     "password_settings_description": "Manage password login settings",
@@ -1436,6 +1449,7 @@
   "oauth": "OAuth",
   "obtainium_configurator": "Obtainium Configurator",
   "obtainium_configurator_instructions": "Use Obtainium to install and update the Android app directly from Immich GitHub's release. Create an API key and select a variant to create your Obtainium configuration link",
+  "ocr": "OCR",
   "official_immich_resources": "Official Immich Resources",
   "offline": "Offline",
   "offset": "Offset",
@@ -1710,6 +1724,8 @@
   "search_by_description_example": "Hiking day in Sapa",
   "search_by_filename": "Search by file name or extension",
   "search_by_filename_example": "i.e. IMG_1234.JPG or PNG",
+  "search_by_ocr": "Search by OCR",
+  "search_by_ocr_example": "Latte",
   "search_camera_make": "Search camera make...",
   "search_camera_model": "Search camera model...",
   "search_city": "Search city...",
@@ -1726,6 +1742,7 @@
   "search_filter_location_title": "Select location",
   "search_filter_media_type": "Media Type",
   "search_filter_media_type_title": "Select media type",
+  "search_filter_ocr": "Search by OCR",
   "search_filter_people_title": "Select people",
   "search_for": "Search for",
   "search_for_existing_person": "Search for existing person",

machine-learning/Dockerfile

@@ -141,7 +141,7 @@ FROM prod-${DEVICE} AS prod
 ARG DEVICE
 
 RUN apt-get update && \
-    apt-get install -y --no-install-recommends tini $(if ! [ "$DEVICE" = "openvino" ] && ! [ "$DEVICE" = "rocm" ]; then echo "libmimalloc2.0"; fi) && \
+    apt-get install -y --no-install-recommends tini ccache libgl1 libglib2.0-0 libgomp1 $(if ! [ "$DEVICE" = "openvino" ] && ! [ "$DEVICE" = "rocm" ]; then echo "libmimalloc2.0"; fi) && \
     apt-get autoremove -yqq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*

machine-learning/immich_ml/config.py

@@ -41,6 +41,7 @@ class PreloadModelData(BaseModel):
 
 class MaxBatchSize(BaseModel):
     facial_recognition: int | None = None
+    text_recognition: int | None = None
 
 
 class Settings(BaseSettings):

machine-learning/immich_ml/main.py

@@ -183,7 +183,9 @@ async def run_inference(payload: Image | str, entries: InferenceEntries) -> Infe
     response: InferenceResponse = {}
 
     async def _run_inference(entry: InferenceEntry) -> None:
-        model = await model_cache.get(entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl)
+        model = await model_cache.get(
+            entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl, **entry["options"]
+        )
         inputs = [payload]
         for dep in model.depends:
             try:

machine-learning/immich_ml/models/__init__.py

@@ -3,6 +3,8 @@
 from immich_ml.models.base import InferenceModel
 from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder
 from immich_ml.models.clip.visual import OpenClipVisualEncoder
+from immich_ml.models.ocr.detection import TextDetector
+from immich_ml.models.ocr.recognition import TextRecognizer
 from immich_ml.schemas import ModelSource, ModelTask, ModelType
 
 from .constants import get_model_source
@@ -28,6 +30,12 @@ def get_model_class(model_name: str, model_type: ModelType, model_task: ModelTas
         case ModelSource.INSIGHTFACE, ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION:
             return FaceRecognizer
 
+        case ModelSource.PADDLE, ModelType.DETECTION, ModelTask.OCR:
+            return TextDetector
+
+        case ModelSource.PADDLE, ModelType.RECOGNITION, ModelTask.OCR:
+            return TextRecognizer
+
         case _:
             raise ValueError(f"Unknown model combination: {source}, {model_type}, {model_task}")
 

machine-learning/immich_ml/models/base.py

@@ -38,9 +38,8 @@ def __init__(
 
     def download(self) -> None:
         if not self.cached:
-            log.info(
-                f"Downloading {self.model_type.replace('-', ' ')} model '{self.model_name}'. This may take a while."
-            )
+            model_type = self.model_type.replace("-", " ")
+            log.info(f"Downloading {model_type} model '{self.model_name}' to {self.model_path}. This may take a while.")
             self._download()
 
     def load(self) -> None:
@@ -58,7 +57,7 @@ def predict(self, *inputs: Any, **model_kwargs: Any) -> Any:
         self.load()
         if model_kwargs:
             self.configure(**model_kwargs)
-        return self._predict(*inputs, **model_kwargs)
+        return self._predict(*inputs)
 
     @abstractmethod
     def _predict(self, *inputs: Any, **model_kwargs: Any) -> Any: ...

machine-learning/immich_ml/models/clip/textual.py

@@ -19,7 +19,7 @@ class BaseCLIPTextualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.TEXTUAL, ModelTask.SEARCH)
 
-    def _predict(self, inputs: str, language: str | None = None, **kwargs: Any) -> str:
+    def _predict(self, inputs: str, language: str | None = None) -> str:
         tokens = self.tokenize(inputs, language=language)
         res: NDArray[np.float32] = self.session.run(None, tokens)[0][0]
         return serialize_np_array(res)

machine-learning/immich_ml/models/clip/visual.py

@@ -26,7 +26,7 @@ class BaseCLIPVisualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.VISUAL, ModelTask.SEARCH)
 
-    def _predict(self, inputs: Image.Image | bytes, **kwargs: Any) -> str:
+    def _predict(self, inputs: Image.Image | bytes) -> str:
         image = decode_pil(inputs)
         res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0]
         return serialize_np_array(res)

machine-learning/immich_ml/models/constants.py

@@ -75,6 +75,11 @@
 }
 
 
+_PADDLE_MODELS = {
+    "PP-OCRv5_server",
+    "PP-OCRv5_mobile",
+}
+
 SUPPORTED_PROVIDERS = [
     "CUDAExecutionProvider",
     "ROCMExecutionProvider",
@@ -159,4 +164,7 @@ def get_model_source(model_name: str) -> ModelSource | None:
     if cleaned_name in _OPENCLIP_MODELS:
         return ModelSource.OPENCLIP
 
+    if cleaned_name in _PADDLE_MODELS:
+        return ModelSource.PADDLE
+
     return None

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

@@ -24,7 +24,7 @@ def _load(self) -> ModelSession:
 
         return session
 
-    def _predict(self, inputs: NDArray[np.uint8] | bytes, **kwargs: Any) -> FaceDetectionOutput:
+    def _predict(self, inputs: NDArray[np.uint8] | bytes) -> FaceDetectionOutput:
         inputs = decode_cv2(inputs)
 
         bboxes, landmarks = self._detect(inputs)

machine-learning/immich_ml/models/facial_recognition/recognition.py

@@ -44,7 +44,7 @@ def _load(self) -> ModelSession:
         return session
 
     def _predict(
-        self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput, **kwargs: Any
+        self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput
     ) -> FacialRecognitionOutput:
         if faces["boxes"].shape[0] == 0:
             return []

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

@@ -0,0 +1,86 @@
+from typing import Any
+
+import numpy as np
+from PIL import Image
+from rapidocr.ch_ppocr_det import TextDetector as RapidTextDetector
+from rapidocr.inference_engine.base import FileInfo, InferSession
+from rapidocr.utils 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.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
+
+
+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.max_resolution = 736
+        self.min_score = 0.5
+        self.score_mode = "fast"
+        self._empty: TextDetectionOutput = {
+            "image": np.empty(0, dtype=np.float32),
+            "boxes": np.empty(0, dtype=np.float32),
+            "scores": np.empty(0, dtype=np.float32),
+        }
+
+    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
+        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
+
+    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:
+            return self._empty
+        return {
+            "image": results.img,
+            "boxes": np.array(results.boxes, dtype=np.float32),
+            "scores": np.array(results.scores, dtype=np.float32),
+        }
+
+    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
+        if (score_mode := kwargs.get("scoreMode")) is not None:
+            self.score_mode = score_mode
+            self.model.postprocess_op.score_mode = score_mode