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.
Evaluated on 54 held-out test images, scored at image level to avoid batch-size bias.
1import torch
2from huggingface_hub import hf_hub_download
3import segmentation_models_pytorch as smp
4
5# Download checkpoint
6ckpt_path = hf_hub_download(
7 repo_id="imranlabs/sem-microstructure-segmentation",
8 filename="best_model.pth"
9)
10
11# Rebuild architecture
12model = smp.Unet(
13 encoder_name = "resnet34",
14 encoder_weights = None, # weights loaded from checkpoint
15 in_channels = 1,
16 classes = 5,
17 activation = None,
18)
19
20# Load weights
21ckpt = torch.load(ckpt_path, map_location="cpu")
22model.load_state_dict(ckpt["model_state"])
23model.eval()
24
25# Inference — input: (1, 1, H, W) float32 tensor normalised to [0, 1]
26with torch.no_grad():
27 logits = model(image_tensor) # (1, 5, H, W)
28 preds = torch.argmax(logits, dim=1) # (1, H, W) class labels
