This project implements a road segmentation model using a U-Net architecture with an EfficientNet-B0 backbone to detect and segment roads from aerial imagery. The project uses the Massachusetts Roads Dataset subset and demonstrates a complete deep learning workflow for image segmentation.
- Source: Massachusetts Roads Dataset (subset)
- Images: 200 aerial images
- Image Size: 1500×1500 pixels
- Coverage: 2.25 square kilometers per image
- Task: Binary segmentation of road networks
- PyTorch
- segmentation-models-pytorch
- albumentations
- OpenCV
- NumPy
- Pandas
- scikit-learn
- matplotlib
-
Data Preparation
- Custom dataset loader
- Train/validation split
- Data augmentation using Albumentations
-
Model Architecture
- U-Net with EfficientNet-B0 backbone
- Binary segmentation
- Uses pre-trained ImageNet weights
-
Training Pipeline
- Adam optimizer
- Dice Loss + Binary Cross Entropy
- 25 training epochs
- Model checkpoint saving
- Learning Rate: 0.003
- Batch Size: 8
- Image Size: 512×512
- Encoder: timm-efficientnet-b0
- Encoder Weights: ImageNet
pip install segmentation-models-pytorch
pip install -U git+https://github.com/albumentations-team/albumentations
pip install --upgrade opencv-contrib-python- Clone the repository
- Prepare your dataset
- Adjust hyperparameters in the script
- Run the training script
- Load the trained model
- Use the
helper.show_image()function to visualize predictions
- Visualizes road segmentation masks
- Demonstrates image augmentation techniques
- Provides a complete deep learning workflow
- Train on full Massachusetts Roads Dataset
- Experiment with different encoders and architectures
- Add more advanced augmentation techniques
- Segmentation Models PyTorch
- Albumentations Documentation
- Massachusetts Roads Dataset
MIT License
Feel free to use this implementation and modify it according to your needs. Contributions are welcome!