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.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including WUT-AI-AI4Mat/microstructure-unet