WUT-AI-AI4Mat/microstructure-swin-unet

Swin-Unet for Quantitative Microstructure Segmentation

Model Description

This repository contains seven dataset-specific Swin-Unet checkpoints. The architecture is Swin-Tiny with patch size 4 and window size 7. Images are resized to 224 x 224 and normalized with ImageNet statistics before inference. 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	Classes	File
Aachen-Heerlen	2	checkpoints/Aachen-Heerlen/best_model.pth
EMPS	2	checkpoints/EMPS/best_model.pth
Grain	2	checkpoints/Grain/best_model.pth
EBC	3	checkpoints/EBC/best_model.pth
Super	3	checkpoints/Super/best_model.pth
MetalDAM	5	checkpoints/MetalDAM/best_model.pth
UHCS	7	checkpoints/UHCS/best_model.pth

Class counts include background.

Training Details

• Input size: 224 x 224
• Patch size: 4
• Window size: 7
• Batch size: 32
• Epochs: 500
• Optimizer: AdamW
• Learning rate: 0.0001
• Weight decay: 0.0001
• Scheduler: CosineAnnealingLR with minimum learning rate 0.000001
• Early-stopping patience: 50
• Objective: 0.5 cross-entropy and 0.5 Dice
• Initialization: ImageNet-pretrained Swin-Tiny

Usage

Clone the upstream Swin-Unet implementation as described in the source repository. Example for EMPS:

hf download NAMESPACE/microstructure-swin-unet checkpoints/EMPS/best_model.pth --local-dir weights/swin-unet
python Swin-Unet/test_data.py \
  --dataset-root /path/to/EMPS \
  --dataset-name EMPS \
  --num-classes 2 \
  --checkpoint weights/swin-unet/checkpoints/EMPS/best_model.pth \
  --output-dir results/swin_unet_emps \
  --device cuda

Evaluation

The testing script saves original-resolution masks and reports mIoU, Dice, precision, recall, accuracy, HD95, Hausdorff distance, NSD, MAE, MBSS, and MBSS_add through Myutils/metrics.py. Numerical benchmark results will be linked after the associated paper becomes publicly available.

Intended Use and Limitations

The checkpoints are intended for research reproduction on the named microstructure datasets. They may not generalize to different materials, imaging settings, magnifications, or annotation policies. 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 Swin-Unet implementation and pretrained initialization remain subject to their upstream terms. A paper citation will be added after publication.