In this post, we will walk through how to train Detectron2 to segment your custom objects from any image by providing our model with example training data.
We will follow these steps to train our custom instance segmentation model:
- Assemble a Custom Instance Segmentation Dataset
- Download and Register a Custom Instance Segmentation Dataset
- Configure a Custom Instance Segmentation Training Pipeline
- Run our Custom Instance Segmentation model
- Evaluate Model Performance on Test Imagery
We will be using the public American Sign Language instance segmentation dataset and the open source How to Train Detectron2 Custom Instance Segmentation Notebook.
We recommend having the notebook up alongside this blog post as you work through the tutorial.
What is Instance Segmentation?
If you have followed our tutorials before, you have probably trained a custom object detection model. Object detection draws a bounding box and makes an object class classification around objects of interest.
Instance segmentation takes object detection a step further and draws the outline around the object of interest.
As a result, modeling is slightly more difficult and instance segmentation should only be used when the exact outline of the object is needed for your downstream application.
Assembling A Custom Instance Segmentation Dataset
In our version of this training tutorial, we use the public American Sign Language instance segmentation dataset from Roboflow Universe. You can download this dataset for free with the Roboflow pip package after authenticating with your Roboflow API key.
VERSION = 2 rf = Roboflow(api_key="YOUR API KEY HERE") project = rf.workspace("paul-guerrie-tang1").project("asl-poly") dataset = project.version(VERSION).download("coco")
To apply this tutorial to your own use case, you will need to gather a representative labeled instance segmentation dataset. Labeling instance segmentations with Roboflow Annotate makes this process easy. Select "Instance Segmentation" as your project type when creating a new project in Roboflow and use the "Polygon tool" to draw shapes.
After labeling, you can generate a dataset version and then export your dataset in COCO JSON format to receive a similar pip package code snippet as above, containing your custom objects.
Register Custom Instance Segmentation Dataset
In order to train Detectron2 on our custom dataset, we must first register our dataset in the Detectron2 dataset registry:
We then run a quick test to make sure that our dataset annotations registered properly:
Configure Custom Instance Segmentation Training Pipeline
Now that we have our dataset properly recognized by the Detectron2 library, we can proceed with configuring our custom training pipeline. From the Detectron2 model zoo, we select the
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml model, which is a basic, tried and true, model for making instance segmentations.
Then we define our evaluator and training loop and kick off training.
For more details on the custom training routine, you can refer to the custom instance segmentation training notebook.
Run Custom Instance Segmentation Training
In our example notebook, we kick off training for 10,000 iterations - you can experiment with shortening and lengthening this number to shorten or extend your model's training time. The routine should early stop if it has found a saturated performance model state.
While the network is training, you can see if your network is learning from your dataset by watching the loss metrics converge (lower is better).
You will also see average precision evaluation metrics in aggregate on your dataset.
Evaluate Model Performance on Test Imagery
After training completes, we can load our model in order to test our model's prediction on example imagery.
After invoking our model on test images, we can indeed see that it learned how to segment our custom objects.
Downloading Our Custom Instance Segmentation Model
As a final step, we download the weight artifacts of our custom instance segmentation model for future use in deployment. After rebuilding the Detectron2 library at the deployment location of your choosing, you can load and invoke your model with the same commands we used for test inference above.
from google.colab import files files.download(os.path.join(cfg.OUTPUT_DIR, "model_final.pth"))
Now you know how to train Detectron2 to recognize your custom objects with the fine granularity provided by instance segmentation. If you're looking for more information on training Detectron2, we have an additional guide for you to checkout.
There are alternative options to training instance segmentation models with this method. Roboflow offers an AutoML solution which lets you train a model in one-click.
We hope you enjoyed this tutorial - happy training, and of course, happy inferencing.