Drowsiness Detection with YOLOv8

The inference given by the final trained model

This project aims to detect drowsiness using YOLOv8, a state-of-the-art object detection model. The goal is to create a custom model by training it on images collected from a webcam to detect signs of drowsiness, such as closed eyes or head drooping. Once drowsiness is detected, an audio alert is triggered to alert the person and prevent potential accidents. The project focuses on enhancing safety for individuals who drive long distances or work in industries where alertness is crucial, such as shift-based jobs.

If you don’t want to continue reading the article, you can go to the repository directly and read the documentation, everything is well explained and detailed there.

One last note before we get into the details of running the project, is that you can use Yolo to create any classifier you want (if you can get the data), although the example is a drowsiness detector.

As the ultralytics documentation says:

Yolov8 the latest version of the acclaimed real-time object detection and image segmentation model.

This model will be the center of our project, it will allow us to create an image detection model quickly and easily. Again, according to the documentation:

Train mode is used for training a YOLOv8 model on a custom dataset. In this mode, the model is trained using the specified dataset and hyperparameters. The training process involves optimizing the model’s parameters so that it can accurately predict the classes and locations of objects in an image.

More details about the model, I’ll let you explore the documentation which is very good! Furthermore, the community is also very strong. :)

Let’s focus on understanding the steps, which will basically be two: data collection and model training. For a more technical reader, see this architectural model designed for the entire project from end to end.

Drowsiness detection architecture

Cool huh? So let’s go to the first step:

I) Data collection

Here, you need a basic function which is the collection of your data by your own webcam.

def collect_images_from_webcam(
        labels: list, number_imgs: int, images_path: str) -> None:
    '''
    Collects images from the webcam and saves them to a folder.

    Args:
        labels (list): List of labels for the images.
        number_imgs (int): Number of images to be collected for each label.
        images_path (str): Path to the folder where the images will be saved.

    Returns:
        None
    '''
    cap = cv2.VideoCapture(0)

    # Loop through labels
    for label in labels:
        print('Collecting images for {}'.format(label))
        time.sleep(5)

        # Loop through image range
        for img_num in range(number_imgs):
            logging.info(f'Collecting images for {label}, image number {img_num}')

            # Webcam feed
            ret, frame = cap.read()

            # Naming our image path
            imgname = os.path.join(images_path, label + '.' + str(uuid.uuid1()) + '.jpg')

            # Write out image to file
            cv2.imwrite(imgname, frame)

            # Render to the screen
            cv2.imshow('Image Collection', frame)

            # 3-second delay between captures
            time.sleep(3)

            if cv2.waitKey(10) & 0xFF == ord('q'):
                logging.info(f'You broke the sequence of collection for {label} images')
                break

    cap.release()
    cv2.destroyAllWindows()

By filling in the parameters you can collect as much data as you want and it will go to the folder of your choice. One thing I did that is very useful for this type of project with Yolo, was to create a function that already directs the collected images to the correct folder, you will soon see why this is useful.

def split_images(images_path: str, labels_path: str):
    '''
    Split the collected images into train, val, and test sets.

    Args:
        images_path (str): Path to the folder where the images are stored.
        labels_path (str): Path to the folder where the labels are stored.

    Returns:
        None
    '''
    dest_folders = ['train', 'val', 'test']
    for folder in dest_folders:
        os.makedirs(os.path.join(images_path, folder), exist_ok=True)
        os.makedirs(os.path.join(labels_path, folder), exist_ok=True)

    file_list = os.listdir(images_path)
    random.shuffle(file_list)

    num_files = len(file_list)
    train_count = int(TRAIN_RATIO * num_files)
    val_count = int(VAL_RATIO * num_files)

    for i, file_name in enumerate(file_list):
        if file_name.endswith('.jpg'):
            src_path = os.path.join(images_path, file_name)

            if i < train_count:
                dest_folder = 'train'
            elif i < train_count + val_count:
                dest_folder = 'val'
            else:
                dest_folder = 'test'
            
            dest_path = os.path.join(images_path, dest_folder)
            shutil.move(src_path, dest_path)

            logging.info(f'Moved {file_name} to {dest_folder} folder.')

Regarding this first step, this is the most important part! Let’s go to the second step.

II) Model training

For training the model with our data collected in the last step, I created a custom YOLO model function so that we can change the hyperparameters.

def train_custom_yolo_model(
    data: str,
    epochs: int,
    batch: int,
    model_name: str,
    lr0: float,
    lrf: float,
    weight_decay: float) -> tuple:
    '''
    Function to train a custom model using YOLO v8.

    Args:
    - data (str): Data configuration file path. Default is 'data.yaml'.
    - epochs (int): Number of training epochs. Default is 20.
    - batch (int): Batch size. Default is 8.
    - augment (bool): Flag to enable data augmentation. Default is True.
    - model_name (str): Output model name. Default is 'yolov8n_drowsiness'.
    - lr0 (float): Initial learning rate for optimizer. Default is 0.01.
    - lrf (float): Final learning rate for optimizer. Default is 0.01.
    - weight_decay (float): Weight decay for optimizer. Default is 0.0005.

    Returns:
    - results (Tuple): A tuple containing the training results.
    '''
    # Load the pretrained model
    model = YOLO('yolov8n.pt')

    # Training settings
    logging.info('Starting training...')
    logging.info(f'Data configuration: {data}')
    logging.info(f'Epochs: {epochs}')
    logging.info(f'Batch size: {batch}')
    logging.info(f'Model name: {model_name}')
    logging.info(f'Initial learning rate: {lr0}')
    logging.info(f'Final learning rate: {lrf}')
    logging.info(f'Weight decay: {weight_decay}')

    # Train the model
    results = model.train(
        data=data,
        epochs=epochs,
        batch=batch,
        name=model_name,
        lr0=lr0,
        lrf=lrf,
        weight_decay=weight_decay,
    )

    logging.info('Training completed.')
    return results

That’s what you need! Relatively easy, right? Calm down, it’s not that easy, we have a crucial part to do and it’s a little more manual and laborious. This is the part of labeling and creating the yaml file to feed the model training function above. See the “data” parameter above, it is fed by a yaml file, let’s detail it.

III) The yaml file

The Ultralytics framework uses a YAML file format to define the dataset and model configuration for training Detection Models. Here is an example of the YAML format used for defining a detection dataset:

path: /{your_local_path}/data
train: images/train  
val: images/val
test: images/test

# Classes
nc: 2
names: ['awake','drowsy']

You must use this file to feed your model, it is like the bridge between the data you collected and stored in these respective folders (train, val, test) and your YOLO model.

One last thing we should talk about is labeling our data, this was the hardest and most annoying part to do, but it’s necessary.

IV) Labelling your own data

This step is quite complicated to explain in words in an article here on medium, so I’ll redirect you to a very nice content that helped me a lot to do this analysis. The package I used to help me in this step is labelImg.

To use this package and label your data, you can follow this link and start watching from minute 52:20, which is where Nicholas Renotte easily explains how to use it. I even recommend watching the whole video, this is an amazing content!

Well, these are the main stages of this type of project and I hope I have clarified and made your life easier if you want to replicate or do a similar project using this incredible tool from ultralytics that is Yolo.

I recommend that you take a look at my drowsiness detection project documentation to better understand how the project was done as a whole. Here I have only explained the steps that I believe are essential, but for a complete end-to-end project it is much bigger than just what has been explained here.

Oh, and after performing all these steps, you will be able to make real-time inferences on your own webcam, you can even trigger an alarm if the “drowsy” tag is detected! Very cool, right?

Potential Applications:

• Driver Safety: The drowsiness detection model can be integrated into vehicles to alert drivers when they exhibit signs of drowsiness, reducing the risk of accidents caused by fatigue.
• Shift-Based Jobs: The model can be deployed in industries where employees work in shifts, such as healthcare or transportation, to ensure that workers are alert and capable of performing their duties effectively.
• Personal Alert System: Individuals can use the model as a personal safety device to prevent drowsiness-related accidents during activities that require attentiveness, such as studying or operating machinery.

References:

• https://docs.ultralytics.com/
• https://www.youtube.com/@NicholasRenotte