- Notifications
You must be signed in to change notification settings - Fork4
🛣️🔍 | Road crack segmentation using UNet in PyTorch > Implementation of different loss functions (i.e Focal, Dice, Dice + CE)
License
yakhyo/crack-segmentation
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
| 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 thisproject we usedUNet to detect the cracks on the road.
In this, Road Crack Segmentation project we have implemented UNet model to segment cracks on the roadusingCrack Segmentation dataset. Weevaluated the model's performance using different loss functions and compared their results. We have implementedfollowing loss functions:
We trained the model using above-mentioned loss functions and evaluated their performance byDice 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...)
Download the project:
git clone https://github.com/yakhyo/crack-segmentation.gitcd crack-segmentationInstall requirements:
pip install -r requirements.txtDownload the weights of the model fromhere intoweightsfolder. Model weights trained using Dice Loss is provided insideweights folder asmodel.pt
To train the model download the dataset and puttrain andtest folders insidedata folder as following:
data-| |-train-| |-images |-masks |-test -| |-images |-maskspython train.pyTraining arguments:
python train.py -husage: 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 argumentsoptional 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 classespython inference.py --weights weights/model.pt --input assets/CFD_001_image.jpgInference arguments
python inference.py -husage: inference.py [-h] [--weights WEIGHTS] [--input INPUT] [--output OUTPUT] [--image-size IMAGE_SIZE] [--view] [--no-save] [--conf-thresh CONF_THRESH]Crack Segmentation inference argumentsoptional 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 maskContributions to improve the crack segmentation project are welcome. Feel free to fork the repository and submit pullrequests,or open issues to suggest features or report bugs.
The project is licensed under theMIT license.
About
🛣️🔍 | Road crack segmentation using UNet in PyTorch > Implementation of different loss functions (i.e Focal, Dice, Dice + CE)