diff --git a/.gitignore b/.gitignore index a29ffee..7584f2f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1 @@ /faces -/models \ No newline at end of file diff --git a/debug_output/003_aligned.png b/debug_output/003_aligned.png new file mode 100644 index 0000000..d73c3f3 Binary files /dev/null and b/debug_output/003_aligned.png differ diff --git a/debug_output/003_aligned_rgb.npy b/debug_output/003_aligned_rgb.npy new file mode 100644 index 0000000..2938d73 Binary files /dev/null and b/debug_output/003_aligned_rgb.npy differ diff --git a/debug_output/003_emb.npy b/debug_output/003_emb.npy new file mode 100644 index 0000000..dcc611b Binary files /dev/null and b/debug_output/003_emb.npy differ diff --git a/models/RetinaFace_mobile320.onnx b/models/RetinaFace_mobile320.onnx new file mode 100644 index 0000000..0cab38e Binary files /dev/null and b/models/RetinaFace_mobile320.onnx differ diff --git a/models/mobilefacenet_arcface_prenorm.onnx b/models/mobilefacenet_arcface_prenorm.onnx new file mode 100644 index 0000000..123658e Binary files /dev/null and b/models/mobilefacenet_arcface_prenorm.onnx differ diff --git a/models/retinaface_mobile320_config.json b/models/retinaface_mobile320_config.json new file mode 100644 index 0000000..5c48142 --- /dev/null +++ b/models/retinaface_mobile320_config.json @@ -0,0 +1,22 @@ +{ + "input": { + "color": "RGB", + "dtype": "float32", + "layout": "NCHW", + "resize": { "mode": "stretch", "size": [320, 320] } + }, + "decoder": { + "type": "retinaface", + "steps": [8, 16, 32], + "min_sizes": [[16, 32], [64, 128], [256, 512]], + "variances": [0.1, 0.2], + "score_index": 1, + "conf_mode": "auto", + "nms_iou_thresh": 0.4 + }, + "outputs": { + "loc": { "name": "output0" }, + "conf": { "name": "572" }, + "landmarks": { "name": "571" } + } +} \ No newline at end of file diff --git a/testpic.py b/testpic.py new file mode 100644 index 0000000..06bb5fd --- /dev/null +++ b/testpic.py @@ -0,0 +1,48 @@ +import os +import cv2 +import numpy as np +import onnxruntime as ort +from gallery_builder.align import align_face_5pts +from gallery_builder.detector import OnnxFaceDetector, load_det_outputs_config + +# ============ 只改这里 ============ +INPUT_IMAGE = "003.jpg" +# ================================= + +DET_MODEL = "./models/RetinaFace_mobile320.onnx" +DET_CONFIG = "./models/retinaface_mobile320_config.json" +RECOG_MODEL = "./models/mobilefacenet_arcface_prenorm.onnx" +OUT_DIR = "./debug_output" + +os.makedirs(OUT_DIR, exist_ok=True) +base_name = os.path.splitext(os.path.basename(INPUT_IMAGE))[0] + +det_cfg = load_det_outputs_config(DET_CONFIG) +detector = OnnxFaceDetector(DET_MODEL, det_cfg, score_thresh=0.5, pick_face="largest") +recog_session = ort.InferenceSession(RECOG_MODEL, providers=["CPUExecutionProvider"]) + +img_bgr = cv2.imread(INPUT_IMAGE) +det = detector.detect_one(img_bgr) +print(f"landmarks5:\n{det.landmarks5}") + +img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) +aligned = align_face_5pts(img_rgb, det.landmarks5) + +# 保存对齐图 +cv2.imwrite(f"{OUT_DIR}/{base_name}_aligned.png", cv2.cvtColor(aligned, cv2.COLOR_RGB2BGR)) +np.save(f"{OUT_DIR}/{base_name}_aligned_rgb.npy", aligned) + +# 提取 embedding +x = (aligned.astype(np.float32) - 127.5) / 128.0 +x = x.transpose(2, 0, 1)[np.newaxis, ...] +out = recog_session.run(None, {recog_session.get_inputs()[0].name: x})[0] +emb = out.flatten() +emb = emb / np.linalg.norm(emb) + +np.save(f"{OUT_DIR}/{base_name}_emb.npy", emb) + +print(f"\n保存到 {OUT_DIR}/") +print(f" {base_name}_aligned.png") +print(f" {base_name}_aligned_rgb.npy") +print(f" {base_name}_emb.npy") +print(f"\nembedding[0:8]: {emb[:8]}") \ No newline at end of file