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

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputClassifier
from tqdm import tqdm

# 提取人脸特征
def extract_face_features(image_path):
    img = cv.imread(image_path)
    detections = detector(img, 1)
    face_features = []
    face_bboxes = []
    for rect in detections:
        shape = predictor(img, rect)
        face_descriptor = face_rec_model.compute_face_descriptor(img, shape)
        face_features.append(face_descriptor)
        bbox = (rect.left(), rect.top(), rect.right(), rect.bottom())
        face_bboxes.append(bbox)
    return face_features, face_bboxes

# 获取图片路径
def get_img_path(directory, extension=None):
    if extension is None:
        extension = ['.jpg', '.jpeg', '.png']
    files = []
    for root, dirs, file_names in os.walk(directory):
        for file_name in file_names:
            if any(file_name.lower().endswith(ext) for ext in extension):
                files.append(os.path.join(root, file_name))
    return files


# 加载dlib模型
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('models/shape_predictor_68_face_landmarks.dat')
face_rec_model = dlib.face_recognition_model_v1('models/dlib_face_recognition_resnet_model_v1.dat')

# 获取图片
img_directory = 'cache/pretrained/'
images = get_img_path(img_directory)

# 提取特征
features = []
labels = []

for image in tqdm(images, desc='提取图片特征中：'):
    extracted_features, face_bboxes = extract_face_features(image)
    for feature,bbox in zip(extracted_features, face_bboxes):
        features.append(feature)
        labels.append(bbox)

X_train = np.array(features)
Y_train = np.array(labels)

# 分割测试集
X_train, X_test, Y_train, Y_test = train_test_split(X_train, Y_train, test_size=0.2, random_state=42)

# 训练SVM模型
print('训练模型中')
clf=MultiOutputClassifier(RandomForestClassifier(n_estimators=100,random_state=42))
clf.fit(X_train, Y_train)

# 评估训练数据
predictions = clf.predict(X_test)
accuracy=(predictions==Y_test).mean()
print(f'分类器准确度：{accuracy * 100:.2f}%')

os.makedirs('models/', exist_ok=True)
joblib.dump(clf, 'models/my_classifier.pkl')