Add support for face recognition via RKNN

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

View File

@ -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

View File

@ -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)