refactor(ml): remove insightface dependency (#30631)

* remove insightface dependency

* update tests

machine-learning/immich_ml/models/base.py

@@ -57,7 +57,7 @@ class InferenceModel(ABC):
         self.load()
         if model_kwargs:
             self.configure(**model_kwargs)
-        return self._predict(*inputs)
+        return self._predict(*inputs, **model_kwargs)
 
     @abstractmethod
     def _predict(self, *inputs: Any, **model_kwargs: Any) -> Any: ...

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

@@ -0,0 +1,74 @@
+"""
+Host-side geometry for the fused face models.
+
+The decode, NMS and Umeyama similarity are ports of insightface (Apache-2.0, https://github.com/deepinsight/insightface)
+"""
+
+from functools import lru_cache
+
+import cv2
+import numpy as np
+from numpy.typing import NDArray
+
+from immich_ml.models.transforms import ensure_dims
+
+DET_SIZE = 640
+ALIGNED_SIZE = 112
+
+# the FPN levels for which the fused detector emits a (scores, boxes, kps) triple,
+# and the anchors each feature-map cell carries, laid out anchor-major
+DET_STRIDES = (8, 16, 32)
+ANCHORS_PER_CELL = 2
+
+# canonical ArcFace 5-point template for a 112x112 crop
+ARCFACE_DST = np.array(
+    [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]],
+    dtype=np.float32,
+)
+
+
+@lru_cache(maxsize=len(DET_STRIDES))
+def _anchor_centers(size: int, stride: int) -> NDArray[np.float32]:
+    ys, xs = np.mgrid[: size // stride, : size // stride]
+    centers = np.stack([xs, ys], axis=-1, dtype=np.float32).reshape(-1, 2) * stride
+    return np.repeat(centers, ANCHORS_PER_CELL, axis=0)
+
+
+def decode_scrfd(
+    heads: list[NDArray[np.float32]], size: int = DET_SIZE
+) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
+    heads = [ensure_dims(head, 4) for head in heads]
+    scores, boxes, kps = [], [], []
+    for level, stride in enumerate(DET_STRIDES):
+        centers = _anchor_centers(size, stride)
+        distance = heads[level + len(DET_STRIDES)][0] * stride
+        offsets = heads[level + 2 * len(DET_STRIDES)][0] * stride
+        scores.append(heads[level][0].squeeze(-1))
+        boxes.append(np.concatenate([centers[None] - distance[:, :, :2], centers[None] + distance[:, :, 2:]], axis=2))
+        kps.append(np.tile(centers, offsets.shape[2] // 2) + offsets)
+
+    return np.concatenate(scores, axis=1), np.concatenate(boxes, axis=1), np.concatenate(kps, axis=1)
+
+
+def nms(boxes: NDArray[np.float32], scores: NDArray[np.float32], threshold: float = 0.4) -> NDArray[np.intp]:
+    wh = np.column_stack([boxes[:, 0], boxes[:, 1], boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]])
+    keep = cv2.dnn.NMSBoxes(wh.tolist(), scores.tolist(), 0.0, threshold)  # NMSBoxes treats the inputs as Sequences
+    return np.asarray(keep, dtype=np.intp).reshape(-1)
+
+
+def umeyama(src: NDArray[np.float32], dst: NDArray[np.float32]) -> NDArray[np.float32]:
+    src_mean, dst_mean = src.mean(0), dst.mean(0)
+    src_c, dst_c = src - src_mean, dst - dst_mean
+    cov = dst_c.T @ src_c / len(src)
+    u, s, vt = np.linalg.svd(cov)
+    d = np.sign(np.linalg.det(u @ vt))
+    diag = np.diag([1.0, d])
+    rotation = u @ diag @ vt
+    scale = np.trace(np.diag(s) @ diag) / (src_c**2).sum() * len(src)
+    translation = dst_mean - scale * rotation @ src_mean
+    return np.hstack([scale * rotation, translation[:, None]], dtype=np.float32)
+
+
+def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32]) -> NDArray[np.float32]:
+    matrix = umeyama(kps, ARCFACE_DST)
+    return cv2.warpAffine(image, matrix, (ALIGNED_SIZE, ALIGNED_SIZE)).astype(np.float32)

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

@@ -1,41 +1,31 @@
-from typing import Any
-
 import numpy as np
-from insightface.model_zoo import RetinaFace
 from numpy.typing import NDArray
 
 from immich_ml.models.base import InferenceModel
-from immich_ml.models.transforms import decode_cv2
-from immich_ml.schemas import FaceDetectionOutput, ModelSession, ModelTask, ModelType
+from immich_ml.models.transforms import decode_pil, letterbox, normalize
+from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType
+
+from ._ops import DET_SIZE, decode_scrfd, nms
 
 
 class FaceDetector(InferenceModel):
     depends = []
     identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)
 
-    def __init__(self, model_name: str, min_score: float = 0.7, **model_kwargs: Any) -> None:
-        self.min_score = model_kwargs.pop("minScore", min_score)
-        super().__init__(model_name, **model_kwargs)
+    def _predict(self, inputs: NDArray[np.uint8] | bytes, minScore: float) -> FaceDetectionOutput:
+        canvas, scale = letterbox(decode_pil(inputs), DET_SIZE)
+        blob = normalize(canvas.astype(np.float32), mean=127.5, std=128).transpose(2, 0, 1)[None]
 
-    def _load(self) -> ModelSession:
-        session = self._make_session(self.model_path)
-        self.model = RetinaFace(session=session)
-        self.model.prepare(ctx_id=0, det_thresh=self.min_score, input_size=(640, 640))
+        input_name = self.session.get_inputs()[0].name
+        heads = self.session.run(None, {input_name: blob})
+        scores, boxes, kps = decode_scrfd(heads, DET_SIZE)
 
-        return session
+        candidates = scores >= minScore
+        scores, boxes, kps = scores[candidates], boxes[candidates] / scale, kps[candidates] / scale
+        keep = nms(boxes, scores)
 
-    def _predict(self, inputs: NDArray[np.uint8] | bytes) -> FaceDetectionOutput:
-        inputs = decode_cv2(inputs)
-
-        bboxes, landmarks = self._detect(inputs)
         return {
-            "boxes": bboxes[:, :4].round(),
-            "scores": bboxes[:, 4],
-            "landmarks": landmarks,
+            "boxes": boxes[keep].round(),
+            "scores": scores[keep],
+            "landmarks": kps[keep].reshape(-1, 5, 2),
         }
-
-    def _detect(self, inputs: NDArray[np.uint8] | bytes) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
-        return self.model.detect(inputs)  # type: ignore
-
-    def configure(self, **kwargs: Any) -> None:
-        self.model.det_thresh = kwargs.pop("minScore", self.model.det_thresh)

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

@@ -4,15 +4,13 @@ from typing import Any
 import numpy as np
 import onnx
 import onnxruntime as ort
-from insightface.model_zoo import ArcFaceONNX
-from insightface.utils.face_align import norm_crop
 from numpy.typing import NDArray
 from onnx.tools.update_model_dims import update_inputs_outputs_dims
 from PIL import Image
 
 from immich_ml.config import log, settings
 from immich_ml.models.base import InferenceModel
-from immich_ml.models.transforms import decode_cv2, serialize_np_array
+from immich_ml.models.transforms import decode_pil, normalize, serialize_np_array
 from immich_ml.schemas import (
     FaceDetectionOutput,
     FacialRecognitionOutput,
@@ -22,6 +20,8 @@ from immich_ml.schemas import (
     ModelType,
 )
 
+from ._ops import align_face
+
 
 class FaceRecognizer(InferenceModel):
     depends = [(ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)]
@@ -37,10 +37,6 @@ class FaceRecognizer(InferenceModel):
         if (not self.batch_size or self.batch_size > 1) and str(session.get_inputs()[0].shape[0]) != "batch":
             self._add_batch_axis(self.model_path)
             session = self._make_session(self.model_path)
-        self.model = ArcFaceONNX(
-            self.model_path_for_format(ModelFormat.ONNX).as_posix(),
-            session=session,
-        )
         return session
 
     def _predict(
@@ -48,20 +44,21 @@ class FaceRecognizer(InferenceModel):
     ) -> FacialRecognitionOutput:
         if faces["boxes"].shape[0] == 0:
             return []
-        inputs = decode_cv2(inputs)
-        cropped_faces = self._crop(inputs, faces)
-        embeddings = self._predict_batch(cropped_faces)
+        image = np.asarray(decode_pil(inputs), dtype=np.uint8)
+        crops = np.stack([align_face(image, kps).transpose(2, 0, 1) for kps in faces["landmarks"]])
+        embeddings = self._predict_batch(normalize(crops, mean=127.5, std=127.5))
         return self.postprocess(faces, embeddings)
 
-    def _predict_batch(self, cropped_faces: list[NDArray[np.uint8]]) -> NDArray[np.float32]:
-        if not self.batch_size or len(cropped_faces) <= self.batch_size:
-            embeddings: NDArray[np.float32] = self.model.get_feat(cropped_faces)
-            return embeddings
+    def _predict_batch(self, crops: NDArray[np.float32]) -> NDArray[np.float32]:
+        input_name = self.session.get_inputs()[0].name
+        if not self.batch_size or crops.shape[0] <= self.batch_size:
+            return self.session.run(None, {input_name: crops})[0]
 
-        batch_embeddings: list[NDArray[np.float32]] = []
-        for i in range(0, len(cropped_faces), self.batch_size):
-            batch_embeddings.append(self.model.get_feat(cropped_faces[i : i + self.batch_size]))
-        return np.concatenate(batch_embeddings, axis=0)
+        batches = [
+            self.session.run(None, {input_name: crops[i : i + self.batch_size]})[0]
+            for i in range(0, crops.shape[0], self.batch_size)
+        ]
+        return np.concatenate(batches, axis=0)
 
     def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32]) -> FacialRecognitionOutput:
         return [
@@ -73,9 +70,6 @@ class FaceRecognizer(InferenceModel):
             for (x1, y1, x2, y2), embedding, score in zip(faces["boxes"], embeddings, faces["scores"])
         ]
 
-    def _crop(self, image: NDArray[np.uint8], faces: FaceDetectionOutput) -> list[NDArray[np.uint8]]:
-        return [norm_crop(image, landmark) for landmark in faces["landmarks"]]
-
     def _add_batch_axis(self, model_path: Path) -> None:
         log.debug(f"Adding batch axis to model {model_path}")
         proto = onnx.load(model_path)

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

@@ -63,7 +63,8 @@ class TextDetector(InferenceModel):
         return OrtSession(self.model_path)
 
     # partly adapted from RapidOCR
-    def _predict(self, inputs: Image.Image) -> TextDetectionOutput:
+    # 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:
             return self._empty

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

@@ -69,7 +69,8 @@ class TextRecognizer(InferenceModel):
         )
         return session
 
-    def _predict(self, img: Image.Image, texts: TextDetectionOutput) -> TextRecognitionOutput:
+    # TODO: take the options as parameters instead of applying them through `configure`
+    def _predict(self, img: Image.Image, texts: TextDetectionOutput, **kwargs: Any) -> TextRecognitionOutput:
         boxes, box_scores = texts["boxes"], texts["scores"]
         if boxes.shape[0] == 0:
             return self._empty

machine-learning/immich_ml/models/transforms.py

@@ -1,6 +1,5 @@
 import string
 from io import BytesIO
-from typing import IO
 
 import cv2
 import numpy as np
@@ -12,11 +11,10 @@ _PIL_RESAMPLING_METHODS = {resampling.name.lower(): resampling for resampling in
 _PUNCTUATION_TRANS = str.maketrans("", "", string.punctuation)
 
 
-def resize_pil(img: Image.Image, size: int) -> Image.Image:
+def resize_pil(img: Image.Image, size: int, resample: Image.Resampling = Image.Resampling.BICUBIC) -> Image.Image:
     if img.width < img.height:
-        return img.resize((size, int((img.height / img.width) * size)), resample=Image.Resampling.BICUBIC)
-    else:
-        return img.resize((int((img.width / img.height) * size), size), resample=Image.Resampling.BICUBIC)
+        return img.resize((size, int((img.height / img.width) * size)), resample=resample)
+    return img.resize((int((img.width / img.height) * size), size), resample=resample)
 
 
 # https://stackoverflow.com/a/60883103
@@ -36,7 +34,9 @@ def to_numpy(img: Image.Image) -> NDArray[np.float32]:
 def normalize(
     img: NDArray[np.float32], mean: float | NDArray[np.float32], std: float | NDArray[np.float32]
 ) -> NDArray[np.float32]:
-    return (img - mean) / std
+    img *= 1.0 / std
+    img -= mean / std
+    return img
 
 
 def get_pil_resampling(resample: str) -> Image.Resampling:
@@ -47,26 +47,21 @@ def pil_to_cv2(image: Image.Image) -> NDArray[np.uint8]:
     return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # type: ignore
 
 
-def decode_pil(image_bytes: bytes | IO[bytes] | Image.Image) -> Image.Image:
-    if isinstance(image_bytes, Image.Image):
-        return image_bytes
-    image: Image.Image = Image.open(BytesIO(image_bytes) if isinstance(image_bytes, bytes) else image_bytes)
+def decode_pil(image_bytes: bytes | Image.Image | NDArray[np.uint8]) -> Image.Image:
+    image: Image.Image
+    match image_bytes:
+        case Image.Image():
+            image = image_bytes
+        case np.ndarray():
+            image = Image.fromarray(image_bytes)
+        case bytes():
+            image = Image.open(BytesIO(image_bytes))
     image.load()
     if not image.mode == "RGB":
         image = image.convert("RGB")
     return image
 
 
-def decode_cv2(image_bytes: NDArray[np.uint8] | bytes | Image.Image) -> NDArray[np.uint8]:
-    match image_bytes:
-        case bytes() | memoryview() | bytearray():
-            return pil_to_cv2(decode_pil(image_bytes))  # pillow is much faster than cv2
-        case Image.Image():
-            return pil_to_cv2(image_bytes)
-        case _:
-            return image_bytes
-
-
 def clean_text(text: str, canonicalize: bool = False) -> str:
     text = " ".join(text.split())
     if canonicalize:
@@ -78,3 +73,21 @@ def clean_text(text: str, canonicalize: bool = False) -> str:
 # TODO: use this in a less invasive way
 def serialize_np_array(arr: NDArray[np.float32]) -> str:
     return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode()
+
+
+def letterbox(image: NDArray[np.uint8] | Image.Image, size: int) -> tuple[NDArray[np.uint8], float]:
+    if isinstance(image, Image.Image):
+        image = np.asarray(image)
+    height, width = image.shape[:2]
+    if height > width:
+        new_height, new_width = size, int(size * width / height)
+    else:
+        new_width, new_height = size, int(size * height / width)
+
+    canvas = np.zeros((size, size, 3), dtype=np.uint8)
+    cv2.resize(image, (new_width, new_height), dst=canvas[:new_height, :new_width])
+    return canvas, new_height / height
+
+
+def ensure_dims(array: NDArray[np.float32], ndim: int) -> NDArray[np.float32]:
+    return array if array.ndim >= ndim else np.expand_dims(array, axis=tuple(range(ndim - array.ndim)))

machine-learning/immich_ml/schemas.py

@@ -14,10 +14,7 @@ class ORJSONResponse(JSONResponse):
 
 
 class StrEnum(str, Enum):
-    value: str
-
-    def __str__(self) -> str:
-        return self.value
+    __str__ = str.__str__
 
 
 class BoundingBox(TypedDict):
@@ -63,7 +60,7 @@ ModelIdentity = tuple[ModelType, ModelTask]
 
 class SessionNode(Protocol):
     @property
-    def name(self) -> str | None: ...
+    def name(self) -> str: ...
 
     @property
     def shape(self) -> tuple[int, ...]: ...

machine-learning/immich_ml/sessions/ann/__init__.py

@@ -37,11 +37,11 @@ class AnnSession:
 
     def get_inputs(self) -> list[SessionNode]:
         shapes = self.ann.input_shapes[self.model]
-        return [AnnNode(None, s) for s in shapes]
+        return [AnnNode(f"input.{i + 1}", s) for i, s in enumerate(shapes)]
 
     def get_outputs(self) -> list[SessionNode]:
         shapes = self.ann.output_shapes[self.model]
-        return [AnnNode(None, s) for s in shapes]
+        return [AnnNode(f"output.{i + 1}", s) for i, s in enumerate(shapes)]
 
     def run(
         self,
@@ -54,5 +54,5 @@ class AnnSession:
 
 
 class AnnNode(NamedTuple):
-    name: str | None
+    name: str
     shape: tuple[int, ...]

machine-learning/immich_ml/sessions/rknn/__init__.py

@@ -69,7 +69,7 @@ class RknnSession:
 
 
 class RknnNode(NamedTuple):
-    name: str | None
+    name: str
     shape: tuple[int, ...]