self-hosted/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2024-12-27 10:58:13 +02:00
Code Issues Releases Activity

fix(ml): batch axis not being added for recognition model (#12588)

Browse Source 
* fix has_batch_axis

* fix typing
This commit is contained in:
Mert 2024-09-11 23:21:12 -04:00 committed by GitHub
parent fa095c3ca0
commit 22dc9bcebb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 1 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
machine-learning/app/models/facial_recognition/recognition.py
View File

@ -13,7 +13,6 @@ from app.config import log
from app.models.base import InferenceModel
from app.models.transforms import decode_cv2
from app.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelFormat, ModelSession, ModelTask, ModelType
from app.sessions import has_batch_axis
class FaceRecognizer(InferenceModel):
@ -27,7 +26,7 @@ class FaceRecognizer(InferenceModel):
def _load(self) -> ModelSession:
session = self._make_session(self.model_path)
if self.batch and not has_batch_axis(session):
if self.batch 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(

5
machine-learning/app/sessions/__init__.py
View File

@ -1,5 +0,0 @@
from app.schemas import ModelSession
def has_batch_axis(session: ModelSession) -> bool:
return not isinstance(session.get_inputs()[0].shape[0], int) or session.get_inputs()[0].shape[0] < 0