How to Train a YOLOv5 Oriented Bounding Box Model

In 2024, Ultralytics released YOLOv8 Oriented Bounding Boxes which outperforms YOLOv5.

Oriented Bounding Boxes

Oriented bounding boxes are bounding boxes rotated to better fit the objects represented on an angle. Take a pill detection dataset for example. Using YOLOv5-obb we are able to detect pills that are rotated on a given frame or image more tightly and accurately, preventing capture of multiple pills or other objects in one bounding box.

How To Use Oriented Object Detection with YOLOv5

In this tutorial, we’ll go over how to use YOLOv5-obb by exporting an already annotated dataset from Roboflow.

You can follow along using your own dataset or the polygon annotated American Sign Language dataset found in the notebook that we’ll be using here. For more computer vision datasets, check out Roboflow Universe.

To follow along, open this google colab notebook in a new tab.

Following the installation instructions on the YOLOv5-OBB repo, the first line of code we'll run in our notebook after confirming version requirements is:

pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
nvcc -V
python

Import torch and confirm CUDA runtime api version is at least 10.0 when prompted.

>>> import torch
>>> torch.version.cuda
>>> exit()

Clone the YOLOv5-OBB repo

!git clone https://github.com/hukaixuan19970627/yolov5_obb.git

Install the requirements

%cd /content/yolov5_obb/
!pip install -r requirements.txt

Run setup.py develop in the utils/nms_rotated folder

%cd utils/nms_rotated
!python setup.py develop  #or "pip install -v -e ."

Export Your Dataset for YOLOv5-OBB

Next we'll make a directory called datasets and cd into it in order to run our !pip install roboflow code snippet.

Below we are using the the instance segmentation American Sign Language dataset found in the notebook. You can substitute this for your own dataset information by exporting your dataset in YOLOv5 Oriented Bounding Boxes format.

Remember to export your annotated dataset by following the steps in the above video. Use your "show download code" snippet to replace the one below.

%cd /content/datasets/

!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="YOUR_API_KEY")
project = rf.workspace("YOUR_WORKSPACE_NAME").project("YOUR_PROJECT_NAME")
dataset = project.version(4).download("yolov5-obb")

How to Train YOLOv5-OBB

After running the code above the expected dataset name when running !python train.py will be roboflow, so be sure to rename your dataset folder to roboflow to prevent any errors from triggering.

In other words, make sure your dataset structure is:

├── yolov5
└── datasets
    └── roboflow
        ├── test
        ├── train
        └── valid
            ├── images
                 |────1.jpg
                 |────...
                 └────10000.jpg
            ├── labelTxt
                 |────1.txt
                 |────...
                 └────10000.txt
        └────README.roboflow.txt
        └────data.yaml
       

%cd /content/yolov5_obb/
!python train.py --data /content/datasets/roboflow/data.yaml --epochs 10 --batch-size 1 --img 1024 --device 0 --exist-ok

Detect with YOLOv5-OBB

After training is complete, run detect.py using the code below.

!python detect.py --weights 'runs/train/exp/weights/best.pt' --source '../datasets/roboflow/valid/images' --conf-thres 0.01 --img 416 --device 0 --agnostic

View an Image with YOLOv5-OBB

After detect.py completes, choose some random images to view the Oriented Bounding Boxes detection on a given image from the exp folder.

import os, random
random_file = random.choice(os.listdir("/content/yolov5_obb/runs/detect/exp"))

from IPython.display import Image
Image(random_file)

Try More Datasets

Visit https://universe.roboflow.com/ to try more datasets to train using YOLOv5's Oriented Object Detection!