SEM Microstructure Semantic Segmentation

UNet with a pretrained ResNet-34 encoder for pixel-wise segmentation of Scanning Electron Microscopy (SEM) images of additively manufactured Ni-WC metal matrix composites. Trained to identify five microstructural phases at pixel level.

Model Details

Property	Detail
Architecture	UNet + ResNet-34 encoder
Encoder weights	ImageNet pretrained
Input channels	1 (grayscale SEM image)
Output classes	5 (pixel-wise segmentation)
Framework	PyTorch + segmentation_models_pytorch
Training device	Apple MPS (M-series)

Classes

ID	Phase
0	Matrix
1	Carbide
2	Void
3	Reprecipitate
4	Dilution Zone

Test Performance

Evaluated on 54 held-out test images, scored at image level to avoid batch-size bias.

Metric	Mean	Std	Min	Max
mIoU	0.872	0.088	0.723	0.958
mDice	0.912	0.079	0.759	0.978
mBF1	0.728	0.027	0.678	0.773

Per-Class Performance

Class	IoU	Dice
Matrix	0.939	0.969
Carbide	0.753	0.859
Void	0.976	0.988
Reprecipitate	0.891	0.942
Dilution Zone	0.881	0.937

Files

File	Description
`best_model.pth`	Best pretrained encoder checkpoint
`ScratchUNet_best_unet.pth`	Best scratch-built UNet checkpoint

How to Use

Install dependencies

pip install torch segmentation-models-pytorch huggingface_hub

Load and run inference

import torch
from huggingface_hub import hf_hub_download
import segmentation_models_pytorch as smp

# Download checkpoint
ckpt_path = hf_hub_download(
    repo_id="imranlabs/sem-microstructure-segmentation",
    filename="best_model.pth"
)

# Rebuild architecture
model = smp.Unet(
    encoder_name    = "resnet34",
    encoder_weights = None,       # weights loaded from checkpoint
    in_channels     = 1,
    classes         = 5,
    activation      = None,
)

# Load weights
ckpt = torch.load(ckpt_path, map_location="cpu")
model.load_state_dict(ckpt["model_state"])
model.eval()

# Inference — input: (1, 1, H, W) float32 tensor normalised to [0, 1]
with torch.no_grad():
    logits = model(image_tensor)          # (1, 5, H, W)
    preds  = torch.argmax(logits, dim=1)  # (1, H, W) class labels

Training Details

Loss: 0.6 × Weighted CrossEntropy + 0.4 × DiceLoss
Optimizer: AdamW (weight decay 1e-5)
Scheduler: ReduceLROnPlateau
Strategy: Encoder frozen for first 4 epochs, then full fine-tuning
Early stopping: Monitored validation mIoU
Best checkpoint: Epoch 12, validation mIoU = 0.876

Dataset

Safdar, M. (2025). Scanning Electron Microscopy (SEM) Dataset of Additively Manufactured Ni-WC Metal Matrix Composites for Semantic Segmentation (Version 1). Zenodo. https://doi.org/10.5281/zenodo.17315241

Author

Imran Khan — Physics PhD · SEM Metrology Engineer · Computer Vision / AI

Downloads last month: -; Downloads are not tracked for this model. How to track

imranlabs
/

sem-microstructure-segmentation