Real-Time Object Detection with YOLOv5: A Step-by-Step Guide

Object detection is a fundamental task in computer vision that involves identifying and localizing objects of interest within an image. It has numerous real-world applications, from self-driving cars and surveillance systems to medical imaging and retail analytics.

In recent years, the YOLO (You Only Look Once) family of object detection models has risen to prominence due to its exceptional speed and accuracy. The latest iteration, YOLOv5, represents the state-of-the-art in real-time object detection. Developed by Ultralytics, YOLOv5 outperforms its predecessors in both inference speed and mean average precision (mAP), while also having a significantly smaller model size.

In this guide, we‘ll walk through how to train a custom YOLOv5 model to detect objects of your choosing. Whether you want to count cars in a parking lot, monitor foot traffic in a retail store, or track wildlife in nature footage, these steps will help you build a powerful detection system tailored to your needs. Let‘s get started!

Understanding the YOLOv5 Architecture

At its core, YOLOv5 is a single-stage object detector, meaning it outputs bounding box coordinates and class probabilities directly from the input image in a single pass. This is more efficient than two-stage detectors like Faster R-CNN that first generate region proposals and then classify each proposal separately.

The YOLOv5 architecture consists of three main components:

1. Backbone: A CSPNet (Cross Stage Partial Network) is used to extract rich features from the input image at multiple scales. The backbone is designed to balance speed and accuracy.

2. Neck: A PANet (Path Aggregation Network) aggregates features from different backbone levels to construct a feature pyramid. This allows the model to detect objects at various sizes and scales.

3. Head: Anchor boxes are used to predict the coordinates, objectness scores, and class probabilities at three different scales. The final detections are obtained through non-maximum suppression.

One of the key strengths of YOLOv5 is its ability to run efficiently on edge devices like smartphones and Raspberry Pis. The YOLOv5s "small" model is just 14MB in size and can achieve 50+ frames per second (FPS) on a mobile GPU. Larger models like YOLOv5x offer even higher accuracy while still maintaining real-time speeds.

Setting Up Your Environment

To train and deploy YOLOv5 models, you‘ll need a system with a CUDA-capable GPU and the following software packages installed:

• Python 3.8 or later
• PyTorch 1.7 or later (with GPU support)
• CUDA 10.1 or later

It‘s recommended to create a new virtual environment to avoid conflicts with other projects. You can do this with conda:

conda create --name yolov5 python=3.8
conda activate yolov5

Next, clone the official YOLOv5 repository and install the required dependencies:

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

Preparing Your Dataset

YOLOv5 requires your training data to be organized in a specific directory structure, with images and labels stored separately. For this example, let‘s say we want to detect three classes of objects: cars, pedestrians, and traffic lights.

First, gather a set of images containing the objects you want to detect. The more diverse the images, the better your model will generalize. Aim to collect at least 500-1000 images per class.

Next, you‘ll need to annotate the objects in each image with bounding boxes and class labels. Tools like LabelImg or VGG Image Annotator can help streamline this process. Export your annotations in YOLO format, which consists of a text file for each image with one bounding box per line:

<class_index> <x_center> <y_center> <width> <height>

The class index should be an integer (0=car, 1=pedestrian, 2=traffic light), while the bounding box coordinates are normalized values between 0 and 1.

Finally, split your annotated dataset into train, validation, and test sets, and organize them like this:

dataset/
  train/
    images/
      img001.jpg
      img002.jpg
      ...
    labels/
      img001.txt
      img002.txt
      ...
  val/
    images/
      ...  
    labels/
      ...
  test/
    images/
      ...
    labels/
      ...

Create a YAML file to specify the dataset configuration:

train: dataset/train/
val: dataset/val/

nc: 3
names: [‘car‘, ‘pedestrian‘, ‘traffic light‘]

Training Your YOLOv5 Model

With your dataset ready, you can now train a custom YOLOv5 model. The YOLOv5 repo provides a handy train.py script to launch the training process.

To train with default settings on a single GPU, run:

python train.py --img 640 --batch 16 --epochs 100 --data dataset.yaml --weights yolov5s.pt

This will train a small YOLOv5s model for 100 epochs with an input image size of 640×640 and a batch size of 16. The model weights will be initialized from the pretrained COCO checkpoint.

You can experiment with different model sizes (s/m/l/x) and hyperparameters to find the optimal configuration for your use case. Multi-GPU training is also supported using PyTorch‘s Distributed Data Parallel.

As the model trains, you can monitor the loss curves and evaluation metrics in real-time using Tensorboard:

tensorboard --logdir runs/

Evaluating Performance

After training, it‘s important to evaluate your model‘s performance on a held-out test set. The val.py script computes standard object detection metrics like mean average precision (mAP) at different intersection over union (IoU) thresholds:

python val.py --data dataset.yaml --weights runs/train/exp/weights/best.pt --img 640

In addition to mAP, also pay attention to your model‘s inference speed (measured in FPS) and memory footprint, especially if you plan to deploy on resource-constrained edge devices. You can profile these using the detect.py script.

Deploying Your Model

Once you‘re satisfied with your model‘s performance, you can use it to make predictions on new images or video streams. The YOLOv5 repo provides several inference scripts for different use cases:

• detect.py: Run inference on images, videos, directories, webcam, YouTube, or IP stream
• export.py: Export trained model to ONNX, TensorRT, CoreML, TFLite, or OpenVINO format
• app.py: Launch a simple web app for interactive inference

For example, to detect objects in a video file:

python detect.py --source video.mp4 --weights runs/train/exp/weights/best.pt

The script will output a new video file with bounding boxes and class labels drawn on each frame.

Example Use Case: Vehicle Counting

To illustrate the power of YOLOv5, let‘s walk through an example application of counting vehicles in traffic camera footage. This could be useful for monitoring congestion, analyzing traffic patterns, or enforcing parking regulations.

We‘ll assume you‘ve already trained a YOLOv5 model to detect cars, trucks, and motorcycles in video frames. The next step is to post-process the raw detections to extract vehicle counts and trajectories.

One approach is to use a simple tracking algorithm like centroid tracking. The idea is to associate detections across consecutive frames based on their Euclidean distance. Each time a detection appears sufficiently far from existing centroids, it‘s registered as a new vehicle. Conversely, centroids that haven‘t been updated for several frames are assumed to have left the scene and are removed.

Here‘s a code snippet showing the centroid tracking logic:

import numpy as np
from scipy.spatial import distance as dist

class CentroidTracker:
    def __init__(self, max_disappeared=30):
        self.nextObjectID = 0
        self.objects = OrderedDict()
        self.disappeared = OrderedDict()
        self.max_disappeared = max_disappeared

    def register(self, centroid):
        self.objects[self.nextObjectID] = centroid
        self.disappeared[self.nextObjectID] = 0
        self.nextObjectID += 1

    def update(self, rects):
        if len(rects) == 0:
            for objectID in list(self.disappeared.keys()):
                self.disappeared[objectID] += 1
                if self.disappeared[objectID] > self.max_disappeared:
                    self.deregister(objectID)
            return self.objects

        inputCentroids = np.zeros((len(rects), 2))
        for (i, (x, y, w, h)) in enumerate(rects):
            cX = int(x + w / 2.0)
            cY = int(y + h / 2.0)
            inputCentroids[i] = (cX, cY)

        if len(self.objects) == 0:
            for i in range(0, len(inputCentroids)):
                self.register(inputCentroids[i])
        else:
            objectIDs = list(self.objects.keys())
            objectCentroids = list(self.objects.values())
            D = dist.cdist(np.array(objectCentroids), inputCentroids)
            rows = D.min(axis=1).argsort()
            cols = D.argmin(axis=1)[rows]
            usedRows = set()
            usedCols = set()

            for (row, col) in zip(rows, cols):
                if row in usedRows or col in usedCols:
                    continue
                objectID = objectIDs[row]
                self.objects[objectID] = inputCentroids[col]
                self.disappeared[objectID] = 0
                usedRows.add(row)
                usedCols.add(col)

            unusedRows = set(range(0, D.shape[0])).difference(usedRows)
            unusedCols = set(range(0, D.shape[1])).difference(usedCols)

            if D.shape[0] >= D.shape[1]:
                for row in unusedRows:
                    objectID = objectIDs[row]
                    self.disappeared[objectID] += 1
                    if self.disappeared[objectID] > self.max_disappeared:
                        self.deregister(objectID)
            else:
                for col in unusedCols:
                    self.register(inputCentroids[col])

        return self.objects

To count vehicles, simply instantiate a CentroidTracker and pass it the bounding boxes output by your YOLOv5 model on each frame. The tracker will return a dictionary mapping unique object IDs to centroid coordinates.

For more advanced vehicle tracking, you could explore algorithms like SORT or Deep SORT that incorporate motion prediction and appearance descriptors. Kalman filters are commonly used to smooth out noisy detections.

By keeping a running count of registered objects, you can tally up the total number of vehicles detected over a given time period. You could also analyze the distribution of vehicle types, speeds, and travel directions to gain further insights.

With a well-trained model and some post-processing logic, YOLOv5 enables building sophisticated traffic monitoring systems. Similar techniques can be applied to a wide range of other object detection and tracking applications.

Conclusion

In this guide, we‘ve seen how to harness the cutting-edge YOLOv5 architecture for custom object detection tasks. By annotating your own dataset and leveraging transfer learning, you can quickly build highly accurate models tailored to your specific use case.

The YOLOv5 framework streamlines training and deployment with a batteries-included approach. Out of the box, it supports distributed training, mixed-precision, pruning, hyperparameter evolution, exporting to edge formats, and more. A vibrant community has also developed extensions for tasks like instance segmentation, pose estimation, and GANs.

Whether you‘re a beginner or an experienced practitioner, YOLOv5 is a powerful tool to have in your computer vision toolkit. Its exceptional performance and ease of use have made it a go-to for applications ranging from industrial defect detection to wildlife conservation.

I encourage you to experiment with YOLOv5 on your own datasets and share your findings with the community. Ultralytics and other researchers continue to innovate on the YOLO architecture, so be sure to keep up with the latest releases and papers. Happy detecting!