jimmodels commited on
Commit
2b86ba2
·
verified ·
1 Parent(s): b77185d

Add model card

Browse files
Files changed (1) hide show
  1. README.md +143 -0
README.md ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ pipeline_tag: image-segmentation
4
+ tags:
5
+ - materials-science
6
+ - metallography
7
+ - microscopy
8
+ - steel
9
+ - u-net
10
+ - pytorch
11
+ ---
12
+
13
+ # microhard UHCS microconstituent segmenter
14
+
15
+ A U-Net that labels each pixel of an SEM micrograph of ultrahigh carbon steel
16
+ as one of four microconstituents: ferritic matrix, proeutectoid cementite
17
+ network, spheroidite, or Widmanstätten cementite. It is the segmentation stage
18
+ of [microhard](https://github.com/jamhan/MicrostructurePredictor), a pipeline
19
+ that goes from a micrograph to microstructure fractions to an estimated
20
+ property such as hardness.
21
+
22
+ The encoder is a resnet50 pretrained on microscopy images (NASA MicroNet) and
23
+ kept frozen; only the U-Net decoder was trained, on the 24 pixel-labeled images
24
+ of the DeCost UHCS segmentation benchmark. The checkpoint bundles the frozen
25
+ encoder weights, so it loads without any external download.
26
+
27
+ ![input, ground truth, and prediction on a held-out benchmark image](example_prediction.png)
28
+
29
+ ## What to expect
30
+
31
+ This is a proof-of-concept trained on 24 images, not a production model. On
32
+ validation samples (split so that no micrograph of a training sample appears in
33
+ validation) it reaches a mean IoU of about 0.50. For reference, the DeCost 2019
34
+ paper reaches roughly 0.7+ by fine-tuning the whole network; training only the
35
+ decoder trades some accuracy for a shared, reusable backbone.
36
+
37
+ Per-class IoU is uneven. Spheroidite and the cementite network segment well
38
+ (around 0.6 to 0.8). Widmanstätten laths are rare in the labeled set and segment
39
+ poorly (often below 0.1). The example figure above shows this directly: the
40
+ network and spheroidite regions are close to the ground truth, while the thin
41
+ Widmanstätten laths on the right are missed.
42
+
43
+ ## Usage
44
+
45
+ The checkpoint is a plain state dict (loads with `weights_only=True`) plus the
46
+ encoder name and the ordered class list. This snippet reproduces the pipeline's
47
+ output exactly and needs only `torch`, `segmentation-models-pytorch`,
48
+ `albumentations`, `pillow`, and `huggingface_hub`.
49
+
50
+ ```python
51
+ import numpy as np, torch
52
+ import segmentation_models_pytorch as smp
53
+ import albumentations as A
54
+ from albumentations.pytorch import ToTensorV2
55
+ from huggingface_hub import hf_hub_download
56
+ from PIL import Image
57
+
58
+ path = hf_hub_download("jimmodels/microhard-uhcs-segmenter", "segmenter.pt")
59
+ ckpt = torch.load(path, map_location="cpu", weights_only=True)
60
+ classes = ckpt["class_nodes"] # ['ferrous/matrix', 'ferrous/network', 'ferrous/spheroidite', 'ferrous/widmanstatten']
61
+
62
+ model = smp.Unet(encoder_name=ckpt["encoder"], encoder_weights=None,
63
+ in_channels=3, classes=len(classes))
64
+ model.load_state_dict(ckpt["state_dict"])
65
+ model.eval()
66
+
67
+ # The model was trained on images padded (not resized) to a multiple of 32,
68
+ # with ImageNet normalisation. Resizing would corrupt the micron-per-pixel scale.
69
+ transform = A.Compose([
70
+ A.PadIfNeeded(min_height=None, min_width=None,
71
+ pad_height_divisor=32, pad_width_divisor=32),
72
+ A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
73
+ ToTensorV2(),
74
+ ])
75
+
76
+ def segment(image_path):
77
+ img = np.asarray(Image.open(image_path).convert("RGB"))
78
+ h, w = img.shape[:2]
79
+ x = transform(image=img)["image"].unsqueeze(0)
80
+ with torch.no_grad():
81
+ pred = model(x).argmax(1)[0].numpy()
82
+ top, left = (pred.shape[0] - h) // 2, (pred.shape[1] - w) // 2
83
+ return pred[top:top + h, left:left + w] # class index per pixel, cropped to input size
84
+ ```
85
+
86
+ Or use the pipeline directly, which also computes area fractions and, where
87
+ calibrated, a property estimate:
88
+
89
+ ```bash
90
+ pip install git+https://github.com/jamhan/MicrostructurePredictor
91
+ ```
92
+
93
+ ```python
94
+ from pathlib import Path
95
+ from huggingface_hub import hf_hub_download
96
+ from microhard.config import Config
97
+ from microhard.segment import load_segmenter, segment_image
98
+
99
+ path = hf_hub_download("jimmodels/microhard-uhcs-segmenter", "segmenter.pt")
100
+ cfg = Config(checkpoint_dir=Path(path).parent)
101
+ model, class_nodes = load_segmenter(cfg)
102
+ ```
103
+
104
+ ## Training data
105
+
106
+ The DeCost UHCS segmentation benchmark: 24 SEM micrographs of a 2C-4Cr
107
+ ultrahigh carbon steel with per-pixel microconstituent labels, originally at
108
+ NIST handle [11256/964](https://hdl.handle.net/11256/964) and mirrored in
109
+ [bdecost/uhcs-segment](https://github.com/bdecost/uhcs-segment). The 38 px
110
+ instrument banner was cropped from every image and label before training.
111
+ Micrographs were collected by Matthew Hecht (Carnegie Mellon University).
112
+
113
+ The decoder was trained for 14 epochs (Dice plus cross-entropy loss, AdamW),
114
+ with the train/validation split grouped by physical sample so that
115
+ near-duplicate micrographs of one sample do not straddle the split.
116
+
117
+ ## Limitations
118
+
119
+ The labeled set is 24 images of a single alloy family, so this model should not
120
+ be expected to transfer to other steels or other imaging conditions without new
121
+ data. The four classes lump several matrix constituents (pearlite, bainite,
122
+ martensite) into one "matrix" label, which limits how much downstream property
123
+ work can distinguish heat treatments. Predictions are least reliable for the
124
+ rare Widmanstätten class and along constituent boundaries.
125
+
126
+ ## License and attribution
127
+
128
+ Released under the MIT license. The encoder weights derive from NASA's
129
+ [pretrained-microscopy-models](https://github.com/nasa/pretrained-microscopy-models)
130
+ (MicroNet, MIT). The training data is the UHCS dataset distributed by NIST under
131
+ a Creative Commons license.
132
+
133
+ If you use this model, please cite the underlying work:
134
+
135
+ - DeCost, Lei, Francis, Holm, "High throughput quantitative metallography for
136
+ complex microstructures using deep learning," *Microscopy and Microanalysis*
137
+ 25 (2019).
138
+ - Stuckner, Harder, Smith, "Microstructure segmentation with deep learning
139
+ encoders pre-trained on a large microscopy dataset," *npj Computational
140
+ Materials* 8, 200 (2022).
141
+ - Hecht, "Effects of Heat Treatments and Compositional Modification on Carbide
142
+ Network and Matrix Microstructure in Ultrahigh Carbon Steels," PhD thesis,
143
+ Carnegie Mellon University (2017).