Road Segmentation Using U-Net

Overview

This project implements semantic road segmentation using a U-Net deep learning model with a ResNet34 encoder. The model performs pixel-level classification to identify road regions from input images.

The project demonstrates a complete computer vision workflow:

• Dataset preparation
• Image preprocessing
• Data augmentation
• Model training
• Evaluation
• Inference
• Prediction visualization

---

Dataset

The project uses a road segmentation dataset containing paired:

• Road images
• Binary segmentation masks

Images were resized and augmented using Albumentations to improve model generalization.

Dataset structure:

dataset/
├── train/
│   ├── images/
│   └── masks/
│
└── val/
    ├── images/
    └── masks/

---

Model

Architecture

• Model: U-Net
• Encoder: ResNet34
• Framework: PyTorch
• Segmentation Library: segmentation_models_pytorch
• Loss Function: Dice Loss

The model outputs a binary segmentation mask identifying road pixels.

---

Features

• U-Net semantic segmentation model
• ResNet34 pretrained encoder
• Dice Loss optimization
• Albumentations data augmentation
• Automatic best-model checkpoint saving
• Binary road mask generation
• Prediction overlays
• Visualization galleries
• Evaluation metrics calculation

---

Project Structure

road_segmentation/

├── checkpoints/
│   └── best_unet.pth
│
├── dataset/
│   ├── train/
│   └── val/
│
├── outputs/
│   ├── comparisons/
│   ├── predictions/
│   ├── masks/
│   ├── comparison_gallery.png
│   └── prediction_gallery.png
│
├── reports/
│   └── metrics.txt
│
├── scripts/
│   ├── model.py
│   ├── train_unet.py
│   ├── predict.py
│   ├── evaluate.py
│   └── visualize_results.py
│
├── requirements.txt
├── LICENSE
└── README.md

---

Evaluation Results

Evaluation was performed on the validation set.

Metric	Score
Accuracy	97.31%
IoU	91.82%
Dice Score	95.67%
Precision	95.45%
Recall	96.04%

Best model:

checkpoints/best_unet.pth

---

Prediction Examples

The model generates road segmentation masks and overlay visualizations.

Comparison Gallery

Prediction Gallery

---

Running the Project

Train Model

python scripts/train_unet.py

Run Prediction

Single image:

python scripts/predict.py --image dataset/val/images/0.png

Folder prediction:

python scripts/predict.py --folder dataset/val/images

Evaluate Model

python scripts/evaluate.py

Evaluation results are saved to:

reports/metrics.txt

---

Technologies Used

• Python
• PyTorch
• segmentation_models_pytorch
• Albumentations
• OpenCV
• NumPy
• Matplotlib
• Git
• GitHub

---

Future Improvements

• Expand training dataset
• Add more diverse road environments
• Experiment with transformer-based segmentation models
• Export model to ONNX
• Optimize inference speed for real-time applications

---

License

This project is licensed under the MIT License. See the LICENSE file for details.