119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
import cv2
|
|
import time
|
|
import argparse
|
|
import os
|
|
from face_utils import FaceDetector, FaceRecognizer, draw_face_box
|
|
from face_db import FaceDatabase
|
|
from voice_prompt import VoicePrompt
|
|
|
|
def parse_args():
|
|
"""解析命令行参数"""
|
|
parser = argparse.ArgumentParser(description='人脸注册程序')
|
|
parser.add_argument('--name', type=str, help='要注册的人员姓名')
|
|
parser.add_argument('--camera', type=int, default=0, help='摄像头索引')
|
|
parser.add_argument('--threshold', type=float, default=0.5, help='人脸检测置信度阈值')
|
|
return parser.parse_args()
|
|
|
|
def register_face():
|
|
"""人脸注册主函数"""
|
|
args = parse_args()
|
|
|
|
# 初始化组件
|
|
voice = VoicePrompt()
|
|
detector = FaceDetector(conf_thres=args.threshold)
|
|
recognizer = FaceRecognizer(detector=detector)
|
|
db = FaceDatabase()
|
|
|
|
# 打开摄像头
|
|
cap = cv2.VideoCapture(args.camera)
|
|
if not cap.isOpened():
|
|
print("无法打开摄像头")
|
|
return
|
|
|
|
# 获取姓名（如果命令行没有提供，则交互输入）
|
|
name = args.name
|
|
if name is None:
|
|
name = input("请输入姓名: ")
|
|
|
|
# 创建新的人员记录
|
|
person_id = db.add_person(name)
|
|
print(f"已创建人员记录: {name} (ID: {person_id})")
|
|
voice.speak(f"开始为{name}注册人脸信息，请正视摄像头", block=True)
|
|
|
|
# 采集计数器和状态
|
|
face_count = 0
|
|
target_count = 5 # 需要采集的人脸图像数量
|
|
last_capture_time = 0
|
|
capture_interval = 1.0 # 每次采集间隔秒数
|
|
|
|
print(f"请看向摄像头，将采集{target_count}张人脸图像...")
|
|
|
|
while face_count < target_count:
|
|
# 读取摄像头帧
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("无法获取摄像头画面")
|
|
break
|
|
|
|
# 镜像翻转以便更直观
|
|
frame = cv2.flip(frame, 1)
|
|
|
|
# 显示帧
|
|
display_frame = frame.copy()
|
|
cv2.putText(
|
|
display_frame,
|
|
f"请直视摄像头 ({face_count}/{target_count})",
|
|
(10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
(0, 255, 0),
|
|
2
|
|
)
|
|
|
|
# 检测人脸
|
|
faces = detector.detect_faces(frame)
|
|
|
|
# 如果检测到人脸
|
|
current_time = time.time()
|
|
if faces and (current_time - last_capture_time) >= capture_interval:
|
|
# 只处理最大的一个人脸
|
|
faces.sort(key=lambda x: (x[2]-x[0])*(x[3]-x[1]), reverse=True)
|
|
face_box = faces[0]
|
|
|
|
# 绘制人脸框
|
|
draw_face_box(display_frame, face_box)
|
|
|
|
# 提取人脸特征
|
|
face_feature = recognizer.extract_face_features(frame, face_box[:4])
|
|
|
|
if face_feature is not None:
|
|
# 保存人脸特征到数据库
|
|
feature_id = db.add_face_feature(person_id, face_feature)
|
|
|
|
face_count += 1
|
|
last_capture_time = current_time
|
|
print(f"已采集第 {face_count}/{target_count} 张人脸")
|
|
voice.speak(f"已完成第{face_count}次采集", block=False)
|
|
|
|
# 显示图像
|
|
cv2.imshow("人脸注册", display_frame)
|
|
|
|
# 按ESC键退出
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == 27:
|
|
break
|
|
|
|
# 释放资源
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
db.close()
|
|
|
|
if face_count >= target_count:
|
|
print(f"成功完成{name}的人脸注册!")
|
|
voice.speak(f"已完成{name}的人脸信息注册", block=True)
|
|
else:
|
|
print("人脸注册未完成")
|
|
|
|
if __name__ == "__main__":
|
|
register_face()
|