This repo lets you easily fine tune pretrained MaskRCNN model with 64 fast image augmentation types using your custom data/annotations, then apply prediction based on the trained model. Training and inference works both on Windows & Linux.

• torchvision is integrated for MaskRCNN training which provides faster convergence with negative sample support
• albumentations is integrated for image augmentation which is much faster than imgaug and supports 64 augmentation types for images, bounding boxes and masks
• torch-optimizer is integrated to support AdaBound, Lamb, RAdam optimizers.
• tensorboard is integrated for visualizing the training/validation losses, category based training/validation coco ap results and iteration based learning rate changes
• Pretrained resnet50 + feature pyramid weights on COCO is downloaded upon training
• COCO evaluation is performed after each epoch, for training and validation sets for each category

• Manually install Miniconda (Python3) for your OS: https://docs.conda.io/en/latest/miniconda.html

Or install Miniconda (Python3) by bash script on Linux:

sudo apt update --yes
sudo apt upgrade --yes
wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
bash ~/miniconda.sh -b -p ~/miniconda
rm ~/miniconda.sh

• Inside the base project directory, open up a terminal/anaconda promt window and create environment:

conda env create -f environment.yml

• After environment setup, activate environment and run the tests to see if everything is ready:

conda activate augmented-maskrcnn
python -m unittest

• In the base project directory, open up a terminal/anaconda promt window, and activate environment:

conda activate augmented-maskrcnn

• Create a config yml file similar to default_config.yml for your needs

• Perform training by giving the config path as argument:

python train.py configs/default_config.yml

• Visualize realtime training/validation losses and accuracies via tensorboard from http://localhost:6006 in your browser:

tensorboard --logdir=experiments

• Perform prediction for image "CA01_01.tif" using model "artifacts/maskrcnn-best.pt":

python predict.py CA01_01.tif artifacts/maskrcnn-best.pt