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
        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(image_path, db_name="face_database.db"):
    conn = sqlite3.connect(db_name)
    c = conn.cursor()

    # 加载待匹配图片并生成编码
    unknown_image = face_recognition.load_image_file(image_path)
    unknown_encoding = face_recognition.face_encodings(unknown_image)[0]

    c.execute("SELECT name, encoding FROM faces")
    matches = []
    for row in c.fetchall():
        name, encoding_blob = row
        known_encoding = np.frombuffer(encoding_blob, dtype=np.float64)
        match = face_recognition.compare_faces([known_encoding], unknown_encoding)
        matches.append((name, match))

    conn.close()
    return matches


# 创建人脸数据库
create_face_database()

# 向数据库中添加人脸
add_face_to_database("Alice", "captured_faces/face_1.jpg")

# 匹配人脸
for captured_image in captured_images:
    matches = match_faces(captured_image)
    for name, match in matches:
        if match[0]:
            print(f"Match found: {name} in image {captured_image}")