import cv2
import face_recognition
import os
import numpy as np
import pickle 
datasetPath = "dataset"
images = []
classNames = []
lisFileTrain = os.listdir(datasetPath)

for file in lisFileTrain:
    currentImg = cv2.imread(f"{datasetPath}/{file}")
    images.append(currentImg)
    classNames.append(os.path.splitext(file)[0].split('_')[0])

print(len(images))

def encodeImgs(images, save_path="encodings.pkl"):
    if os.path.exists(save_path):
        print(f"Loading encodings from {save_path}...")
        with open(save_path, "rb") as f:
            return pickle.load(f)
        
    encodeList = []
    for i, img in enumerate(images):
        print(i+1)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        encode = face_recognition.face_encodings(img)
       
        if encode:  # Check if encodings list is not empty
            encodeList.append(encode[0])
        else:
            print("No face detected in an image. Skipping...")
            os.remove(f"{datasetPath}/{lisFileTrain[i]}")
    # Lưu encodeList vào file
    print(f"Saving encodings to {save_path}...")
    with open(save_path, "wb") as f:
        pickle.dump(encodeList, f)

    return encodeList

encodeListKnow = encodeImgs(images)
print("Load data success")
print(len(encodeListKnow))

cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    frameS = cv2.resize(frame, (0,0), None, fx=1, fy=1)
    frameS = cv2.cvtColor(frameS, cv2.COLOR_BGR2RGB)

    faceCurFrame = face_recognition.face_locations(frameS)
    encodeCurFrame = face_recognition.face_encodings(frameS)

    for encodeFace, faceLoc in zip(encodeCurFrame, faceCurFrame):
        matches = face_recognition.compare_faces(encodeListKnow, encodeFace)
        faceDis = face_recognition.face_distance(encodeListKnow, encodeFace)
        print(faceDis)
        matchIndex = np.argmin(faceDis)

        if faceDis[matchIndex] < 0.3:
            name = classNames[matchIndex].upper()
        else:
            name = "Unknow"
        
        y1, x2, y2, x1 = faceLoc
        y1, x2, y2, x1 = y1, x2, y2, x1
        cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2)
        cv2.putText(frame, name + f"({(1 - round(faceDis[matchIndex], 2))*100}%)", (x2, y2), cv2.FONT_HERSHEY_COMPLEX, 1, (255,255,255), 2)
    cv2.imshow('Face decting', frame)

    if cv2.waitKey(1) == ord("q"):
        break

cap.release()
cv2.destroyAllWindows()