- FCN32
- VGG Segnet
- VGG U-Net
You need to make two folders
- Images Folder - For all the training images
- Annotations Folder - For the corresponding ground truth segmentation images
The filenames of the annotation images should be same as the filenames of the RGB images.
The size of the annotation image for the corresponding RGB image should be same.
For each pixel in the RGB image, the class label of that pixel in the annotation image would be the value of the blue pixel.
You need to download the pretrained VGG-16 weights trained on imagenet
cd data
wget "https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_th_dim_ordering_th_kernels.h5"
To train the model run the following command:
python train.py \
--save_weights_path=model \
--train_images="data/dataset1/images_prepped_train/" \
--train_annotations="data/dataset1/annotations_prepped_train/" \
--val_images="data/dataset1/images_prepped_test/" \
--val_annotations="data/dataset1/annotations_prepped_test/" \
--n_classes=10 \
--epochs=5 \
--input_height=320 \
--input_width=640 \
--model_name="vgg_segnet"
Choose model_name from vgg_segnet vgg_unet, fcn32
To get the predictions of a trained model
python predict.py \
--save_weights_path=model \
--test_images="data/dataset1/images_prepped_test/" \
--output_path="data/predictions/" \
--n_classes=10 \
--input_height=320 \
--input_width=640 \
--model_name="vgg_segnet"