jimmodels/microhard-uhcs-segmenter

microhard UHCS microconstituent segmenter

A U-Net that labels each pixel of an SEM micrograph of ultrahigh carbon steel as one of four microconstituents: ferritic matrix, proeutectoid cementite network, spheroidite, or Widmanstätten cementite. It is the segmentation stage of microhard, a pipeline that goes from a micrograph to microstructure fractions to an estimated property such as hardness.

The encoder is a resnet50 pretrained on microscopy images (NASA MicroNet) and kept frozen; only the U-Net decoder was trained, on the 24 pixel-labeled images of the DeCost UHCS segmentation benchmark. The checkpoint bundles the frozen encoder weights, so it loads without any external download.

What to expect

This is a proof-of-concept trained on 24 images, not a production model. On validation samples (split so that no micrograph of a training sample appears in validation) it reaches a mean IoU of about 0.50. For reference, the DeCost 2019 paper reaches roughly 0.7+ by fine-tuning the whole network; training only the decoder trades some accuracy for a shared, reusable backbone.

Per-class IoU is uneven. Spheroidite and the cementite network segment well (around 0.6 to 0.8). Widmanstätten laths are rare in the labeled set and segment poorly (often below 0.1). The example figure above shows this directly: the network and spheroidite regions are close to the ground truth, while the thin Widmanstätten laths on the right are missed.

Usage

The checkpoint is a plain state dict (loads with weights_only=True) plus the encoder name and the ordered class list. This snippet reproduces the pipeline's output exactly and needs only torch, segmentation-models-pytorch, albumentations, pillow, and huggingface_hub.

import numpy as np, torch
import segmentation_models_pytorch as smp
import albumentations as A
from albumentations.pytorch import ToTensorV2
from huggingface_hub import hf_hub_download
from PIL import Image

path = hf_hub_download("jimmodels/microhard-uhcs-segmenter", "segmenter.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=True)
classes = ckpt["class_nodes"]  # ['ferrous/matrix', 'ferrous/network', 'ferrous/spheroidite', 'ferrous/widmanstatten']

model = smp.Unet(encoder_name=ckpt["encoder"], encoder_weights=None,
                 in_channels=3, classes=len(classes))
model.load_state_dict(ckpt["state_dict"])
model.eval()

# The model was trained on images padded (not resized) to a multiple of 32,
# with ImageNet normalisation. Resizing would corrupt the micron-per-pixel scale.
transform = A.Compose([
    A.PadIfNeeded(min_height=None, min_width=None,
                  pad_height_divisor=32, pad_width_divisor=32),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

def segment(image_path):
    img = np.asarray(Image.open(image_path).convert("RGB"))
    h, w = img.shape[:2]
    x = transform(image=img)["image"].unsqueeze(0)
    with torch.no_grad():
        pred = model(x).argmax(1)[0].numpy()
    top, left = (pred.shape[0] - h) // 2, (pred.shape[1] - w) // 2
    return pred[top:top + h, left:left + w]  # class index per pixel, cropped to input size

Or use the pipeline directly, which also computes area fractions and, where calibrated, a property estimate:

pip install git+https://github.com/jamhan/MicrostructurePredictor

from pathlib import Path
from huggingface_hub import hf_hub_download
from microhard.config import Config
from microhard.segment import load_segmenter, segment_image

path = hf_hub_download("jimmodels/microhard-uhcs-segmenter", "segmenter.pt")
cfg = Config(checkpoint_dir=Path(path).parent)
model, class_nodes = load_segmenter(cfg)

Training data

The DeCost UHCS segmentation benchmark: 24 SEM micrographs of a 2C-4Cr ultrahigh carbon steel with per-pixel microconstituent labels, originally at NIST handle 11256/964 and mirrored in bdecost/uhcs-segment. The 38 px instrument banner was cropped from every image and label before training. Micrographs were collected by Matthew Hecht (Carnegie Mellon University).

The decoder was trained for 14 epochs (Dice plus cross-entropy loss, AdamW), with the train/validation split grouped by physical sample so that near-duplicate micrographs of one sample do not straddle the split.

Limitations

The labeled set is 24 images of a single alloy family, so this model should not be expected to transfer to other steels or other imaging conditions without new data. The four classes lump several matrix constituents (pearlite, bainite, martensite) into one "matrix" label, which limits how much downstream property work can distinguish heat treatments. Predictions are least reliable for the rare Widmanstätten class and along constituent boundaries.

License and attribution

Released under the MIT license. The encoder weights derive from NASA's pretrained-microscopy-models (MicroNet, MIT). The training data is the UHCS dataset distributed by NIST under a Creative Commons license.

If you use this model, please cite the underlying work:

• DeCost, Lei, Francis, Holm, "High throughput quantitative metallography for complex microstructures using deep learning," Microscopy and Microanalysis 25 (2019).
• Stuckner, Harder, Smith, "Microstructure segmentation with deep learning encoders pre-trained on a large microscopy dataset," npj Computational Materials 8, 200 (2022).
• Hecht, "Effects of Heat Treatments and Compositional Modification on Carbide Network and Matrix Microstructure in Ultrahigh Carbon Steels," PhD thesis, Carnegie Mellon University (2017).