Semantic Segmentation is process to classify each pixel of the image into some class and then in the final image obtained has various classes marked with different colors.
Drone Cam Images from bird view are taken and segmented into 7 different classes. The Kaggle dataset can be downloaded from the this website
The dataset contains 400 images which are of dimensions (3,4000,6000). Out of which 390 images have been taken for training and remaining 10 haven been use for testing the model. The dataset contains the following,
- Images from the Cam
- Label Images(RGB type)
- Colored Masks Images(RGB type)
- CSV file containing color channels values for RGB images
For training the model within available hardware resources, the images are resized to a size of (3,256,256).
Label Title
Num
0) others
1) paved-area
2) vegetation
3) water
4) house
5) person
6) vehicles
For their corresponding RGB colors check color scheme.csv file
The U-net model has been implemented using Pytorch with some minor changes for producing best results. A BatchNorm layer has been introduced in the model for faster training which was not present in the original model since it was discovered a year earlier before BatchNorm.
Also, the number of channels have been reduced to curb with memory issues. The number of channels were reduced to one-fourth of the optimal values obtained in U-net architecture.
Dice Loss function has been used for training the model. The dice loss function along with one-hot encoding rather than using torch.argmax() function since the argmax function is not differentiable.
| Camera Image | Segmented Image |
 |
 |
Google Colab's Tesla T4 GPU has been used to train the model. The model has been trained for 50 epochs.
Credit : Just drop a star as a credit if you use any part of the implementation