insightface_pipeline.py

import cv2
import numpy as np
from insightface.app import FaceAnalysis
import argparse
import os
from tqdm import tqdm

class FaceRecognitionPipeline:
    def __init__(self, model_path='./models/axmodel'):
        """
        初始化人脸识别管道  
        :param providers: ONNX 推理后端，支持 GPU/CPU
        """
        # 1. 初始化 FaceAnalysis（包含 detection + landmark）
        self.app = FaceAnalysis(root=model_path)
        self.app.prepare(ctx_id=0, det_size=(640, 640))  # det_size 可调

        # 2. （可选）显式加载 ArcFace 模型（其实 FaceAnalysis 已包含）
        # self.rec_model = get_model('antelopev2', providers=providers)  # 或 'buffalo_l'
        # self.rec_model.prepare(ctx_id=0)

    def extract_features(self, image):
        """
        从图像中提取所有人脸的特征向量
        :param image: BGR 图像 (H, W, 3)
        :return: faces_info = [{'bbox', 'kps', 'embedding'}, ...]
        """
        faces = self.app.get(image)
        return faces  # 每个 face 是 dict，包含 embedding (512-d)

    def compare(self, emb1, emb2, threshold=0.25):
        """
        计算两个 512 维特征向量的余弦相似度（InsightFace 使用余弦距离）
        :param emb1, emb2: shape=(512,)
        :param threshold: 相似度阈值（antelopev2 推荐 0.35，buffalo_l 推荐 0.25）
        :return: (similarity, is_same)
        """
        similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
        is_same = similarity > threshold
        return similarity, is_same

    def recognize(self, query_image, gallery_embeddings, gallery_names, threshold=0.25):
        """
        在图库中识别查询图像中的人脸
        :param query_image: BGR 图像
        :param gallery_embeddings: list of embeddings (n x 512)
        :param gallery_names: list of names (n,)
        :param threshold: 识别阈值
        :return: [{'name', 'similarity', 'bbox'}, ...]
        """
        query_faces = self.extract_features(query_image)
        results = []
        for face in query_faces:
            best_sim = -1
            best_name = "Unknown"
            for emb, name in zip(gallery_embeddings, gallery_names):
                sim, _ = self.compare(face['embedding'], emb, threshold=0.25)  # 先算相似度
                if sim > best_sim:
                    best_sim = sim
                    best_name = name if sim > threshold else "Unknown"
            results.append({
                'name': best_name,
                'similarity': best_sim,
                'face': face
            })
        return results

    def draw_results(self, image, results):
        """在图像上绘制识别结果"""
        img_draw = image.copy()
        for res in results:
            img_draw = self.app.draw_on(img_draw, [res['face']])

            x1, y1, x2, y2 = res['face']['bbox'].astype(int)
            color = (0, 255, 0) if res['name'] != "Unknown" else (0, 0, 255)
            cv2.putText(img_draw, f"{res['name']}: {res['similarity']:.2f}",
                        (x1, y2 + 15), cv2.FONT_HERSHEY_COMPLEX, 0.7, color, 1)
        return img_draw

# ==================== 使用示例 ====================

if __name__ == "__main__":
    args = argparse.ArgumentParser(description="Face Recognition Pipeline Example")
    args.add_argument("--model_path", "-m", type=str, default="./models/buffalo_l", help="Path to the model directory")
    args.add_argument("--type", "-t", type=int, default=0, help="Type of operation: 1: 1v1 compare, 2: 1vN recognize")
    args.add_argument("--gallery_path", "-g", type=str, default=None, help="Path to the gallery image for image file")
    args.add_argument("--query_path", "-q", type=str, default=None, help="Path to the query image")
    args.add_argument("--draw", "-d", action='store_true', help="Whether to draw results on the image")
    args = args.parse_args()

    # 初始化 Pipeline
    pipeline = FaceRecognitionPipeline(model_path=args.model_path)

    # 示例1: 1:1 比对
    if args.type == 0:
        assert args.gallery_path is not None and args.query_path is not None, "请提供 gallery_path 和 query_path"
        gallery_img = cv2.imread(args.gallery_path)
        faces1 = pipeline.extract_features(gallery_img)
        if faces1:
            emb1 = faces1[0]['embedding']
        else:
            print(f"警告: {args.gallery_path} 未检测到人脸")
            exit(0)

        query_img = cv2.imread(args.query_path)
        faces2 = pipeline.extract_features(query_img)
        if faces1 and faces2:
            sim, is_same = pipeline.compare(emb1, faces2[0]['embedding'], threshold=0.25)
            print(f"相似度: {sim:.4f}, 是否同一人: {is_same}")
        else:
            print(f"警告: {args.query_path} 未检测到人脸")
            exit(0)

        if args.draw:
            os.makedirs("./output", exist_ok=True)
            output_img = pipeline.app.draw_on(query_img, [faces2[0]])
            cv2.imwrite(f"./output/{os.path.basename(args.query_path)}", output_img)
            print(f"结果已保存到 ./output/{os.path.basename(args.query_path)}")

    elif args.type == 1:

        assert args.gallery_path is not None, "请提供 gallery_path"

        # 构建特征库
        gallery_names = []
        gallery_embeddings = []
        for fname in tqdm(os.listdir(args.gallery_path)):
            name = os.path.splitext(os.path.basename(fname))[0]
            gallery_img = cv2.imread(os.path.join(args.gallery_path, fname))
            faces = pipeline.extract_features(gallery_img)
            if faces:
                gallery_names.append(name)
                gallery_embeddings.append(faces[0]['embedding'])
            else:
                print(f"警告: {fname} 未检测到人脸")

        # 识别查询图像
        print("特征库构建完成，包含以下人员：", gallery_names)
        if args.draw:
            os.makedirs("./output", exist_ok=True)
        while True:
            print("请输入查询图像路径 (输入 'exit' 退出): ")
            user_input = input()
            if user_input.lower() == 'exit':
                break
            if not os.path.isfile(user_input):
                print("输入的路径不是有效的文件，请重新输入。")
                continue
            query_img = cv2.imread(user_input)
            results = pipeline.recognize(query_img, gallery_embeddings, gallery_names, threshold=0.25)
            
            if results is None or len(results) == 0:
                print(f"{user_input} 未检测到人脸")
                continue
            for res in results:
                print(f"识别结果: {res['name']}, 相似度(0-1): {res['similarity']:.4f}")
            if args.draw:
                output_img = pipeline.draw_results(query_img, results)
                cv2.imwrite(f"./output/{os.path.basename(user_input)}", output_img)
                print(f"结果已保存到 ./output/{os.path.basename(user_input)}")