"""
Name : tired_detect.py
Author  : guiyu.feng
Contect : 1204748417@qq.com
Time    : 2022/5/29 16:23
Desc:
"""
import numpy as np
from scipy.spatial import distance as dist
from imutils import face_utils
import dlib
import imutils
import cv2
import os
import sklearn.preprocessing as sp
import sklearn.ensemble as se
import sklearn.model_selection as ms
from sklearn import tree
import pickle

# 计算眼睛长宽比
def eye_aspect_ratio(eye):
    # 眼部特征点上下宽度有2对特征点表示
    A = dist.euclidean(eye[1], eye[5])# 计算两个集合之间的欧式距离
    B = dist.euclidean(eye[2], eye[4])
    # 计算水平之间的欧几里得距离
    # 部特征点左右长度有1对特征点表示
    C = dist.euclidean(eye[0], eye[3])
    # 眼睛长宽比的计算
    ear = (A + B) / (2.0 * C)
    # 返回眼睛的长宽比
    return ear

# 计算嘴巴长宽比
def mouth_aspect_ratio(mouth):
    #上下宽度
    A = np.linalg.norm(mouth[2] - mouth[9])  # 51, 59
    B = np.linalg.norm(mouth[4] - mouth[7])  # 53, 57
    #左右长度
    C = np.linalg.norm(mouth[0] - mouth[6])  # 49, 55
    mar = (A + B) / (2.0 * C)
    return mar

# 计算前脸检测框面积
def Calarea(rect):
    # width = rect[0][1][0] - rect[0][0][0]
    # height = rect[0][1][1] - rect[0][0][1]

    left = rect.left()
    top = rect.top()
    right = rect.right()
    bottom = rect.bottom()
    width = right -left
    height = bottom - top
    area = width * height
    return area

# 从多个前脸检测框中得到驾驶员的检测框
def GetDriverRect(rects):
    max_rect = rects[0]
    maxa_area = Calarea(max_rect)
    for rect in rects:
        area = Calarea(rect)
        if area > maxa_area:
            max_rect =rect
    return max_rect

# 调用dlib获得分类模型输入数据和输出真值
def Detection(image_path, save_file):
    file = open(save_file, 'w')
    # 眼睛长宽比
    # 闪烁阈值
    EYE_AR_THRESH = 0.2; EYE_AR_CONSEC_FRAMES = 2
    # 打哈欠长宽比
    # 闪烁阈值
    MAR_THRESH = 0.5; MOUTH_AR_CONSEC_FRAMES = 1
    # 初始化前脸丢失检测次数
    MISS_DET_COUNTER = 0
    # 前脸丢失检测阈值
    MISS_DET_THRE = 1
    # 初始化帧计数器和眨眼总数
    COUNTER = 0; TOTAL = 0
    # 初始化帧计数器和打哈欠总数
    mCOUNTER = 0; mTOTAL = 0
    # 获取人脸数据
    # 获得前脸检测模型
    detector = dlib.get_frontal_face_detector()
    # 获得脸部形状特征点检测模型，需要shape_predictor_68_face_landmarks.dat文件
    predictor = dlib.shape_predictor("./shape_predictor_68_face_landmarks.dat")
    # 人脸中左眼、右眼、嘴巴的标记点在整个人脸的标记点中的索引区间
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
    (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
    if len(os.listdir(image_path)) > 0:
        for img in os.listdir(image_path):
            # 读取单张图片
            # 单帧数据
            data = {"ear": 0,"blink_counter": 0, "mar": 0, "yawn_counter": 0, "miss_detect": 0,
                    "danger": 0, "very_danger": 0}
            image = cv2.imread(os.path.join(image_path, img))
            image = imutils.resize(image, width=500)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            # 获得前脸框
            rects = detector(gray, 1)
            if len(rects) > 0:
                # 重置丢失检测计数
                MISS_DET_COUNTER = 0
                # 只取司机的即框面积最大的检测rect
                if len(rects) > 1:
                    rect = GetDriverRect(rects)
                else:
                    rect = rects[0]
                # 利用前脸检测得到的人脸框来检测人脸的标记点
                shape = predictor(gray, rect)  # 标记人脸中的68个landmark点
                shape = face_utils.shape_to_np(shape)  # shape转换成68个坐标点矩阵
                # 提取左眼和右眼坐标
                leftEye = shape[lStart:lEnd]
                rightEye = shape[rStart:rEnd]
                # 嘴巴坐标
                mouth = shape[mStart:mEnd]
                # 构造函数计算左右眼的EAR值，使用平均值作为最终的EAR
                leftEAR = eye_aspect_ratio(leftEye)
                rightEAR = eye_aspect_ratio(rightEye)
                ear = (leftEAR + rightEAR) / 2.0
                if ear < EYE_AR_THRESH:  # 眼睛长宽比：0.2
                    COUNTER += 1
                else:
                    # 如果连续2次都小于阈值，则表示进行了一次眨眼活动
                    if COUNTER >= EYE_AR_CONSEC_FRAMES:  # 阈值：2
                        TOTAL += 1
                    # 重置眼帧计数器
                    COUNTER = 0
                data["ear"] = float(ear)
                data["blink_counter"] = TOTAL
                # 是否打哈欠判断
                mar = mouth_aspect_ratio(mouth)
                if mar > MAR_THRESH:  # 张嘴阈值0.5
                    mCOUNTER += 1
                    # cv2.putText(image, "Yawning!", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                else:
                    # 如果连续3次都大于阈值，则表示打了一次哈欠
                    if mCOUNTER >= MOUTH_AR_CONSEC_FRAMES:  # 阈值：1
                        mTOTAL += 1
                    # 重置嘴帧计数器
                    mCOUNTER = 0
                data["mar"] = float(mar)
                data["yawn_counter"] = mTOTAL
                # 疲劳驾驶造成的危险驾驶，疲劳驾驶后仍然打哈欠判断为特别危险驾驶
                if TOTAL >= 35 or mTOTAL >= 15:
                    data["danger"] = 1
                    if mar > 0.7 or ear < 0.1:
                        data["very_danger"] = 1

                file.write(str(list(data.values())).replace('[', '').replace(']', '') + "\n")
                (x, y, w, h) = face_utils.rect_to_bb(rect)  # 返回人脸框的左上角坐标和矩形框的尺寸
                cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

                cv2.putText(image, "Face #{}".format("driver"), (x - 10, y - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                for (x, y) in shape:
                    cv2.circle(image, (x, y), 2, (0, 0, 255), -1)
                cv2.imshow("Output", image)
                cv2.waitKey(1)
            else:
                MISS_DET_COUNTER += 1
                #驾驶员没有面朝前方造成的危险驾驶
                data["miss_detect"] = MISS_DET_COUNTER
                if MISS_DET_COUNTER > MISS_DET_THRE: # 1
                    data["danger"] = 1
                if MISS_DET_COUNTER > 2 * MISS_DET_THRE:
                    data["very_danger"] = 1
                # file.write(str(data).replace('{', '').replace('}', '') + '\n')
                file.write(str(list(data.values())).replace('[', '').replace(']', '') + "\n")
                cv2.imshow("Output", image)
                cv2.waitKey(1)
    else:
        print("over!")
        file.close()

# 可视化随机森林
def draw_randomforest_trees(model, feature_names: list, class_names: list,
                            save_dir: str, model_name: str) -> list:
    print("可视化随机森林的树")
    resimgpaths = []  # 图片保存路径
    for i in range(len(model.estimators_)):
        # 将树导出为.dot文件和jpg
        savedotfile = os.path.join(save_dir, model_name + "_" + str(i) + ".dot")
        saveimgfile = os.path.join(save_dir, model_name + "_" + str(i) + ".jpg")
        print(f"第{i}棵树 {saveimgfile}")
        tree.export_graphviz(model.estimators_[i],
                        out_file=savedotfile,
                        feature_names=feature_names,
                        class_names=class_names,
                        rounded=True,
                        proportion=False,
                        precision=6,
                        filled=True)

        # 使用系统命令把dot转换为jpg图片 dot -Tjpg tree.dot -o tree.jpg
        os.system(f"dot -Tpng {savedotfile} -o {saveimgfile}")
        resimgpaths.append(saveimgfile)
    return resimgpaths

# 训练随机森林模型并测试模型得分
def DecisionTree(data_file, datatest_file, model_full_name):
    # 数据读取
    data = np.loadtxt(data_file, dtype=float, delimiter=',')
    data_test = np.loadtxt(datatest_file, dtype=float, delimiter=',')

    # 数据清洗
    # 防止有全为0的数据或异常小的数据, 删除缺少元素的数据
    del_index = []
    for index, item in enumerate(data):
        if sum(item) <= 0.000001:
            del_index.append(index)
        if len(item) != 7:
            del_index.append(index)
    if len(del_index) > 0:
        data = np.delete(data, del_index)
    # print(data)
    test_del_index = []
    for index, item in enumerate(data_test):
        if sum(item) <= 0.000001:
            test_del_index.append(index)
        if len(item) != 7:
            test_del_index.append(index)
    if len(test_del_index) > 0:
        data_test = np.delete(data_test, test_del_index)
    print(data_test)

    # 数据处理
    data = data.T
    data_test = data_test.T
    train_x, train_y = [], []
    test_x, test_y = [], []
    for row in range(len(data)):
        # 创建适用于当前特征的标签编码器
        encoder = sp.LabelEncoder()
        # 获得输入
        if row < len(data) - 2:
            train_x.append(
                encoder.fit_transform(data[row]))
        else:
        # 获得输出真实值
            train_y.append(encoder.fit_transform(data[row]))
    for row in range(len(data_test)):
        # 创建适用于当前特征的标签编码器
        encoder_test = sp.LabelEncoder()
        # 获得输入
        if row < len(data_test) - 2:
            test_x.append(
                encoder_test.fit_transform(data_test[row]))
        else:
        # 获得输出真实值
            test_y.append(encoder_test.fit_transform(data_test[row]))
    train_x = np.array(train_x).T
    train_y = np.array(train_y).T
    test_x = np.array(test_x).T
    test_y = np.array(test_y).T

    # 构造模型,选择随机森林分类模型
    # max_depth:树的最大深度
    # random_state：控制生成随机森林的模式。随机性越大，模型效果越好，当然这样可能就不是很稳定，不便于调试。想要模型稳定，可以设置random_state参数
    # n_estimators：森林里决策树的数目，不考虑过拟合，应该n_estimators越大越好，但占用的内存与训练和预测的时间也会相应增长
    model = se.RandomForestClassifier(max_depth=8,
                                      n_estimators=4, random_state=7)
    # 模型训练
    model.fit(train_x, train_y)
    print(train_x.shape, train_y.shape)

    # 模型可视化
    draw_randomforest_trees(model,["ear","blink_counter", "mar", "yawn_counter", "miss_detect"],
                            ["danger", "very_danger"], "./model", "dangerdrive_detect_model")



    # 模型预测
    y_predict = model.predict(test_x)
    # print('真实值\n', test_y)
    # print('预测值\n', y_predict)
    # 交叉验证，输出训练集，测试集分类F1得分
    acc_train = ms.cross_val_score(model, train_x,
                               train_y, cv=2, scoring='accuracy')
    score_train = ms.cross_val_score(model, train_x,
                               train_y, cv=2, scoring='f1_weighted')
    print('train accuracy:', acc_train.mean())
    print('train F1 score:',score_train.mean())
    acc_test = ms.cross_val_score(model, test_x,
                                   test_y, cv=2, scoring='accuracy')
    score_test = ms.cross_val_score(model, test_x,
                               test_y, cv=2, scoring='f1_weighted')
    print('test accuracy:', acc_test.mean())
    print('test F1 score:',score_test.mean())
    print('模型得分 ', model.score(test_x, test_y))

    # 保存模型
    with open(model_full_name, 'wb') as f:
        pickle.dump(model, f)

# 实际场景测试
def DecisionTreeTest(test_image_path, model_full_name):
    with open(model_full_name, 'rb') as f:
        model = pickle.load(f)
    # 眼睛长宽比
    # 闪烁阈值
    EYE_AR_THRESH = 0.2;
    EYE_AR_CONSEC_FRAMES = 2
    # 打哈欠长宽比
    # 闪烁阈值
    MAR_THRESH = 0.5;
    MOUTH_AR_CONSEC_FRAMES = 1
    # 初始化前脸丢失检测次数
    MISS_DET_COUNTER = 0
    # 前脸丢失检测阈值
    MISS_DET_THRE = 1
    # 初始化帧计数器和眨眼总数
    COUNTER = 0;
    TOTAL = 0
    # 初始化帧计数器和打哈欠总数
    mCOUNTER = 0;
    mTOTAL = 0
    # 获取人脸数据
    # 获得前脸检测模型
    detector = dlib.get_frontal_face_detector()
    # 获得脸部形状特征点检测模型
    predictor = dlib.shape_predictor("./shape_predictor_68_face_landmarks.dat")
    # 人脸中左眼、右眼、嘴巴的标记点在整个人脸的标记点中的索引区间
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
    (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
    if len(os.listdir(image_path)) > 0:
        for img in os.listdir(image_path):
            # 读取单张图片
            # 单帧数据
            data = {"ear": 0, "blink_counter": 0, "mar": 0, "yawn_counter": 0, "miss_detect": 0,
                    "danger": 0, "very_danger": 0}
            image = cv2.imread(os.path.join(image_path, img))
            image = imutils.resize(image, width=500)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            # 获得前脸框
            rects = detector(gray, 1)
            if len(rects) > 0:
                # 重置丢失检测计数
                MISS_DET_COUNTER = 0
                # 只取司机的即框面积最大的检测rect
                if len(rects) > 1:
                    rect = GetDriverRect(rects)
                else:
                    rect = rects[0]
                # 利用前脸检测得到的人脸框来检测人脸的标记点
                shape = predictor(gray, rect)  # 标记人脸中的68个landmark点
                shape = face_utils.shape_to_np(shape)  # shape转换成68个坐标点矩阵
                # 提取左眼和右眼坐标
                leftEye = shape[lStart:lEnd]
                rightEye = shape[rStart:rEnd]
                # 嘴巴坐标
                mouth = shape[mStart:mEnd]
                # 构造函数计算左右眼的EAR值，使用平均值作为最终的EAR
                leftEAR = eye_aspect_ratio(leftEye)
                rightEAR = eye_aspect_ratio(rightEye)
                ear = (leftEAR + rightEAR) / 2.0
                if ear < EYE_AR_THRESH:  # 眼睛长宽比：0.2
                    COUNTER += 1
                else:
                    # 如果连续2次都小于阈值，则表示进行了一次眨眼活动
                    if COUNTER >= EYE_AR_CONSEC_FRAMES:  # 阈值：2
                        TOTAL += 1
                    # 重置眼帧计数器
                    COUNTER = 0
                data["ear"] = float(ear)
                data["blink_counter"] = TOTAL
                # 是否打哈欠判断
                mar = mouth_aspect_ratio(mouth)
                if mar > MAR_THRESH:  # 张嘴阈值0.5
                    mCOUNTER += 1

                else:
                    # 如果连续3次都大于阈值，则表示打了一次哈欠
                    if mCOUNTER >= MOUTH_AR_CONSEC_FRAMES:  # 阈值：1
                        mTOTAL += 1
                    # 重置嘴帧计数器
                    mCOUNTER = 0
                data["mar"] = float(mar)
                data["yawn_counter"] = mTOTAL
                (x, y, w, h) = face_utils.rect_to_bb(rect)  # 返回人脸框的左上角坐标和矩形框的尺寸
                cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

                cv2.putText(image, "Face #{}".format("driver"), (x - 10, y - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                for (x, y) in shape:
                    cv2.circle(image, (x, y), 2, (0, 0, 255), -1)
                test_x = list(data.values())[0:5]
                test_y = model.predict([test_x])
                data["danger"] = test_y[0][0]
                data["very_danger"] = test_y[0][1]
            else:
                MISS_DET_COUNTER += 1
                # 驾驶员没有面朝前方造成的危险驾驶
                data["miss_detect"] = MISS_DET_COUNTER
                test_x = list(data.values())[0:5]
                test_y = model.predict([test_x])
                data["danger"] = test_y[0][0]
                data["very_danger"] = test_y[0][1]
            if data["danger"] == 1:
                cv2.putText(image, "danger!", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            if data["very_danger"] == 1:
                cv2.putText(image, "very_danger!", (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            cv2.imshow("Output", image)
            # 按空格按帧查看，或者修改0为1连续查看
            cv2.waitKey(0)
    else:
        print("over!")

if __name__ == '__main__':
    # 数据保存地址
    data_file = './data.txt'
    datatest_file = './datatest.txt'
    # 读取图片地址，在已经有data.txt和datatest.txt数据的时候不需要也可以
    image_path = './train'
    test_image_path = './test'
    # 训练模型保存地址
    model_save_name = "./model/save2.pickle"

    # 检测图片获得特征以及脸部活动来获得训练数据data_file和测试数据datatest_file
    # 在已经有data.txt和datatest.txt的时候，不需要运行也可以
    # Detection(image_path, data_file)
    # Detection(test_image_path, datatest_file)

    # 训练随机森林分类模型以及模型得分
    # DecisionTree(data_file, datatest_file, model_save_name)

    # 应用场景测试，调用训练保存的.pickle文件作为分类模型，需要有图像输入
    DecisionTreeTest(test_image_path, model_save_name)