feat(ml): add face models (#4952)

added models to config dropdown

fixed downloading

updated tests

use hf for face models

formatting

machine-learning/app/models/facial_recognition.py

@@ -1,4 +1,3 @@
-import zipfile
 from pathlib import Path
 from typing import Any
 
@@ -7,8 +6,8 @@ import numpy as np
 import onnxruntime as ort
 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 app.config import clean_name
 from app.schemas import ModelType, ndarray_f32
 
 from .base import InferenceModel
@@ -25,37 +24,21 @@ class FaceRecognizer(InferenceModel):
         **model_kwargs: Any,
     ) -> None:
         self.min_score = model_kwargs.pop("minScore", min_score)
-        super().__init__(model_name, cache_dir, **model_kwargs)
-
-    def _download(self) -> 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()
+        super().__init__(clean_name(model_name), cache_dir, **model_kwargs)
 
     def _load(self) -> 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(
             session=ort.InferenceSession(
-                det_file.as_posix(),
+                self.det_file.as_posix(),
                 sess_options=self.sess_options,
                 providers=self.providers,
                 provider_options=self.provider_options,
             ),
         )
         self.rec_model = ArcFaceONNX(
-            rec_file.as_posix(),
+            self.rec_file.as_posix(),
             session=ort.InferenceSession(
-                rec_file.as_posix(),
+                self.rec_file.as_posix(),
                 sess_options=self.sess_options,
                 providers=self.providers,
                 provider_options=self.provider_options,
@@ -103,7 +86,15 @@ class FaceRecognizer(InferenceModel):
 
     @property
     def cached(self) -> bool:
-        return self.cache_dir.is_dir() and any(self.cache_dir.glob("*.onnx"))
+        return self.det_file.is_file() and self.rec_file.is_file()
+
+    @property
+    def det_file(self) -> Path:
+        return self.cache_dir / "detection" / "model.onnx"
+
+    @property
+    def rec_file(self) -> Path:
+        return self.cache_dir / "recognition" / "model.onnx"
 
     def configure(self, **model_kwargs: Any) -> None:
         self.det_model.det_thresh = model_kwargs.pop("minScore", self.det_model.det_thresh)