Image Background Remover Using Detectron2
Detectron2 is a ground-up rewrite of Detectron that started with maskrcnn-benchmark. The platform is now implemented in PyTorch. With a new, more modular design, Detectron2 is flexible and extensible, and able to provide fast training on single or multiple GPU servers. Detectron2 includes high-quality implementations of state-of-the-art object detection algorithms, including DensePose, panoptic feature Pyramid networks, and numerous variants of the pioneering Mask-R-CNN model family also developed by FAIR. Its extensible design makes it easy to implement cutting-edge research projects without having to fork the entire codebase. for more information you can visit here.
In this blog, I have shown installation and working of detectron2.
System Requirements
- Linux or macOS with Python ≥ 3.6
- PyTorch ≥ 1.5 and torchvision that matches the PyTorch installation. You can install them together at pytorch.org to make sure of this
- OpenCV is optional and needed by demo and visualization
we can also use the Google colab to use detectron2. In this blog I have shown the implementation in google colab.
Installation of Detectron2
In Colab Notebook, just run below code to install detectron2 and other pre-requisites.
Dataset
I have used the COCO dataset for the training of the model. COCO dataset is the dataset used for object detection purpose and it has 80 classes ,80000 training images and 40000 validation images. Below is the code for the downloading the dataset.
Import libraries
Registering dataset
In this step I have registered the dataset using inbuilt detectron2 function that is register_coco_instances.
below i have shown the code which you need to run in colab notebook.
Train the model
In this phase, I trained dataset using detectron2 default trainer and make the the model for predicting the object and making the mask over the object. I have trained the model for 300 iterations.
Below is the code for plotting the trained models some figures.
I have shown the graph for accuracy of model which shows the plot on accuracy / iteration using mask-RCNN.
Validating the model
First I read the image using OpenCV and then predicted the object from image using default predictor. Then I have created the mask from the predicted output on image and Then remove the background of object by checking the value of every pixel and if the value of object at pixel is 0 that means pixel is not of an object.
it will read the image and plot it as shown below.
just run the below code in colab notebook. It will produce the image from which background has been removed.
here is the removed background image.
So Detectron2 make it simple to train segmentation model for custom dataset. You can also visit the detectron2 GitHub repository.
