import cv2 as cv
import numpy as np
import dlib
import joblib


# 画框函数
def draw_rect(img, faces):
    for face in faces:
        cv.rectangle(img, (face.left(), face.top()), (face.right(), face.bottom()), (0, 0, 255), 32)
        return img


# 检测人脸并画框
def detect_face(img):
    img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    result = detector(img_gray, 1)
    img_result = draw_rect(img, result)
    return img_result


# 提取人脸特征
def extract_face_feature(img):
    faces = detector(img)
    if len(faces) > 0:
        face = faces[0]
        landmarks = predictor(img, face)
        face_descriptor = face_rec_model.compute_face_descriptor(img, landmarks)
        return np.array(face_descriptor)
    else:
        return None


# 加载分类器
classifier = joblib.load('models/classifier.pkl')

# 读取测试图片
positive_sample = cv.imread('cache/positive_test.jpg')
positive_sample_rgb = cv.cvtColor(positive_sample, cv.COLOR_BGR2RGB)

negative_sample = cv.imread('cache/love.jpg')
negative_sample_rgb = cv.cvtColor(negative_sample, cv.COLOR_BGR2RGB)

# 加载dlib人脸检测模型
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("models/shape_predictor_68_face_landmarks_GTX.dat")
face_rec_model = dlib.face_recognition_model_v1('models/dlib_face_recognition_resnet_model_v1.dat')

# 检测人脸
positive_result = detect_face(positive_sample)
negative_result = detect_face(negative_sample)

# 计算人脸特征
positive_feature = extract_face_feature(positive_sample_rgb)
positive_feature = np.array(positive_feature).reshape(1, -1)

negative_feature = extract_face_feature(negative_sample_rgb)
negative_feature = np.array(negative_feature).reshape(1, -1)

# 通过分类器预测
positive_predict = classifier.predict(positive_feature)
negative_predict = classifier.predict(negative_feature)

# 展示
positive_result_resized = cv.resize(positive_result,
                                    (int(positive_result.shape[1] * 0.1), int(positive_result.shape[0] * 0.1)))
positive_result_resized = cv.putText(positive_result_resized, str(positive_predict), (50, 50), cv.FONT_HERSHEY_SIMPLEX,
                                     1, (0, 0, 255), 2)

negative_result_resized = cv.resize(negative_result,
                                    (int(negative_result.shape[1] * 0.1), int(negative_result.shape[0] * 0.1)))
negative_result_resized = cv.putText(negative_result_resized, str(negative_predict), (50, 50), cv.FONT_HERSHEY_SIMPLEX,
                                     1, (0, 0, 255), 2)

cv.imshow('positive_result', positive_result_resized)
cv.imshow('negative_result', negative_result_resized)
cv.waitKey(0)
cv.destroyAllWindows()