Road Crack Segmentation using UNet

Input Image	Ground Truth	CE Loss	Dice Loss	DiceCE Loss	Focal Loss
Dice Score➡️		0.9719	0.9804	0.9754	0.9679

Road crack segmentation is the task of identifying and segmenting road cracks in images or videos of roads. In this project we used UNet to detect the cracks on the road.

Table of Contents

• Project Description
• Installation
• Usage
• Contributing
• License

Project Description

In this, Road Crack Segmentation project we have implemented UNet model to segment cracks on the road using Crack Segmentation dataset. We evaluated the model's performance using different loss functions and compared their results. We have implemented following loss functions:

• Cross Entropy Loss
• Dice Loss
• Dice Cross Entropy Loss
• Focal Loss

We trained the model using above-mentioned loss functions and evaluated their performance by Dice coefficient (dice score = 1 - dice loss, so our evaluation metric may not be right metric to compare those models with each other)

*see here to check full implementation of more loss functions ( updating...)

Installation

Download the project:

git clone https://github.com/yakhyo/crack-segmentation.git
cd crack-segmentation

Install requirements:

pip install -r requirements.txt

Download the weights of the model from here into weights folder. Model weights trained using Dice Loss is provided inside weights folder as model.pt

Usage

Dataset

To train the model download the dataset and put train and test folders inside data folder as following:

data-|
     |-train-|
             |-images
             |-masks
                
     |-test -|
             |-images
             |-masks

Train

python train.py

Training arguments:

python train.py -h
usage: train.py [-h] [--data DATA] [--image_size IMAGE_SIZE] [--save-dir SAVE_DIR] [--epochs EPOCHS] [--batch-size BATCH_SIZE] [--lr LR] [--weights WEIGHTS] [--amp] [--num-classes NUM_CLASSES]

Crack Segmentation training arguments

optional arguments:
  -h, --help            show this help message and exit
  --data DATA           Path to root folder of data
  --image_size IMAGE_SIZE
                        Input image size, default: 512
  --save-dir SAVE_DIR   Directory to save weights
  --epochs EPOCHS       Number of epochs, default: 5
  --batch-size BATCH_SIZE
                        Batch size, default: 12
  --lr LR               Learning rate, default: 1e-5
  --weights WEIGHTS     Pretrained model, default: None
  --amp                 Use mixed precision
  --num-classes NUM_CLASSES
                        Number of classes

Inference

python inference.py --weights weights/model.pt --input assets/CFD_001_image.jpg

Inference arguments

python inference.py -h
usage: inference.py [-h] [--weights WEIGHTS] [--input INPUT] [--output OUTPUT] [--image-size IMAGE_SIZE] [--view] [--no-save] [--conf-thresh CONF_THRESH]

Crack Segmentation inference arguments

optional arguments:
  -h, --help            show this help message and exit
  --weights WEIGHTS     Path to weight file (default: last.pt)
  --input INPUT         Path to input image
  --output OUTPUT       Path to save mask image
  --image-size IMAGE_SIZE
                        Input image size
  --view                Visualize image and mask
  --no-save             Do not save the output masks
  --conf-thresh CONF_THRESH
                        Confidence threshold for mask

Contributing

Contributions to improve the crack segmentation project are welcome. Feel free to fork the repository and submit pull requests,or open issues to suggest features or report bugs.

License

The project is licensed under the MIT license.