DriftFox/voronoi-microstructure-resnet

Voronoi microstructure ResNet

Synthetic-pretrained ResNet encoders for estimating microstructure geometry from 256 x 256 grayscale images. This release contains four frozen pretraining checkpoints, their parameter normalization, and an inference loader. The models predict 13 geometric parameters and their uncertainties; they also provide encoder embeddings for downstream models.

Released checkpoints

Variant	Synthetic training images	Training seed	Embedding dimension
resnet18_seed1019	3,072,000	1019	512
resnet18_seed2026	3,072,000	2026	512
resnet34_seed1019	3,072,000	1019	512
resnet50_seed1019	3,072,000	1019	2048

Each variant has a .safetensors file. config.json records the source checkpoint checksum, released file checksum, best epoch, target normalization, and synthetic validation metric. SHA256SUMS covers all release files. The release preserves every original model tensor exactly. The original normalization arrays are included as float32 buffers for inference.

Training and data provenance

All four encoders were initialized using Torchvision ImageNet-1K V1 weights, then trained to recover known parameters of procedurally rendered Voronoi microstructures. ImageNet pretraining is upstream training; this is not a from-scratch model. The initialization URLs and original copyright notices are in THIRD_PARTY_NOTICES.md.

Training settings: 25 epochs, batch size 64, AdamW with learning rate 3e-4 and weight decay 0.001, cosine decay to 1e-5, gradient clipping at 5, BF16 mixed precision, and random 90-degree rotations / horizontal flips. There is no scale augmentation. Target normalization is fitted on synthetic training rows only. The loss is Gaussian negative log likelihood plus 0.1 times normalized mean squared error. The best checkpoint is selected by the mean of the 13 normalized synthetic validation RMSEs.

Architecture and outputs

The original ResNet classifier is replaced with an identity layer. A dimension -> 128 -> 26 head with a SiLU activation produces 13 normalized means and 13 log variances, with log variances clamped to [-4, 3]. Grayscale inputs are repeated across three channels and normalized using ImageNet mean [0.485, 0.456, 0.406] and standard deviation [0.229, 0.224, 0.225].

The output order is fixed:

Name	Synthetic target definition
log_grains	log(1 + N), where N is the number of Voronoi grains present in the image, including cropped border grains
mean_ecd	Unweighted mean equivalent circular diameter, 2 * sqrt(grain_area / pi), in pixels
area_weighted_ecd	Mean ECD weighted by each grain's pixel area
ecd_cv	Population standard deviation of ECD divided by mean ECD
hall_petch_mean_ecd	1 / sqrt(mean_ecd), with ECD measured in pixels
phase_area_fraction	Area-weighted fraction of grains assigned to the darker phase
phase_grain_fraction	Unweighted fraction of grains assigned to the darker phase
log_stretch	Natural logarithm of the renderer's anisotropic coordinate stretch
log_defects	log(1 + K), where K is the number of generated defect ellipses, including overlapping/cropped ellipses
defect_area_fraction	Fraction of image pixels covered by the union of defect masks
mean_defect_radius	Mean r * sqrt(aspect) of generated ellipses in pixels; zero when K is zero
boundary_fraction	Fraction of pixels marked by find_boundaries(labels, mode='inner') before blur
interior_grain_fraction	Fraction of present grains that do not touch any image edge

predict() reverses target standardization, but preserves the three explicitly named log-domain targets. For example, recover estimated counts with expm1(log_grains) / expm1(log_defects) and stretch with exp(log_stretch). Predictions are continuous estimates and are not clipped to physical bounds.

Usage

Tested with PyTorch 2.9.1, Torchvision 0.24.1, and safetensors. CPU inference is supported; a GPU is optional.

pip install torch torchvision safetensors huggingface_hub pillow numpy
hf download DriftFox/voronoi-microstructure-resnet --local-dir voronoi-microstructure-resnet
cd voronoi-microstructure-resnet

import numpy as np
from PIL import Image
import torch
from model import load_model, TARGET_NAMES

image = np.asarray(Image.open('microstructure.png').convert('L'), dtype=np.float32)
assert image.shape == (256, 256)  # Use native resolution; resizing changes grain size.
x = torch.from_numpy(image.copy())[None, None] / 255.0
model = load_model('.', variant='resnet18_seed1019', device='cpu')
parameters, uncertainty, embedding = model.predict(x, tta=True)
print(dict(zip(TARGET_NAMES, parameters[0].tolist())))
print(uncertainty.shape, embedding.shape)

Default TTA averages all eight 90-degree rotation / horizontal flip views. The reported standard deviation combines mean predicted variance and between-view variance in normalized target space, then rescales it. model(x) returns normalized means, log variances, and the embedding directly. python model.py runs a small load/inference check.

Limitations

The model estimates renderer-defined properties, not segmented grain or defect masks. Generated defect count is not the number of visible connected dark regions when ellipses overlap. Cropped grains affect pixel-based size estimates. Real microscopy can differ substantially from the renderer, so physical parameter accuracy and uncertainty calibration require independent real-image validation. ECDs are in pixels, not micrometers.

License and attribution

This release is distributed under BSD-3-Clause, as stated in LICENSE. Redistributions must also retain the original Torchvision notice, license conditions, and disclaimer in THIRD_PARTY_NOTICES.md. Original author and contributor names must not be used to imply endorsement.