187 lines
8.3 KiB
Python
187 lines
8.3 KiB
Python
import os
|
||
import cv2
|
||
import numpy as np
|
||
import insightface
|
||
from insightface.app import FaceAnalysis
|
||
|
||
class FaceRecognizer:
|
||
"""
|
||
封装InsightFace人脸识别功能，支持从文件夹加载已知人脸。
|
||
"""
|
||
def __init__(self, known_faces_dir: str):
|
||
self.known_faces_dir = known_faces_dir
|
||
self.app = self._initialize_insightface()
|
||
self.known_faces_embeddings = {}
|
||
self.known_faces_names = []
|
||
self._load_known_faces()
|
||
|
||
def _initialize_insightface(self):
|
||
"""
|
||
初始化InsightFace FaceAnalysis应用。
|
||
默认使用CPU，如果检测到CUDA，会自动使用GPU。
|
||
"""
|
||
print("正在初始化InsightFace人脸识别引擎...")
|
||
try:
|
||
# 默认模型是 'buffalo_l'，包含检测、对齐、识别功能
|
||
# 如果需要更小的模型，可以尝试 'buffalo_s' 或 'buffalo_m'
|
||
# ctx_id=0 表示使用GPU，ctx_id=-1 表示使用CPU
|
||
# InsightFace会自动检测CUDA并选择GPU，所以通常不需要手动设置ctx_id
|
||
app = FaceAnalysis(name='buffalo_l', root='~/.insightface') # 模型下载到用户目录
|
||
app.prepare(ctx_id=0, det_size=(640, 640)) # det_size影响检测性能和精度
|
||
print("InsightFace人脸识别引擎初始化成功。")
|
||
return app
|
||
except Exception as e:
|
||
print(f"InsightFace人脸识别引擎初始化失败: {e}")
|
||
print("请确保已安装insightface和onnxruntime，并且模型文件已下载或可访问。")
|
||
return None
|
||
|
||
def _load_known_faces(self):
|
||
"""
|
||
扫描已知人脸目录，加载每个人的照片并计算人脸特征。
|
||
"""
|
||
if not os.path.exists(self.known_faces_dir):
|
||
print(f"警告: 已知人脸目录 '{self.known_faces_dir}' 不存在。请创建并放入照片。")
|
||
os.makedirs(self.known_faces_dir, exist_ok=True)
|
||
return
|
||
|
||
print(f"正在加载已知人脸特征从: '{self.known_faces_dir}'...")
|
||
for person_name in os.listdir(self.known_faces_dir):
|
||
|
||
person_dir = os.path.join(self.known_faces_dir, person_name)
|
||
if os.path.isdir(person_dir):
|
||
print(f" 加载人物: {person_name}")
|
||
embeddings = []
|
||
for filename in os.listdir(person_dir):
|
||
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
|
||
image_path = os.path.join(person_dir, filename)
|
||
try:
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
print(f" 警告: 无法读取图片 '{image_path}'，已跳过。")
|
||
continue
|
||
|
||
# 查找人脸并提取特征
|
||
faces = self.app.get(img)
|
||
if faces:
|
||
# 通常一张照片只有一个人脸，取第一个
|
||
embeddings.append(faces[0].embedding)
|
||
print(f" 成功提取 '{filename}' 的人脸特征。")
|
||
else:
|
||
print(f" 警告: 在图片 '{filename}' 中未检测到人脸，已跳过。")
|
||
except Exception as e:
|
||
print(f" 处理图片 '{image_path}' 时发生错误: {e}")
|
||
|
||
if embeddings:
|
||
# 将多张照片的特征取平均，作为该人物的最终特征
|
||
self.known_faces_embeddings[person_name] = np.array(embeddings).mean(axis=0)
|
||
self.known_faces_names.append(person_name)
|
||
print(f" 人物 '{person_name}' 加载完成，共 {len(embeddings)} 张照片。")
|
||
else:
|
||
print(f" 警告: 人物 '{person_name}' 没有有效的人脸特征，已跳过。")
|
||
print(f"已知人脸加载完成。共 {len(self.known_faces_names)} 个人物。")
|
||
|
||
def recognize(self, frame, threshold=0.4):
|
||
"""
|
||
在视频帧中识别人脸。
|
||
|
||
Args:
|
||
frame: 输入的图像帧 (NumPy数组, BGR格式)。
|
||
threshold (float): 识别相似度阈值。0.0到1.0，越高越严格。
|
||
|
||
Returns:
|
||
tuple[bool, str|None, float|None]: 是否检测到已知人脸，检测到的人名，以及相似度。
|
||
"""
|
||
if not self.app or not self.known_faces_names:
|
||
return False, None, None
|
||
|
||
faces = self.app.get(frame) # 在帧中检测并提取所有人的脸
|
||
if not faces:
|
||
return False, None, None
|
||
|
||
for face in faces:
|
||
# 遍历已知人脸，进行比对
|
||
for known_name in self.known_faces_names:
|
||
known_embedding = self.known_faces_embeddings[known_name]
|
||
|
||
# --- 关键修改：手动计算余弦相似度 ---
|
||
# 确保embedding是float32类型，避免潜在的类型不匹配问题
|
||
embedding1 = face.embedding.astype(np.float32)
|
||
embedding2 = known_embedding.astype(np.float32)
|
||
|
||
# 计算点积
|
||
dot_product = np.dot(embedding1, embedding2)
|
||
# 计算L2范数（向量长度）
|
||
norm_embedding1 = np.linalg.norm(embedding1)
|
||
norm_embedding2 = np.linalg.norm(embedding2)
|
||
|
||
# 避免除以零
|
||
if norm_embedding1 == 0 or norm_embedding2 == 0:
|
||
similarity = 0.0
|
||
else:
|
||
similarity = dot_product / (norm_embedding1 * norm_embedding2)
|
||
# -------------------------------------
|
||
|
||
if similarity >= threshold:
|
||
print(f"!!! 人脸识别检测到已知人物: '{known_name}' (相似度: {similarity:.4f}) !!!")
|
||
return True, known_name, similarity # 只要检测到一个就立即返回
|
||
|
||
return False, None, None # 没有检测到已知人脸
|
||
|
||
|
||
# def test_single_image(self, image_path: str, threshold=0.4):
|
||
# """
|
||
# 测试单张图片的人脸识别效果
|
||
#
|
||
# Args:
|
||
# image_path: 图片路径
|
||
# threshold: 识别阈值
|
||
#
|
||
# Returns:
|
||
# tuple[bool, str|None, float|None]: 是否检测到已知人脸，检测到的人名，以及相似度
|
||
# """
|
||
# if not os.path.exists(image_path):
|
||
# print(f"错误: 图片 '{image_path}' 不存在")
|
||
# return False, None, None
|
||
#
|
||
# # 读取图片
|
||
# frame = cv2.imread(image_path)
|
||
# if frame is None:
|
||
# print(f"错误: 无法读取图片 '{image_path}'")
|
||
# return False, None, None
|
||
#
|
||
# # 调用识别方法
|
||
# result, name, similarity = self.recognize(frame, threshold)
|
||
#
|
||
# # 显示结果
|
||
# if result:
|
||
# print(f"测试结果: 在图片中识别到 {name}，相似度: {similarity:.4f}")
|
||
#
|
||
# # 绘制识别结果并显示图片
|
||
# faces = self.app.get(frame)
|
||
# for face in faces:
|
||
# bbox = face.bbox.astype(int)
|
||
# # 绘制 bounding box
|
||
# cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
|
||
# # 绘制姓名和相似度
|
||
# text = f"{name}: {similarity:.2f}"
|
||
# cv2.putText(frame, text, (bbox[0], bbox[1] - 10),
|
||
# cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
|
||
#
|
||
# # 显示图片
|
||
# cv2.imshow('Recognition Result', frame)
|
||
# print("按任意键关闭图片窗口...")
|
||
# cv2.waitKey(0)
|
||
# cv2.destroyAllWindows()
|
||
# else:
|
||
# print("测试结果: 未在图片中识别到已知人脸")
|
||
#
|
||
# return result, name, similarity
|
||
|
||
|
||
# if __name__ == "__main__":
|
||
# # 初始化人脸识别器，指定已知人脸目录
|
||
# recognizer = FaceRecognizer(known_faces_dir="known_faces")
|
||
#
|
||
# # 测试单张图片
|
||
# test_image_path = r"F:\OCR\RapidOCR-main\known_faces\B\14sino-qiu02-master1050.jpg" # 替换为你的测试图片路径
|
||
# recognizer.test_single_image(test_image_path, threshold=0.4) |