---
license: mit
library_name: segmentation_models_pytorch
pipeline_tag: image-segmentation
tags:
- semantic-segmentation
- microstructure
- materials-science
- pytorch
- unet
---

# 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:

```bash
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.