KevinMidboe/immich
1
0
Fork 0
mirror of https://github.com/KevinMidboe/immich.git synced 2025-12-08 20:29:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(ml): model downloading (#3545)

Browse Source 
* download facial recognition models

* download hf models

* simplified logic

* updated `predict` for facial recognition

* ensure download method is called

* fixed repo_id for clip

* fixed download destination

* use st's own `snapshot_download`

* conditional download

* fixed predict method

* check if loaded

* minor fixes

* updated mypy overrides

* added pytest-mock

* updated tests

* updated lock
This commit is contained in:
Mert
2023-08-05 22:45:13 -04:00
committed by GitHub
parent 2f26a7edae
commit c73832bd9c
10 changed files with 350 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
machine-learning/app/models/facial_recognition.py
View File

@@ -1,8 +1,12 @@
import zipfile
from pathlib import Path
from typing import Any
import cv2
from insightface.app import FaceAnalysis
import numpy as np
from insightface.model_zoo import ArcFaceONNX, RetinaFace
from insightface.utils.face_align import norm_crop
from insightface.utils.storage import BASE_REPO_URL, download_file
from ..config import settings
from ..schemas import ModelType
@@ -22,39 +26,62 @@ class FaceRecognizer(InferenceModel):
self.min_score = min_score
super().__init__(model_name, cache_dir, **model_kwargs)
def load(self, **model_kwargs: Any) -> None:
self.model = FaceAnalysis(
name=self.model_name,
root=self.cache_dir.as_posix(),
allowed_modules=["detection", "recognition"],
**model_kwargs,
)
self.model.prepare(
ctx_id=0,
def _download(self, **model_kwargs: Any) -> None:
zip_file = self.cache_dir / f"{self.model_name}.zip"
download_file(f"{BASE_REPO_URL}/{self.model_name}.zip", zip_file)
with zipfile.ZipFile(zip_file, "r") as zip:
members = zip.namelist()
det_file = next(model for model in members if model.startswith("det_"))
rec_file = next(model for model in members if model.startswith("w600k_"))
zip.extractall(self.cache_dir, members=[det_file, rec_file])
zip_file.unlink()
def _load(self, **model_kwargs: Any) -> None:
try:
det_file = next(self.cache_dir.glob("det_*.onnx"))
rec_file = next(self.cache_dir.glob("w600k_*.onnx"))
except StopIteration:
raise FileNotFoundError("Facial recognition models not found in cache directory")
self.det_model = RetinaFace(det_file.as_posix())
self.rec_model = ArcFaceONNX(rec_file.as_posix())
self.det_model.prepare(
ctx_id=-1,
det_thresh=self.min_score,
det_size=(640, 640),
input_size=(640, 640),
)
self.rec_model.prepare(ctx_id=-1)
def _predict(self, image: cv2.Mat) -> list[dict[str, Any]]:
bboxes, kpss = self.det_model.detect(image)
if bboxes.size == 0:
return []
assert isinstance(kpss, np.ndarray)
scores = bboxes[:, 4].tolist()
bboxes = bboxes[:, :4].round().tolist()
def predict(self, image: cv2.Mat) -> list[dict[str, Any]]:
height, width, _ = image.shape
results = []
faces = self.model.get(image)
for face in faces:
x1, y1, x2, y2 = face.bbox
height, width, _ = image.shape
for (x1, y1, x2, y2), score, kps in zip(bboxes, scores, kpss):
cropped_img = norm_crop(image, kps)
embedding = self.rec_model.get_feat(cropped_img)[0].tolist()
results.append(
{
"imageWidth": width,
"imageHeight": height,
"boundingBox": {
"x1": round(x1),
"y1": round(y1),
"x2": round(x2),
"y2": round(y2),
"x1": x1,
"y1": y1,
"x2": x2,
"y2": y2,
},
"score": face.det_score.item(),
"embedding": face.normed_embedding.tolist(),
"score": score,
"embedding": embedding,
}
)
return results
@property
def cached(self) -> bool:
return self.cache_dir.is_dir() and any(self.cache_dir.glob("*.onnx"))