forked from Cutlery/immich
66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from insightface.model_zoo import ArcFaceONNX
|
|
from insightface.utils.face_align import norm_crop
|
|
from numpy.typing import NDArray
|
|
|
|
from app.config import clean_name
|
|
from app.models.transforms import crop_np, crop_bounding_box, resize_np
|
|
from app.schemas import DetectedFace, ModelTask, RecognizedFace, ModelSession, ModelType, is_ndarray
|
|
|
|
from ..base import InferenceModel
|
|
|
|
|
|
class FaceRecognizer(InferenceModel):
|
|
_model_task = ModelTask.FACIAL_RECOGNITION
|
|
_model_type = ModelType.RECOGNITION
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
min_score: float = 0.7,
|
|
cache_dir: Path | str | None = None,
|
|
**model_kwargs: Any,
|
|
) -> None:
|
|
self.min_score = model_kwargs.pop("minScore", min_score)
|
|
super().__init__(clean_name(model_name), cache_dir, **model_kwargs)
|
|
|
|
def _load(self) -> ModelSession:
|
|
session = self._make_session(self.model_path)
|
|
self.model = ArcFaceONNX(
|
|
self.model_path.with_suffix(".onnx").as_posix(),
|
|
session=session,
|
|
)
|
|
return session
|
|
|
|
# def _predict(self, img: Any, **kwargs: Any) -> Any:
|
|
def _predict(
|
|
self, inputs: NDArray[np.uint8] | bytes, faces: list[DetectedFace] = [], **kwargs: Any
|
|
) -> list[RecognizedFace]:
|
|
if isinstance(inputs, bytes):
|
|
decoded_image = cv2.imdecode(np.frombuffer(inputs, np.uint8), cv2.IMREAD_COLOR)
|
|
else:
|
|
decoded_image = inputs
|
|
assert is_ndarray(decoded_image, np.float32)
|
|
|
|
results: list[RecognizedFace] = []
|
|
for detected_face in faces:
|
|
landmarks = detected_face.get("landmarks", None)
|
|
if landmarks is not None:
|
|
cropped_img = norm_crop(decoded_image, np.asarray(landmarks))
|
|
else:
|
|
cropped_img = crop_bounding_box(decoded_image, detected_face["box"])
|
|
cropped_img = crop_np(resize_np(cropped_img, 112), 112)
|
|
assert is_ndarray(cropped_img, np.uint8)
|
|
|
|
embedding = self.model.get_feat(cropped_img)[0]
|
|
assert is_ndarray(embedding, np.float32)
|
|
|
|
face: RecognizedFace = {"box": detected_face["box"], "embedding": embedding}
|
|
results.append(face)
|
|
|
|
return results
|