add batch dim

machine-learning/app/models/facial_recognition/recognition.py

@@ -2,17 +2,19 @@
 from typing import Any
 
 import numpy as np
+import onnx
+import onnxruntime as ort
 from insightface.model_zoo import ArcFaceONNX
 from insightface.utils.face_align import norm_crop
 from numpy.typing import NDArray
+from onnx.tools.update_model_dims import update_inputs_outputs_dims
 from PIL import Image
 
-from app.config import clean_name
+from app.config import clean_name, log
+from app.models.base import InferenceModel
 from app.models.transforms import decode_cv2
 from app.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelSession, ModelTask, ModelType
 
-from ..base import InferenceModel
-
 
 class FaceRecognizer(InferenceModel):
     depends = [(ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)]
@@ -30,6 +32,9 @@ def __init__(
 
     def _load(self) -> ModelSession:
         session = self._make_session(self.model_path)
+        if not self._has_batch_dim(session):
+            self._add_batch_dim(self.model_path)
+            session = self._make_session(self.model_path)
         self.model = ArcFaceONNX(
             self.model_path.with_suffix(".onnx").as_posix(),
             session=session,
@@ -57,3 +62,16 @@ def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32
 
     def _crop(self, image: NDArray[np.uint8], faces: FaceDetectionOutput) -> list[NDArray[np.uint8]]:
         return [norm_crop(image, landmark) for landmark in faces["landmarks"]]
+
+    def _has_batch_dim(self, session: ort.InferenceSession) -> bool:
+        return not isinstance(session, ort.InferenceSession) or session.get_inputs()[0].shape[0] == "batch"
+
+    def _add_batch_dim(self, model_path: Path) -> None:
+        log.debug(f"Adding batch dimension to model {model_path}")
+        proto = onnx.load(model_path)
+        static_input_dims = [shape.dim_value for shape in proto.graph.input[0].type.tensor_type.shape.dim[1:]]
+        static_output_dims = [shape.dim_value for shape in proto.graph.output[0].type.tensor_type.shape.dim[1:]]
+        input_dims = {proto.graph.input[0].name: ["batch"] + static_input_dims}
+        output_dims = {proto.graph.output[0].name: ["batch"] + static_output_dims}
+        updated_proto = update_inputs_outputs_dims(proto, input_dims, output_dims)
+        onnx.save(updated_proto, model_path)