139 lines
5.5 KiB
Python
139 lines
5.5 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应用"""
|
||
print("初始化InsightFace引擎...")
|
||
try:
|
||
app = FaceAnalysis(name='buffalo_l', root='~/.insightface')
|
||
app.prepare(ctx_id=0, det_size=(640, 640))
|
||
print("InsightFace引擎初始化完成")
|
||
return app
|
||
except Exception as e:
|
||
print(f"InsightFace初始化失败: {e}")
|
||
print("请检查依赖是否安装及模型是否可访问")
|
||
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):
|
||
"""识别人脸并返回结果"""
|
||
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]
|
||
|
||
embedding1 = face.embedding.astype(np.float32)
|
||
embedding2 = known_embedding.astype(np.float32)
|
||
|
||
dot_product = np.dot(embedding1, embedding2)
|
||
norm_embedding1 = np.linalg.norm(embedding1)
|
||
norm_embedding2 = np.linalg.norm(embedding2)
|
||
|
||
similarity = 0.0 if (norm_embedding1 == 0 or norm_embedding2 == 0) else (
|
||
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):
|
||
"""测试单张图片识别"""
|
||
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)
|
||
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('识别结果', 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) |