YOLOv5 is usually associated with object detection and is one of the most popular networks in the world for that task. Recently, image classification was added to YOLOv5, and it keeps getting better!

As of September 2022, YOLOv5 also supports instance segmentation tasks.

This blog will walk through how to train YOLOv5 for instance segmentation on a custom dataset. If you're interested other tasks, we have guides for those:

Follow along with the How To Train YOLOv5 Instance Segmentation Colab Notebook

This post will walk through:

  • Setting up the environment
  • Prepare a custom dataset for instance segmentation
  • Training YOLOv5 for instance segmentation
  • Testing and validating the custom model

Setting Up The Environment for YOLOv5 Segmentation

First, we need to clone the Ultralytics YOLOv5 repo and install all its dependancies.

!git clone https://github.com/ultralytics/yolov5  # clone
%cd yolov5
%pip install -qr requirements.txt  # install

import torch
import utils
display = utils.notebook_init()  # checks

Prepare a Custom Dataset for Instance Segmentation

In order to train YOLOv5 with a custom dataset, you'll need to gather a dataset, label the data, and export the data in the proper format for YOLOv5 to understand your annotated data.

Roboflow Annotate makes each of these steps easy and is the tool we will use in this tutorial.

To get started, create a free Roboflow account. We will be using this American Sign Language dataset from Roboflow Universe as our example dataset. An American Sign Language (ASL) model could be used to automatically interpret ASL.

Instance segmentation can be used in many of the same applications as object detection with the added bonus of getting a detailed polygon for each detected object, instead of just a bounding box.

This can be useful when more detail is required about each detection, like the object's area or shape. Instance segmentation can be helpful in many applications including agricultural, OCR, insurance, and others.

Feel free to follow along with the same dataset or find another dataset in Universe (a community of 66M+ open source images) to use if you don't already have your own data.

If you have your own images ready to go from another source, here's what we'll do.

Create a new Project in Roboflow and select Instance Segmentation.

Next, add the data to your Project via API or through our web interface. If you're using a dataset from Roboflow Universe, you can download the data with the annotations already done for you which means you do not need to manually label your images.

If you drag and drop a directory with a dataset in a supported format, Roboflow will automatically read the images and annotations together. If you only have images, you can label them in Roboflow Annotate. Add your images and assign them to a user in your account to label them (even if this assigning to yourself).

When labeling for instance segmentation jobs, it's important to use polygon annotations, since we want the model to learn the shape of each object.

After labeling the data, we can apply preprocessing and augmentation to increase the size of our dataset and account for cases that may give our model difficulty predicting the object.

Preprocessing options ensure our data is formatted consistently for model training.

For this tutorial, let's use Auto-Orient and Resize because the photo may have been taken in portrait or landscape and we have different image sizes in our dataset.

We can apply augmentations that could be possible in real world scenarios. Flipping, rotating, and shear will help when our objects can be in multiple orientations so we will apply those. Blur, noise, and cutout can help when anything is blocking the full view of our object or if other objects might enter our images.

Once we've applied our preprocessing and augmentations, we can generate our dataset which is the final step in creating a dataset to train the model.

After clicking generate, our dataset will be prepared and can then be exported for training in our custom YOLOv5 instance segmentation notebook pipeline.

We can train with Roboflow Train to enable model assisted labeling and test our model. 

You have multiple options for exporting the dataset and for this tutorial you'll select 'show download code' which you'll use to then access this dataset in the Colab notebook.

Be sure to copy this code snippet that Roboflow generates for us.

Return to the YOLOv5 Instance Segmentation Colab Notebook and paste this generated code snippet. Be sure to paste the code snippet  (which includes our Workspace, Project name, and private API key) to below %cd ../datasets/ so that the data downloads to the right place.

!pip install roboflow

%cd ../datasets/

from roboflow import Roboflow
rf = Roboflow(api_key="YOUR API KEY")
project = rf.workspace("workspace-name").project("project-name")
dataset = project.version(1).download("yolov5")

Train YOLOv5 For Instance Segmentation on a Custom Dataset

Before training on our own custom dataset, we'll need do download some pre-trained weights so we aren't starting from scratch. This will accelerate the training process and improve the quality of our results.

%cd ../yolov5
from utils.downloads import attempt_download

p5 = ['n', 's', 'm', 'l', 'x']  # P5 models
cls = [f'{x}-seg' for x in p5]  # segmentation models

for x in cls:
    attempt_download(f'weights/yolov5{x}.pt')

Then, we are ready to train! We can follow the same steps as above with some minor modifications to the input parameters.

!python segment/train.py --img 320 --batch 16 --epochs 100 --data {dataset.location}/data.yaml --weights yolov5s-seg.pt
display.Image(filename=f'runs/train-seg/exp/results.png', width=1200)

Here's the (truncated) output to expect:

Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
       0/99      4.69G    0.09511     0.0684    0.01977    0.08956         30   
                 Class     Images  Instances      Box(P          R      mAP50  m
                   all        334        334     0.0129      0.648     0.0367     0.0114     0.0146      0.693      0.047     0.0291

      Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
       1/99      4.76G    0.05718    0.02761    0.02013    0.08223         35   
                 Class     Images  Instances      Box(P          R      mAP50  m
                   all        334        334     0.0038      0.987     0.0576     0.0204    0.00377       0.98     0.0644     0.0457

      Epoch    GPU_mem   box_loss   seg_loss   obj_loss   cls_loss  Instances       Size
       2/99      4.76G    0.05415    0.02485    0.01717    0.08046         30   
                 Class     Images  Instances      Box(P          R      mAP50  m
                   all        334        334     0.0283      0.403     0.0459     0.0154      0.031      0.445     0.0486     0.0302

Test and Validate the Custom Model

We can test and validate with our newly trained custom model. This script reports the results of the instance segmentation training.

!python segment/val.py --weights runs/train-seg/exp/weights/best.pt --data {dataset.location}/data.yaml --img 320

Next, we can try out inferring with our custom model. We'll pass an example image. In the notebook, note that we grab one of the images that comes in the test set folder that comes from the Roboflow data loader. This ensures that we're truly testing on an image our model did not see in training.

!python segment/predict.py --img 320 --weights runs/train-seg/exp/weights/best.pt --source {example_image_path}
display.Image(filename=f'runs/predict-seg/exp2/{example_image_name}', width=600)

We can see the output of the model overlayed on the test image above, so it looks like our model is beginning to learn, though it can likely benefit from more data and active learning.

Improving the YOLOv5 Segmentation Model with Active Learning

Once you have trained the model, you may want to improve the results by adding additional data going through your computer vision pipeline. Use the Roboflow pip package to make the most of your trained model.

Using the pip package you can download or export images from your dataset, upload images and annotations, run inference on a trained version of your dataset, and improve your model's performance by adding more data to train the model again.

from roboflow import Roboflow
import json

# private api key found in Roboflow > YOURWORKSPACE > Roboflow API
# NOTE: this is your private key, not publishable key!
# https://docs.roboflow.com/rest-api#obtaining-your-api-key
private_api_key = "INSERT API KEY HERE"

# gain access to your workspace
rf = Roboflow(api_key=private_api_key)
workspace = rf.workspace()

# you can obtain your model path from your project URL, it is located
# after the name of the workspace within the URL - you can also find your
# model path on the Example Web App for any dataset version trained
# with Roboflow Train
# https://docs.roboflow.com/inference/hosted-api#obtaining-your-model-endpoint
model_path = "INSERT MODEL PATH HERE"
project = workspace.project(f"{model_path}")

# be sure to replace with a path to your file
# if you run this in Colab, be sure to upload the file to colab, hover over
# the file name, and select the 3 dots on the right of the file name to copy
# the file path, and paste it as the value for "imgpath"
img_path = "INSERT IMAGE PATH HERE"

# establish number of retries to use in case of upload failure
project.upload(f"{img_path}", num_retry_uploads=3)

Below you'll methods to upload images that can help select the right data to then improve the performance of your model. These are out of the box options for you, but feel free to create custom logic that makes the most sense for your domain.

By increasing the available data to train with and focusing on data your model is not handling well, you can improve model performance quickly while in production. We recommend deploying with containers to help make the process of updating your model easy to do.

from roboflow import Roboflow
# obtaining your API key: https://docs.roboflow.com/rest-api#obtaining-your-api-key
rf = Roboflow(api_key="INSERT_PRIVATE_API_KEY")
workspace = rf.workspace()


raw_data_location = "INSERT_PATH_TO_IMAGES"
raw_data_extension = ".jpg" # or ".png", ".jpeg" depending on file type

# replace * with your model version number for inference
inference_endpoint = ["INSERT_MODEL_ID", *]
upload_destination = "INSERT_MODEL_ID"
# set the conditionals values as necessary for your active learning needs
conditionals = {
    "required_objects_count" : 1,
    "required_class_count": 1,
    "target_classes": [],
    "minimum_size_requirement" : float('-inf'),
    "maximum_size_requirement" : float('inf'),
    "confidence_interval" : [10,90],
}

# filtering out images for upload by similarity is available for paid plans
# contact the Roboflow team for access: https://roboflow.com/sales
# conditionals = {
#     "required_objects_count" : 1,
#     "required_class_count": 1,
#     "target_classes": [],
#     "minimum_size_requirement" : float('-inf'),
#     "maximum_size_requirement" : float('inf'),
#     "confidence_interval" : [10,90],
#     "similarity_confidence_threshold": .3,
#     "similarity_timeout_limit": 3
# }

workspace.active_learning(raw_data_location=raw_data_location, 
    raw_data_extension=raw_data_extension,
    inference_endpoint=inference_endpoint,
    upload_destination=upload_destination,
    conditionals=conditionals)

Start Running Inference with YOLOv5 Instance Segmentation

You trained YOLOv5 for Instance Segmentation!

Don't forget to checkout the How To Train YOLOv5 Instance Segmentation Colab.

Happy segmenting!