How to train Mask RCNN model for custom dataset using google colab

In this blog we will implement mask rcnn model for custom dataset. mask rcnn is a instance Segmentation. First we need dataset. dataset is more important part of artificial intelligence. Mask R-CNN, returns class name and bounding box coordinates for each object,object mask values.

Now come to the point. i have already covered instance segmentation data preparation blogs. first you Need to learned how to annotate images data for instance segmentation. these blog will content all the Details of instance segmentation data preparation. also you can use my sample kangaroo images dataset Its available on my these blog.these blog i have used labelme tools. just invest your 30 minutes on this Blogs and learn how to prepare dataset for mask rcnn Click here .

After finish dataset preparation steps you need to download my project folder on google drive. i have Mentioned all the important folder and python files etc in my project folder also include pretrained mask_rcnn_coco.h5 models.after downloading you need to copy/past your dataset folder in downloaded Project folder. after finished this steps we are ready to train mask rcnn model on custom dataset.

You must be download this folder.

We will used google colab for model training. also you can try on your local machine. just replace the Path of our files and folder i have mentioned all the details in python script. 

After this step you need to upload your project on your google drive. 

Then search google colab and mount the drive and also setup the GPU. its not a rocket science just click Mount drive button and sign up, and for GPU you need to click runtime then click change runtime type Then select GPU.

After this you need to install these given packages on google colab for GPU training.Just copy all Packages and past on colab. if you used your system so please install this packages under the python 3 Virtual environment.

%tensorflow_version 2.x

!pip uninstall -y tensorflow

!pip install tensorflow-gpu==1.14.0

!pip install Keras==2.2.4

!pip install mrcnn

After this steps you need to copy your project on google colab. just create mask_rcnn_train folder and run this script.

cmd :- !mkdir mask_rcnn_train

Source code :

from distutils.dir_util import copy_tree

copy_tree("/content/drive/My Drive/mask_rcnn_custom_train", "/content/mask_rcnn_train/")

Now your drive project folder (which is uploaded) copied in your google colab mask_rcnn_train folder.

After this run

cmd :- !pwd

You seen this type of result 

 /content/

So you need to run this command 

cmd :- ! cd mask_rcnn_train

After this we are ready to train our model on custom dataset 

One more point if you train this model on your local system so you need to change the path of  our logs Folder. these folder will store our train model .so please make sure you need to change.

Also you can change the epochs according to your labels, and also you need to change the class name And number of class name. in this blogs we are using one class name is kangaroo. you can insert your Class name according your dataset.

line number 58 for logs folder location change (these folder will store our train model)

line number 66 for class name change 

line number 221 for epochs change 

Also you can follow this google colab notebook. i have mentioned all the code and necessary command This is the Google colab notebook click here then modified the paths and follow this steps for training.

Click here Colab Notebook

After run this python script on colab or if you try on local so please remove this ! operator 

cmd :- !python3 train.py train --dataset=/content/mask_rcnn_train/dataset --weights=coco

After run this script it take time. in logs folder received model files. we will use last train model files. 

I have already mentioned the image test files. on your project folder you can use it. but you need to Change this lines according your class name model path 

line number 23 for number of classes 

line number 38 for path for new train model (logs folder .h5)

line number 37 for model directory folder (you can used any path for this )

line number 40 for class name 

Then run this script 

cmd :- !python test_images.py --i /content/mask_rcnn_train/input_test_images/00054.jpg --o /content/mask_rcnn_train

I have already mentioned the input_test_images folder in our project folder. you can past your images in Input_test_images folder then replace the path of your images. if you used local system so please Change the path accordingly.

--i : - Path of input images

--o : - Path of output images

After run this script you will receive detected images in your output folder.

FIG (A Result)

FIG (B Result)

You will received this type of result. i know this images not generated mask because i have used only 50 Images. for demo purpose you can used more data and increase the epochs. it will show definitely mask And good result.

After this we will  test our model on video. don't worry i have already mentioned script on folder you Can used it , also i have mentioned .mp4 files for kangaroo.you can test your model on videos according Your class name. just replaces the class name and number of class.

line number 29 for number of classes 

line number 34 for path for new train model (logs folder .h5)

line number 33 for model directory folder (you can used any path for this )

line number 36 for class name 

Then run this script 

cmd :- !python test_videos.py --i /content/mask_rcnn_train/videos_test/kangaroo_videos_test.mp4

--i :- path of input videos

After run this script you will receive the output video output.avi.

Short videos (youtube)

If you need Flask api for this so please visit my blog How to create Flask Machine Learning API for Mask RCNN Detection model

If you have any doubt so please comment 

Thanks