Drowsiness Detection using Custom YOLO Classification Model

This project implements a real-time drowsiness detection system using a custom-trained YOLO classification model. The system classifies whether a person's eyes are open or closed and alerts the user when drowsiness is detected.

Table of Contents

• Project Overview
• Dataset
• Installation
• Training a Custom Model
• Drowsiness Detection Application
• Usage
• Model Architecture
• Contributing
• License

Project Overview

This project consists of two main components:

1. Custom YOLO Model Training: Train a YOLO model to detect open and closed eyes using the MRL Eye Dataset
2. Drowsiness Detection Application: Real-time application that uses the trained model to detect drowsiness and alert the user

The system uses computer vision and machine learning to monitor eye state and determine if a person is becoming drowsy, which is particularly useful for driver safety and fatigue monitoring.

Dataset

This project uses the MRL Eye Dataset from Kaggle, which contains images of open and closed eyes. The dataset is used to train a custom YOLO classification model to distinguish between open and closed eyes.

Important Note: The images in this dataset do not require traditional bounding box labels (like in object detection). Instead, this classification model uses folder-based labeling where the folder name serves as the class label.

Dataset Structure

The dataset should be organized with folders named after the classes:

dataset/
├── train/
│   ├── open_eyes/     # All images of open eyes
│   └── closed_eyes/   # All images of closed eyes
├── valid/
│   ├── open_eyes/
│   └── closed_eyes/
└── test/
    ├── open_eyes/
    └── closed_eyes/

This folder-based approach allows the model to automatically learn the class labels from the folder names, eliminating the need for manual annotation with bounding boxes.

Installation

Prerequisites

• Python 3.8 or higher
• pip package manager

Setup

1. Clone this repository:

git clone <repository-url>
cd drowsiness-detection-yolo

2. Install required dependencies:

pip install -r requirements.txt

3. Download the MRL Eye Dataset from Kaggle and organize it in the dataset/ directory with the folder-based structure as shown above.

4. Download a pre-trained YOLO classification model (e.g., yolo11m-cls.pt ) and place it in the project root directory.

Training a Custom Model

Dataset Preparation

1. Download the MRL Eye Dataset from Kaggle
2. Organize your dataset using folder-based labeling:

dataset/
├── train/
│   ├── open_eyes/     # All images of open eyes
│   └── closed_eyes/   # All images of closed eyes
├── valid/
│   ├── open_eyes/
│   └── closed_eyes/
└── test/
    ├── open_eyes/
    └── closed_eyes/

Training Process

Option 1: Using the Python script

1. Place a pre-trained YOLO classification model (e.g., yolo11m-cls.pt) in the project root directory
2. Run the training script:

python train_yolo.py

Option 2: Using the CLI command You can also train the model directly using the Ultralytics CLI without modifying Python code:

yolo classify train data=dataset model=yolo11m-cls.pt epochs=75 imgsz=224 batch=8 project=runs/classify name=train_custom_eye

Note: The default training script is configured to use GPU (device=0). If you don't have a CUDA-compatible GPU, modify the train_yolo.py file to use CPU by changing device=0 to device='cpu'.

The training script will:

• Load the pre-trained YOLOv11 classification model
• Train it on your custom eye dataset using folder-based labels
• Validate the model performance
• Save the trained model

Training Configuration

The training script uses the following configuration:

• Model: YOLOv11 medium classification (can be changed in the script)
• Epochs: 75
• Image Size: 224x224 (standard for classification)
• Batch Size: 8
• Device: GPU (device=0) - change to 'cpu' if no GPU available
• Patience: 100 (early stopping)
• Workers: 8

You can modify these parameters in train_yolo.py as needed.

Drowsiness Detection Application

The drowsiness detection application uses the trained classification model to monitor eye state in real-time and alert the user when drowsiness is detected.

How It Works

1. The application captures video from the default webcam
2. It processes each frame using the trained YOLO classification model to classify whether eyes are open or closed
3. It analyzes the eye state over time to determine if the user is becoming drowsy
4. When drowsiness is detected, it triggers both visual and audio alerts

Detection Logic

The system uses two thresholds to determine drowsiness:

1. Percentage-based: If 60% or more of the last 15 frames show closed eyes
2. Consecutive frames: If 8 or more consecutive frames show closed eyes

Alerts

When drowsiness is detected, the system:

• Shows a visual alert on the screen
• Plays an audio beep
• Sends a desktop notification

Usage

Training the Model

python train_yolo.py

Running the Drowsiness Detection

python drowsiness_detection.py

Press 'q' to quit the application.

Model Architecture

The system uses a YOLO (You Only Look Once) classification model, specifically YOLOv11 classification, for real-time image classification. This classification approach is chosen for its speed and accuracy in categorizing images in video streams, making it ideal for real-time drowsiness detection based on eye state classification.

Configuration

The drowsiness detection system can be configured by modifying parameters in the drowsiness_detection.py file:

• confidence_threshold: Minimum confidence for classification (default: 0.5)
• drowsiness_threshold: Percentage of closed-eye frames to trigger alert (default: 60%)
• consecutive_closed_threshold: Consecutive closed frames for immediate alert (default: 8)
• notification_cooldown: Minimum time between alerts (default: 10 seconds)

Make sure to update the model_path parameter in the main() function to point to your trained classification model (e.g., runs/classify/train_custom_eye/weights/best.pt).

Contributing

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add some amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request

License

This project is open source and available under the MIT License.

Acknowledgments

• MRL Eye Dataset contributors on Kaggle
• Ultralytics for the YOLO implementation
• OpenCV for computer vision capabilities