This repository is an open source tool to build and train image segmentation models on custom data with the ability to inference existing models!
Begin by creating a conda environment and activating it,
conda create -n <ENV NAME> python=3.10; conda activate <ENV NAME>
and install all of the requirements,
python -m pip install -r /path/to/requirements.txt
The scripts will work fine provided you have all the installed requirements. For the alliance canada HPC cluster, the following commands executed in this order allows for proper environment set-up,
virtualenv --no-download ENV; source ENV/bin/activate
then load up the requirements,
module load python/3.10;
python -m pip install -r requirements.txt;
module load scipy-stack;
module spider gcc cuda opencv/4.9.0
Organize your dataset in the project directory so that it has the following structure,
├── data
│ ├── train
│ │ ├── <train0>.jpg
│ │ ├── ...
│ │ ├── <trainX>.jpg
│ │ ├── annotations.json
│ ├── val
│ │ ├── <val0>.jpg
│ │ ├── ...
│ │ ├── <valY>.jpg
│ │ ├── annotations.json
│ ├── test
│ │ ├── <test0>.jpg
│ │ ├── ...
│ │ ├── <testZ>.jpg
│ │ ├── annotations.json
Note: You will have to generate the annotations yourself; this can be done easily through an online annotation tool such as VGG Image Annotator and exporting it such that it follows standard formatting,
In the end, each of your annotations.json
file should look like,
{
"<train0>.jpg": {
"filename": "<image0>.jpg",
"size": 1234,
"regions": [
{
"shape_attributes": {
"name": "polygon",
"all_points_x": [
119,
94,
200,
200
],
"all_points_y": [
114,
197,
199,
90
]
},
"region_attributes": {
"label": "<image0: region 1 label>"
}
},
... more regions ...
],
},
... more images ...
}
In addition to your dataset, you need a configuration file that tells detectron2 how to build and train your Mask-RCNN model. See the example below from ./examples/dot_configuration/configuration.yaml.
info:
name: quantum dot configuration
# Need to specify class dictionary for custom datasets,
# 'nd' -> 'no dot'
# 'ld' -> 'left dot'
# 'cd' -> 'central dot'
# 'rd' -> 'right dot'
# 'dd' -> 'double dot'
# Dictionary keys MUST match your annotations.json labels
# Dictionary values can be anything as long as they are unique
classes: {"nd": 0, "ld": 1, "cd": 2, "rd": 3, "dd": 4}
hyperparameters:
learning_rate: 2.0e-4
num_epochs: 7000
batch_num: 10
batch_size_per_img: 128
To begin training your model on your data/train
folder, simply execute the following command,
python train.py --config <configuration>.yaml --data_dir <path/to/data>
Note: Use
--help
to see all of the flags available to you.
In order to validate your model's performance, you can inference the model on the data/val
folder. It is recommended to first do it interactively using the src/inference.ipynb
notebook.
With your model validated, you can inference the model on the data/test
folder. It is recommended to first do it interactively using the src/inference.ipynb
notebook.
Before running inference, you need to provide two files along with your image of interest,
-
<model>.pth
: This file contains the trained model weights. -
configuration.yaml
: This file contains the exact configuration settings that were used for the trained model you are trying to inference (see Training section). -
<image>.jpg
: This file is the image you wish to inference
To begin training simply execute the following command,
python inference.py --model_weights <model>.pth --config <configuration>.yaml --image <image>.jpg
Note: Use
--help
to see all of the flags available to you.