import cv2
import face_recognition
import sqlite3
import numpy as np

from match_face import save_face_encoding


def add_new_face(name):
    cap = cv2.VideoCapture(0)

    while True:
        ret, frame = cap.read()
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        face_locations = face_recognition.face_locations(rgb_frame)
        face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)

        if face_encodings:
            # 假设只处理第一张检测到的人脸
            encoding = face_encodings[0]
            save_face_encoding(name, encoding)
            print(f"Face of {name} saved!")
            break

        cv2.imshow('Capture New Face', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


# 示例用法
add_new_face("John Doe")