README.md

---
license: mit
pipeline_tag: image-segmentation
tags:
  - materials-science
  - metallography
  - microscopy
  - steel
  - u-net
  - pytorch
---

# 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](https://github.com/jamhan/MicrostructurePredictor), 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.

![input, ground truth, and prediction on a held-out benchmark image](example_prediction.png)

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

```python
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:

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

```python
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](https://hdl.handle.net/11256/964) and mirrored in
[bdecost/uhcs-segment](https://github.com/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](https://github.com/nasa/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).