import cv2
import face_recognition
import os
import sqlite3
import numpy as np

# 初始化摄像头
cap = cv2.VideoCapture(0)
photo_count = 0
max_photos = 10
captured_images = []

# 创建目录以保存照片
save_path = "./captured_faces"
os.makedirs(save_path, exist_ok=True)

while photo_count < max_photos:
    ret, frame = cap.read()
    if not ret:
        break

    # 将图像转换为RGB颜色
    rgb_frame = frame[:, :, ::-1]

    # 检测人脸
    face_locations = face_recognition.face_locations(rgb_frame)

    for face_location in face_locations:
        top, right, bottom, left = face_location

        # 在图像上绘制绿框
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)

        # 提取人脸区域
        face_image = frame[top:bottom, left:right]

        # 保存抓拍的照片
        image_path = os.path.join(save_path, f"face_{photo_count + 1}.jpg")
        cv2.imwrite(image_path, face_image)
        captured_images.append(image_path)

        photo_count += 1
        if photo_count >= max_photos:
            break

    # 显示结果
    cv2.imshow("Capturing Faces", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

print(f"Captured {photo_count} images.")

def create_face_database(db_name="face_database.db"):
    conn = sqlite3.connect(db_name)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS faces
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                 name TEXT NOT NULL,
                 encoding BLOB NOT NULL)''')
    conn.commit()
    conn.close()

def add_face_to_database(name, image_path, db_name="face_database.db"):
    conn = sqlite3.connect(db_name)
    c = conn.cursor()

    # 加载图片并生成编码
    image = face_recognition.load_image_file(image_path)
    face_encodings = face_recognition.face_encodings(image)

    if face_encodings:
        face_encoding = face_encodings[0]
        # 将编码转换为可以存储的格式
        encoding_blob = np.array(face_encoding).tobytes()
        c.execute("INSERT INTO faces (name, encoding) VALUES (?, ?)",
                  (name, encoding_blob))
        conn.commit()
    conn.close()

def match_faces(captured_images, db_name="face_database.db", tolerance=0.4):
    conn = sqlite3.connect(db_name)
    c = conn.cursor()

    # 获取数据库中所有存储的人脸编码
    c.execute("SELECT name, encoding FROM faces")
    known_faces = c.fetchall()

    for image_path in captured_images:
        # 加载待匹配图片并生成编码
        unknown_image = face_recognition.load_image_file(image_path)
        face_encodings = face_recognition.face_encodings(unknown_image)

        if len(face_encodings) == 0:
            print(f"没有人脸 {image_path}")
            continue  # 如果没有检测到人脸，跳过该图片

        unknown_encoding = face_encodings[0]

        for name, encoding_blob in known_faces:
            known_encoding = np.frombuffer(encoding_blob, dtype=np.float64)
            match = face_recognition.compare_faces([known_encoding], unknown_encoding, tolerance=tolerance)

            if match[0]:  # 如果匹配成功
                print(f"发现匹配: {name} 在 {image_path}")
                conn.close()
                return True  # 一旦找到匹配，返回成功
    print(f"没发现匹配: 在 {image_path}")
    conn.close()
    return False  # 如果所有比较都没有匹配，返回失败

# 创建人脸数据库
create_face_database()

# 向数据库中添加人脸
#add_face_to_database("屈礼", "./db_image/test.jpg")

# 逐张匹配抓拍的照片
if match_faces(captured_images):
    print("至少一张匹配")
else:
    print("没有匹配")