U-Net for Quantitative Microstructure Segmentation
Model Description
This repository contains seven dataset-specific U-Net checkpoints used in a
benchmark of microstructure segmentation methods. The model is implemented
with segmentation_models_pytorch and uses a ResNet-50 encoder initialized
from ImageNet weights. Inputs are resized to 512 x 512 pixels and predictions
are restored to the original image size with nearest-neighbor interpolation.
Source code: https://github.com/WUT-AI-AI4Mat/Segmentation-methods-evaluation-for-quantitative-microstructure-analysis
Checkpoints
| Dataset | Task | Output channels | File |
|---|---|---|---|
| Aachen-Heerlen | Binary segmentation | 1 | checkpoints/Aachen-Heerlen/Best_Model.pth |
| EMPS | Binary segmentation | 1 | checkpoints/EMPS/Best_Model.pth |
| Grain | Binary segmentation | 1 | checkpoints/Grain/Best_Model.pth |
| EBC | 3-class segmentation | 3 | checkpoints/EBC/Best_Model.pth |
| Super | 3-class segmentation | 3 | checkpoints/Super/Best_Model.pth |
| MetalDAM | 5-class segmentation | 5 | checkpoints/MetalDAM/Best_Model.pth |
| UHCS | 7-class segmentation | 7 | checkpoints/UHCS/Best_Model.pth |
Binary checkpoints produce one foreground logit; background is represented by the complementary binary label. Multiclass counts include background.
Training Details
- Input size: 512 x 512
- Batch size: 32
- Epochs: 500
- Optimizer: AdamW
- Learning rate: 0.0003
- Weight decay: 0.001
- Scheduler: CosineAnnealingLR with minimum learning rate 0.000001
- Early-stopping patience: 50
- Binary objective: Dice and BCEWithLogits
- Multiclass objective: Dice and cross-entropy
Usage
Download a checkpoint and run the matching test entry point from the source repository. Example for binary EMPS segmentation:
hf download NAMESPACE/microstructure-unet checkpoints/EMPS/Best_Model.pth --local-dir weights/unet
python CNN/U-net/predict.py \
--dataset-root /path/to/EMPS \
--checkpoint weights/unet/checkpoints/EMPS/Best_Model.pth \
--output-dir results/unet_emps
For EBC, Super, MetalDAM, and UHCS, use CNN/U-net/test_mutil.py and pass the
matching --num-classes value from the table.
Evaluation
The public evaluation scripts restore predictions to the original resolution,
save raw masks, and calculate mIoU, Dice, precision, recall, accuracy, HD95,
Hausdorff distance, NSD, MAE, MBSS, and MBSS_add using Myutils/metrics.py.
Numerical benchmark results will be linked after the associated paper becomes
publicly available.
Intended Use and Limitations
These checkpoints are intended for research reproduction and comparative evaluation on the named microstructure datasets. They are dataset-specific and should not be treated as general-purpose materials segmentation models. Image acquisition conditions, magnification, annotation conventions, and unseen microstructures may cause substantial performance degradation. They are not validated for safety-critical or industrial quality-control decisions.
License and Citation
The released experiment files are provided under the MIT license. The model implementation and ImageNet initialization also remain subject to their upstream terms. A paper citation will be added after publication.