import cv2
import cv2 as cv
import dlib
import joblib
import numpy as np
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from tqdm import tqdm


# 方法区
## 转换为二值图和RGB图
def img_cvt(img):
    img_gray = cv.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, img_gray = cv.threshold(img_gray, 127, 255, cv.THRESH_BINARY)
    img_gray = cv.blur(img_gray, (3, 3))
    img_rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    return img_gray, img_rgb


## 提取人脸特征
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


## 生成伪样品
def gen_negative_samples(features_dim=128):
    return np.random.uniform(-1, 1, features_dim)


# 变量区
dataset_video = cv.VideoCapture('cache/dataset.mp4')
datasets = []
datasets_rgb = []
features = []
negative_samples = []

# 收集数据集
## 采用视频的格式收集图片
while (True):
    ret, frame = dataset_video.read()
    if not ret:
        break
    datasets.append(frame)

# 加载人脸检测模型
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')

# 预处理数据集
for dataset in tqdm(datasets, desc='正在预处理数据集'):
    _, dataset_rgb = img_cvt(dataset)
    datasets_rgb.append(dataset_rgb)

# 求人脸特征
for dataset in tqdm(datasets_rgb, desc='正在求人脸特征'):
    face_descriptor = extract_face_feature(dataset)
    if face_descriptor is not None:
        features.append(face_descriptor)

# 生成负样本
for i in tqdm(range(len(features) * 2048), desc='正在生成负样本'):
    negative_sample = gen_negative_samples()
    negative_samples.append(negative_sample)

# 编码数据集
positive_labels = [1] * len(features)
negative_labels = [0] * len(negative_samples)

samples = features + negative_samples
labels = positive_labels + negative_labels

X_train = np.array(samples)
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, stratify=Y_train)

# 训练
print('正在训练分类器')
classifier = SVC(kernel='linear', C=1.0, random_state=42)
classifier.fit(X_train, Y_train)

# 评估分类器性能
Y_pred = []
for test_sample in tqdm(X_test, desc='评估分类器'):
    Y_pred.append(classifier.predict(test_sample.reshape(1, -1)))
accuracy = accuracy_score(Y_test, Y_pred)
print(f'准确率：{accuracy * 100:.5f}%')

# 保存模型
joblib.dump(classifier, 'models/classifier.pkl')