From ee90fd2ca502b9894d829448ea32aa0487035065 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 21 Aug 2025 06:03:49 -0600
Subject: [PATCH] Add support for face recognition via RKNN

---
 frigate/embeddings/onnx/runner.py | 26 ++++----------------------
 frigate/util/rknn_converter.py    | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/frigate/embeddings/onnx/runner.py b/frigate/embeddings/onnx/runner.py
index 3a2acc7f6..9158be9f4 100644
--- a/frigate/embeddings/onnx/runner.py
+++ b/frigate/embeddings/onnx/runner.py
@@ -184,6 +184,8 @@ class RKNNModelRunner:
 
         if "vision" in model_name:
             return ["pixel_values"]
+        elif "arcface" in model_name:
+            return ["data"]
         else:
             # Default fallback - try to infer from model type
             if self.model_type and "jina-clip" in self.model_type:
@@ -199,6 +201,8 @@ class RKNNModelRunner:
         model_name = os.path.basename(self.model_path).lower()
         if "vision" in model_name:
             return 224  # CLIP V1 uses 224x224
+        elif "arcface" in model_name:
+            return 112
         return -1
 
     def run(self, inputs: dict[str, Any]) -> Any:
@@ -222,28 +226,6 @@ class RKNNModelRunner:
                         rknn_inputs.append(pixel_data)
                     else:
                         rknn_inputs.append(inputs[name])
-                else:
-                    logger.warning(f"Input '{name}' not found in inputs, using default")
-
-                    if name == "pixel_values":
-                        batch_size = 1
-                        if inputs:
-                            for val in inputs.values():
-                                if hasattr(val, "shape") and len(val.shape) > 0:
-                                    batch_size = val.shape[0]
-                                    break
-                        # Create default in NHWC format as expected by RKNN
-                        rknn_inputs.append(
-                            np.zeros((batch_size, 224, 224, 3), dtype=np.float32)
-                        )
-                    else:
-                        batch_size = 1
-                        if inputs:
-                            for val in inputs.values():
-                                if hasattr(val, "shape") and len(val.shape) > 0:
-                                    batch_size = val.shape[0]
-                                    break
-                        rknn_inputs.append(np.zeros((batch_size, 1), dtype=np.float32))
 
             outputs = self.rknn.inference(inputs=rknn_inputs)
             return outputs
diff --git a/frigate/util/rknn_converter.py b/frigate/util/rknn_converter.py
index e42547320..6c8a0d1fc 100644
--- a/frigate/util/rknn_converter.py
+++ b/frigate/util/rknn_converter.py
@@ -32,6 +32,11 @@ MODEL_TYPE_CONFIGS = {
         "std_values": [[0.26862954 * 255, 0.26130258 * 255, 0.27577711 * 255]],
         "target_platform": None,  # Will be set dynamically
     },
+    "arcface-r100": {
+        "mean_values": [[127.5,127.5,127.5]],
+        "std_values": [[127.5,127.5,127.5]],
+        "target_platform": None,  # Will be set dynamically
+    }
 }
 
 
@@ -41,6 +46,9 @@ def get_rknn_model_type(model_path: str) -> str | None:
 
     model_name = os.path.basename(str(model_path)).lower()
 
+    if "arcface" in model_name:
+        return "arcface-r100"
+
     if any(keyword in model_name for keyword in ["yolo", "yolox", "yolonas"]):
         return model_name
 
@@ -184,6 +192,12 @@ def convert_onnx_to_rknn(
                 inputs=["pixel_values"],
                 input_size_list=[[1, 3, 224, 224]],
             )
+        elif model_type == "arcface-r100":
+            load_output = rknn.load_onnx(
+                model=onnx_path,
+                inputs=["data"],
+                input_size_list=[[1, 3, 112, 112]],
+            )
         else:
             load_output = rknn.load_onnx(model=onnx_path)