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
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:
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/crackseg.git
cd crackseg
Install requirements:
pip install -r requirements
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 -m tools.main
Training arguments:
python -m tools.main -h
usage: main.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 -m tools.inference --weights weights/model.pt --input assets/CFD_001_image.jpg
Inference arguments
python -m tools.inference -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
If you're going to contribute to this project, please see the Contributing policy and Code of Conduct. Do not hesitate to report bugs or request new features.
License
The project is licensed under the MIT license.