administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-04-23 07:14:35 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Add support for face recognition via RKNN

Browse Source
This commit is contained in:
Nicolas Mowen 2025-08-21 06:03:49 -06:00
parent 1be84d6833
commit ee90fd2ca5
2 changed files with 18 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
frigate/embeddings/onnx/runner.py
View File

@ -184,6 +184,8 @@ class RKNNModelRunner:
if "vision" in model_name: if "vision" in model_name:
return ["pixel_values"] return ["pixel_values"]
elif "arcface" in model_name:
return ["data"]
else: else:
# Default fallback - try to infer from model type # Default fallback - try to infer from model type
if self.model_type and "jina-clip" in self.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() model_name = os.path.basename(self.model_path).lower()
if "vision" in model_name: if "vision" in model_name:
return 224 # CLIP V1 uses 224x224 return 224 # CLIP V1 uses 224x224
elif "arcface" in model_name:
return 112
return -1 return -1
def run(self, inputs: dict[str, Any]) -> Any: def run(self, inputs: dict[str, Any]) -> Any:
@ -222,28 +226,6 @@ class RKNNModelRunner:
rknn_inputs.append(pixel_data) rknn_inputs.append(pixel_data)
else: else:
rknn_inputs.append(inputs[name]) 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) outputs = self.rknn.inference(inputs=rknn_inputs)
return outputs return outputs

14
frigate/util/rknn_converter.py
View File

@ -32,6 +32,11 @@ MODEL_TYPE_CONFIGS = {
"std_values": [[0.26862954 * 255, 0.26130258 * 255, 0.27577711 * 255]], "std_values": [[0.26862954 * 255, 0.26130258 * 255, 0.27577711 * 255]],
"target_platform": None, # Will be set dynamically "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() 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"]): if any(keyword in model_name for keyword in ["yolo", "yolox", "yolonas"]):
return model_name return model_name
@ -184,6 +192,12 @@ def convert_onnx_to_rknn(
inputs=["pixel_values"], inputs=["pixel_values"],
input_size_list=[[1, 3, 224, 224]], 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: else:
load_output = rknn.load_onnx(model=onnx_path) load_output = rknn.load_onnx(model=onnx_path)