import cv2
import numpy as np
from keras.models import load_model
from statistics import mode
from PIL import Image, ImageDraw, ImageFont

# 加载模型
model = load_model('modelv2.h5')
#

# 创建分类标签

class_labels = ['生气', '厌恶', '害怕', '开心', '悲伤', '惊讶', '正常']

# 设置人脸检测器
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

# 开启摄像头
cap = cv2.VideoCapture(0)

# 设置计数器和标志位
counter = 0
flag = False

# 设置字体和字体大小
font = ImageFont.truetype('eujm781krhhpl31ju1c1ur0b7vs4mhz6.ttf', 36)

while True:
    # 读取图像
    ret, frame = cap.read()
    if not ret:
        break

    # 转换为灰度图像并进行人脸检测
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    # 预测人脸情绪并显示标签
    for (x, y, w, h) in faces:
        face_roi = gray[y:y + h, x:x + w]
        face_roi = cv2.resize(face_roi, (48, 48))
        face_roi = np.expand_dims(face_roi, -1)
        face_roi = np.expand_dims(face_roi, 0)
        result = model.predict(face_roi)[0]
        label = class_labels[np.argmax(result)]

        # 创建PIL图像对象
        pil_img = Image.fromarray(frame)

        # 创建绘图对象
        draw = ImageDraw.Draw(pil_img)

        # 绘制文字
        draw.text((x, y), label, font=font, fill=(0, 255, 0))

        # 将PIL图像转换为OpenCV图像
        frame = np.array(pil_img)

        # 计算走神人数
        if label in ['悲伤', '正常']:
            counter += 1
            if counter >= 5 * 30:
                flag = True
        else:
            counter = 0
            flag = False

    # 绘制人脸矩形框
    for (x, y, w, h) in faces:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

    # 判断是否走神
    num_faces = len(faces)
    if num_faces > 0:
        if flag:
            cv2.putText(frame, f'{num_faces} people are drowsy', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2,
                        cv2.LINE_AA)
        else:

            cv2.putText(frame, f'{num_faces} people are not drowsy', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0),
                        2, cv2.LINE_AA)

    # 显示图像
    cv2.imshow('Drowsiness Detection', frame)

    # 按'q'键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像