使用 YOLO 检测眼睛闭合 | 设定警报

pip install opencv-python
pip install ultralytics
1.
2.

import cv2
from ultralytics import YOLO


# Load the pre-trained YOLO model
model = YOLO("best.pt")
names = model.names
# Open the video file
cap = cv2.VideoCapture("video.mp4")
assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
# Initialize video writer to save the output
video_writer = cv2.VideoWriter("output.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Initialize variables for eye closure detection
eye_closed_frames = 0
eye_closed_threshold_seconds = 1  # Threshold in seconds
eye_closed_threshold_frames = eye_closed_threshold_seconds * fps  # Convert seconds to frames
while cap.isOpened():
    success, im0 = cap.read()
    if not success:
        print("Video frame is empty or video processing has been successfully completed.")
        break
    # Predict the state of the eyes using YOLO
    results = model.predict(im0, show=False)
    boxes = results[0].boxes.xyxy.cpu().tolist()
    clss = results[0].boxes.cls.cpu().tolist()
    annotator = Annotator(im0, line_width=2, example=names)
    eye_closed = False  # Flag to check if the eye is closed in the current frame
If boxes are not None:
        for box, cls in zip(boxes, clss):
            clsName = names[int(cls)]
            xmax = int(box[0])
            ymin = int(box[1])
            xmin = int(box[2])
            ymax = int(box[3])
            # Set color based on the class name
            if clsName == 'closed':
                clr = (0, 0, 255)
                eye_closed = True  # Mark eye as closed
            elif clsName == 'opened':
                clr = (0, 255, 0)
            # Draw the bounding box and label
cv2.FONT_HERSHEY_SIMPLEX
            Font_scale = 1
            Font_thickness = 2
            tw, th = cv2.getTextSize(clsName, font, font_scale, font_thickness)[0]
            cv2.rectangle(im0, (xmin, ymin), (xmax, ymax), color=clr, thickness=2)
            cv2.putText(im0, clsName, (xmax, ymin - 5), font, font_scale, color=clr, thickness=font_thickness)
    # Check for eye closure duration
    if eye_closed:
        Eye_closed_frames += 1
    else:
# Reset counter if the eye is not closed
    # Display warning if eye has been closed for more than the threshold
    if eye_closed_frames > eye_closed_threshold_frames:
        print("Warning: Eye has been closed for more than 2 seconds!")
        cv2.putText(im0, "WARNING: Eye closed for more than 2 seconds!", (50, 50), font, font_scale, (0, 0, 255), font_thickness)
    # Write the processed frame to the output video
    video_writer.write(im0)
# Release resources
cap.release()
video_writer.release()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.

cap = cv2.VideoCapture(0)  # Use 0 for the default camera, or 1, 2, etc. for other cameras
1.

cv2.imshow("Eye Closure Detection", im0) if cv2.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit     break
1.