mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-04-16 03:52:09 +03:00
Add support for face recognition via RKNN
This commit is contained in:
parent
1be84d6833
commit
ee90fd2ca5
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user