ocr1.0
This commit is contained in:
@ -4,10 +4,12 @@ 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()
|
||||
@ -16,40 +18,30 @@ class FaceRecognizer:
|
||||
self._load_known_faces()
|
||||
|
||||
def _initialize_insightface(self):
|
||||
"""
|
||||
初始化InsightFace FaceAnalysis应用。
|
||||
默认使用CPU，如果检测到CUDA，会自动使用GPU。
|
||||
"""
|
||||
print("正在初始化InsightFace人脸识别引擎...")
|
||||
"""初始化InsightFace FaceAnalysis应用"""
|
||||
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人脸识别引擎初始化成功。")
|
||||
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("请确保已安装insightface和onnxruntime，并且模型文件已下载或可访问。")
|
||||
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}' 不存在。请创建并放入照片。")
|
||||
print(f"已知人脸目录不存在，已创建: {self.known_faces_dir}")
|
||||
os.makedirs(self.known_faces_dir, exist_ok=True)
|
||||
return
|
||||
|
||||
print(f"正在加载已知人脸特征从: '{self.known_faces_dir}'...")
|
||||
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}")
|
||||
print(f"处理人物: {person_name}")
|
||||
embeddings = []
|
||||
for filename in os.listdir(person_dir):
|
||||
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
|
||||
@ -57,131 +49,91 @@ class FaceRecognizer:
|
||||
try:
|
||||
img = cv2.imread(image_path)
|
||||
if img is None:
|
||||
print(f" 警告: 无法读取图片 '{image_path}'，已跳过。")
|
||||
print(f"无法读取图片: {image_path}，已跳过")
|
||||
continue
|
||||
|
||||
# 查找人脸并提取特征
|
||||
|
||||
faces = self.app.get(img)
|
||||
if faces:
|
||||
# 通常一张照片只有一个人脸，取第一个
|
||||
embeddings.append(faces[0].embedding)
|
||||
print(f" 成功提取 '{filename}' 的人脸特征。")
|
||||
print(f"提取特征成功: {filename}")
|
||||
else:
|
||||
print(f" 警告: 在图片 '{filename}' 中未检测到人脸，已跳过。")
|
||||
print(f"未检测到人脸: {filename}，已跳过")
|
||||
except Exception as e:
|
||||
print(f" 处理图片 '{image_path}' 时发生错误: {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)} 张照片。")
|
||||
print(f"人物 {person_name} 加载完成，共 {len(embeddings)} 张照片")
|
||||
else:
|
||||
print(f" 警告: 人物 '{person_name}' 没有有效的人脸特征，已跳过。")
|
||||
print(f"已知人脸加载完成。共 {len(self.known_faces_names)} 个人物。")
|
||||
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) # 在帧中检测并提取所有人的脸
|
||||
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)
|
||||
# -------------------------------------
|
||||
|
||||
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 # 没有检测到已知人脸
|
||||
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
|
||||
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" # 替换为你的测试图片路径
|
||||
# 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)
|
||||
Reference in New Issue
Block a user