Custom Instance Segmentation Training With 7 Lines Of Code. | by Ayoola Olafenwa (she/her) | Towards Data Science

Image Segmentation is an important field in computer vision, it is applied in different fields of life. PixelLib is a library created to allow easy application of segmentation to real life problems. It supports instance segmentation of objects with Coco model. Segmentation with coco model is limited as you cannot perform segmentation beyond the 80 classes available in coco. It is now possible to train your custom objects’ segmentation model with PixelLib Library with just 7 Lines of Code.
Install PixelLib and its dependencies:
Install Tensorflow with:(PixelLib supports tensorflow 2.0 and above)
pip3 install tensorflow
Install imgaug with:
pip3 install imgaug
Install PixelLib with
pip3 install pixellib
If installed upgrade to the latest version using:
pip3 install pixellib — upgrade
STEP1:
Prepare your dataset:
Our goal is to create a model that can perform instance segmentation and object detection on butterflies and squirrels.
Collect images for the objects you want to detect and annotate your dataset for custom training. Labelme is the tool employed to perform polygon annotation of objects. Create a root directory or folder and within it create train and test folder. Separate the images required for training (a minimum of 300) and test. Put the images you want to use for training in the train folder and put the images you want to use for testing in the test folder. You will annotate both images in the train and test folder. Download Nature’s dataset used as a sample dataset in this article, unzip it to extract the images’’ folder. This dataset will serve as a guide for you to know how to organize your images. Ensure that the format of your dataset’s directory is not different from it. Nature is a dataset with two classes butterfly and squirrel. There are 300 images for each class for training and 100 images for each class for testing i.e 600 images for training and 200 images for validation. Nature is a dataset with 800 images.
Read this article to learn how to use labelme for annotating objects.
medium.com
Nature >>train>>>>>>>>>>>> image1.jpg                           image1.json                           image2.jpg                           image2.json             >>test>>>>>>>>>>>>  img1.jpg                           img1.json                           img2.jpg                           img2.json   
Sample of folder directory after annotation.
Note: PixelLib supports annotation with Labelme. If you make use of another annotation tool it will not be compatible with the library.
STEP 2:
Visualize Dataset
Visualize a sample image before training to confirm that the masks and bounding boxes are well generated.
import pixellibfrom pixellib.custom_train import instance_custom_trainingvis_img = instance_custom_training()
We imported pixellib, from pixellib we imported the class instance_custom_training and created an instance of the class.
vis_img.load_dataset("Nature”)
We loaded the dataset using load_dataset function. PixelLib requires polygon annotations to be in coco format. When you call the load_dataset function, the individual json files in the train and test folder will be converted into a single train.json and test.json respectively. The train and test json files will be located in the root directory as the train and test folder. The new folder directory will now look like this:
Nature >>>>>>>>train>>>>>>>>>>>>>>> image1.jpg               train.json           image1.json                                    image2.jpg                                    image2.json       >>>>>>>>>>>test>>>>>>>>>>>>>>> img1.jpg                  test.json           img1.json                                      img2.jpg                                      img2.json
Inside the load_dataset function annotations are extracted from the jsons’s files. Masks are generated from the polygon points of the annotations and bounding boxes are generated from the masks. The smallest box that encapsulates all the pixels of a mask is used as a bounding box.
vis_img.visualize_sample()
When you call this function, it shows a sample image with a mask and bounding box.
Great! the dataset is fit for training, the load_dataset function successfully generates mask and bounding box for each object in the image. Random colors are generated for the masks in HSV space and then converted to RGB.
FINAL STEP:
Train a custom model using your dataset
This is the code for performing training. In just seven lines of code, you train your dataset.
train_maskrcnn.modelConfig(network_backbone = "resnet101", num_classes= 2, batch_size = 4)                       
We called the function modelConfig, i.e model’s configuration. It takes the following parameters:
network_backbone: This is the CNN network used as a feature extractor for mask-rcnn. The feature extractor used is resnet101.
num_classes: We set the number of classes to the categories of objects in the dataset. In this case, we have two classes(butterfly and squirrel) in nature’s dataset.
batch_size: This is the batch size for training the model. It is set to 4.
train_maskrcnn.load_pretrained_model("mask_rcnn_coco.h5")train_maskrcnn.load_dataset(“Nature”)
We are going to employ the technique of transfer learning for training the model. Coco model has been trained on 8O categories of objects, it has learnt a lot of features that will help in training the model. We called the function load_pretrained_model function to load the mask-rcnn coco model. We loaded the dataset using load_dataset function.
Download coco model from here.
train_maskrcnn.train_model(num_epochs = 300, augmentation=True,path_trained_models = “mask_rcnn_models”)
Finally, we called the train function for training mask r-cnn model. We called train_model function. The function took the following parameters:
num_epochs:The number of epochs required for training the model. It is set to 300.
augmentation: Data augmentation is applied on the dataset. This is because we want the model to learn different representations of the objects.
path_trained_models: This is the path to save the trained models during training. Models with the lowest validation losses are saved.
Using resnet101 as network backbone For Mask R-CNN modelTrain 600 images Validate 200 images Applying augmentation on dataset Checkpoint Path: mask_rcnn_modelsSelecting layers to trainEpoch 1/200100/100 - 164s - loss: 2.2184 - rpn_class_loss: 0.0174 - rpn_bbox_loss: 0.8019 - mrcnn_class_loss: 0.1655 - mrcnn_bbox_loss: 0.7274 - mrcnn_mask_loss: 0.5062 - val_loss: 2.5806 - val_rpn_class_loss: 0.0221 - val_rpn_bbox_loss: 1.4358 - val_mrcnn_class_loss: 0.1574 - val_mrcnn_bbox_loss: 0.6080 - val_mrcnn_mask_loss: 0.3572Epoch 2/200100/100 - 150s - loss: 1.4641 - rpn_class_loss: 0.0126 - rpn_bbox_loss: 0.5438 - mrcnn_class_loss: 0.1510 - mrcnn_bbox_loss: 0.4177 - mrcnn_mask_loss: 0.3390 - val_loss: 1.2217 - val_rpn_class_loss: 0.0115 - val_rpn_bbox_loss: 0.4896 - val_mrcnn_class_loss: 0.1542 - val_mrcnn_bbox_loss: 0.3111 - val_mrcnn_mask_loss: 0.2554Epoch 3/200100/100 - 145s - loss: 1.0980 - rpn_class_loss: 0.0118 - rpn_bbox_loss: 0.4122 - mrcnn_class_loss: 0.1271 - mrcnn_bbox_loss: 0.2860 - mrcnn_mask_loss: 0.2609 - val_loss: 1.0708 - val_rpn_class_loss: 0.0149 - val_rpn_bbox_loss: 0.3645 - val_mrcnn_class_loss: 0.1360 - val_mrcnn_bbox_loss: 0.3059 - val_mrcnn_mask_loss: 0.2493
This is the training log. It shows the network backbone used for training mask-rcnn which is resnet101, the number of images used for training and number of images used for validation. In the path_to_trained models’ directory, the models are saved based on decrease in validation loss, and typical model name will appear like this: mask_rcnn_model_25–0.55678, it is saved with its epoch number and its corresponding validation loss.
Network Backbones:
There are two network backbones for training mask-rcnn
Resnet101
Resnet50
Google colab: Google Colab provides a single 12GB NVIDIA Tesla K80 GPU that can be used up to 12 hours continuously.
Using Resnet101: Training Mask-RCNN consumes a lot of memory. On google colab using resnet101 as network backbone, you will be able to train with a batchsize of 4. The default network backbone is resnet101. Resnet101 is used as a default backbone because it appears to reach a lower validation loss during training faster than resnet50. It also works better for a dataset with multiple classes and much more images.
Using Resnet50: The advantage with resnet50 is that it consumes lesser memory. Thus, you can use a batch_size of 6 or 8 on google colab depending on how colab randomly allocates GPU.
The modified code supporting resnet50 will be like this.
Full code
The main differences from the original code is that in the model’s configuration function, we set network_backbone to be resnet50 and changed the batch size to 6.
The only difference in the training log is this:
Using resnet50 as network backbone For Mask R-CNN model
It shows that we are using resnet50 for training.
Note: The batch_sizes given are samples used for google colab. If you are using a less powerful GPU, reduce your batch size. For example, for a PC with a 4G RAM GPU, you should use a batch size of 1 for both resnet50 or resnet101. I used a batch size of 1 to train my model on my PC’s GPU, trained for less than 100 epochs and it produced a validation loss of 0.263. This is favourable because my dataset is not large. A PC with a more powerful GPU should be able to use a batch size of 2. If you have a large dataset with more classes and much more images, you can use google colab where you have free access to GPU. Most importantly, try to use a more powerful GPU and train for more epochs to produce a custom model that will perform efficiently across multiple classes.
Achieve better results by training with much more images. 300 images for each each class is recommended to be the minimum requirement for training.
Model Evaluation
When we are done with training, we should evaluate models with lowest validation losses. Model evaluation is used to access the performance of the trained model on the test dataset. Download the trained model from here.
output
mask_rcnn_models/Nature_model_resnet101.h5 evaluation using iou_threshold 0.5 is 0.890000
The mAP(Mean Avearge Precision) of the model is 0.89.
You can evaluate multiple models at once, what you just need is to pass in the folder directory of the models.
Output log
mask_rcnn_models\Nature_model_resnet101.h5 evaluation using iou_threshold 0.5 is 0.890000mask_rcnn_models\mask_rcnn_model_055.h5 evaluation using iou_threshold 0.5 is 0.867500mask_rcnn_models\mask_rcnn_model_058.h5 evaluation using iou_threshold 0.5 is 0.8507500
Evaluation with Resnet50
Note: Change the network_backbone to resnet50 if you are evaluating a resnet50 model.
Inference with a custom model
We have trained and evaluated the model. The next step is to see the performance of the model on unknown images.
We are going to test the model on the classes we have trained it on.
Sample1.jpg
import pixellibfrom pixellib.instance import custom_segmentation segment_image =custom_segmentation()segment_image.inferConfig(num_classes= 2, class_names= ["BG", "butterfly", "squirrel"])
We imported the class custom_segmentation, the class for performing inference and created an instance of the class. We called the model configuration’s function and introduced an extra parameter class_names.
class_names= ["BG", "butterfly", "squirrel"])
class_names:It is a list containing the names of classes the model is trained with. “BG” refers to the background of the image. It is the first class and must be available along the names of the classes.
Note: If you have multiple classes and you are confused of how to arrange the classes’s names according to their class ids, in your test.json in the dataset’s folder, check the categories’ list.
{"images": [{"height": 205,"width": 246,"id": 1,"file_name": "C:\\Users\\olafe\\Documents\\Ayoola\\PIXELLIB\\Final\\Nature\\test\\butterfly (1).png"},],"categories": [{"supercategory": "butterfly","id": 1,"name": "butterfly"},{"supercategory": "squirrel","id": 2,"name": "squirrel"}],
You can observe from the sample of the directory of test.json above that after the images’ list in the test.json is object categories’ list. The classes’ names are there with their corresponding class ids. Butterfly has a class id 1 and squirrel has a class id 2. Remember, the first id “0” is kept in reserve for the background.
segment_image.load_model("mask_rcnn_model/Nature_model_resnet101.h5)segment_image.segmentImage("sample1.jpg", show_bboxes=True, output_image_name="sample_out.jpg")
The custom model is loaded and we called the function to segment the image.
When parameter show_bboxes is set to True, we will be able to perform instance segmentation with object detection.
test_maskrcnn.segmentImage(“sample1.jpg”,show_bboxes = False, output_image_name=”sample_out.jpg”)
When show_bboxes is set to False, we obtain only segmentation masks.
Sample2.jpg
test_maskrcnn.segmentImage(“sample2.jpg”,show_bboxes = True, output_image_name=”sample_out.jpg”)
WOW! We have successfully trained a custom model for performing instance segmentation and object detection on butterflies and squirrels.
sample_video1
We want to perform segmentation on the butterflies in this video.
test_video.process_video("video.mp4", show_bboxes = True,  output_video_name="video_out.mp4", frames_per_second=15)
The function process_video is called to perform segmentation on objects in a video.
It takes the following parameters:-
video_path: This is the path to the video file we want to segment.
frames_per_second: This is the parameter used to set the number of frames per second for the saved video file. In this case, it is set to 15 i.e the saved video file will have 15 frames per second.
output_video_name: This is the name of the saved segmented video. The output video will be saved in your current working directory.
Output_video
A sample of another segmented video with our custom model.
You can perform live camera segmentation with your custom model making use of this code:
You will replace the process_video funtion with process_camera function.In the function, we replaced the video’s filepath to capture i.e we are processing a stream of frames captured by the camera instead of a video file. We added extra parameters for the purpose of showing the camera frames:
show_frames: This parameter handles the showing of segmented camera’s frames.
frame_name: This is the name given to the shown camera’s frame.
Visit Google colab’s notebook set up for training a custom dataset.
colab.research.google.com
Visit Pixellib’s Official Github Repository
Visit Pixellib’s Official Documentation
Reach to me via:
Email: olafenwaayoola@gmail.com
Linkedin: Ayoola Olafenwa
Twitter: @AyoolaOlafenwa
Facebook: Ayoola Olafenwa