public/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2025-12-05 12:23:58 +09:00
Code Issues Actions 8 Packages Projects Releases Wiki 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
2 changed files with 1 additions and 7 deletions
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(