Defect-aware MAE
A Vision Transformer encoder pretrained without labels on steel-surface images with a masked autoencoder adapted to texture defects. Frozen features are clustered into defect categories.
No class labels, bounding boxes or detection objective are used in pretraining, and no contrastive loss.
Code, training scripts, stored results and reproduction notebooks: https://github.com/ShreyPatel1311/NEU-Defect-Aware-Autoencoder
What is different from a standard MAE
| Component | Standard MAE | This model |
|---|---|---|
| Reconstruction target | per-patch normalised pixels | per-patch HOG descriptor (9 orientation bins × 4 cells = 36-d, L2-normalised) |
| Masking | uniform random, 75% | defect-aware, 75%: patches with high local texture energy (pixel std) are masked with higher probability (Gumbel-top-k, β = 1) |
| Augmentation | flips, 90° rotation, jitter | plus random-resized-crop and blur |
| Decoder | 4 layers | 2 layers (256-d, 8 heads) |
| Readout | CLS token | mean-pooled patch tokens of blocks 4–6, each instance-normalised, then averaged |
Encoder: 224×224 grayscale input, 16×16 patches (196 tokens), 6 pre-norm Transformer layers, 12 heads, 768-d.
Files
| File | Contents |
|---|---|
defect_aware_mae_seed0.pth |
{"encoder", "decoder", "epoch"} state dicts, 1,600 epochs, seed 0 |
simclr_baseline_seed0.pth |
{"encoder", "epoch"}: contrastive baseline on the same encoder, 800 epochs, seed 0 |
*_args.json |
exact training arguments for each checkpoint |
defect_aware_mae_seed0_history.json |
per-epoch training loss |
model.py |
model definition, preprocessing and feature extraction |
Usage
import torch
from model import Encoder, embed # model.py from this repo
enc = Encoder()
enc.load_state_dict(torch.load("defect_aware_mae_seed0.pth", map_location="cpu")["encoder"])
enc.cuda().eval()
# x_uint8: (n, 224, 224) uint8 grayscale images after CLAHE (clipLimit=2.0, tileGridSize=8x8) and resize to 224
features = embed(enc, x_uint8, mode="mid") # (n, 768); use mode="cls" for the SimCLR baseline
The decoder in the MAE checkpoint has 2 layers and a 36-d output head:
Decoder(depth=2) with dec.pred = torch.nn.Linear(256, 36).
Training data
NEU-DET (hot-rolled steel strip, 6 defect classes, 1,800 images). Pretraining used only the 1,440-image training split. The 360-image validation split was held out from pretraining.
Results
Held-out validation split (360 images never seen in pretraining). Features are standardised, projected to 2-D with UMAP (cosine, 25 neighbours, min_dist 0.1) and grouped by Ward agglomerative clustering (K = 6); accuracy is Hungarian-matched. Mean ± std over 5 seeds.
| Model | Accuracy (%) | NMI | ARI | Silhouette |
|---|---|---|---|---|
| SimCLR baseline | 97.8 ± 1.3 | 0.955 ± 0.021 | 0.951 ± 0.027 | 0.736 ± 0.019 |
| Defect-aware MAE | 98.6 ± 0.2 | 0.968 ± 0.003 | 0.967 ± 0.004 | 0.838 ± 0.019 |
The released seed-0 checkpoints score 98.33% (Defect-aware MAE: NMI 0.963, ARI 0.961, silhouette 0.827; SimCLR: NMI 0.955, ARI 0.961, silhouette 0.738).
License
Weights and code in this repository: MIT (see LICENSE). The NEU-DET dataset belongs to its original authors
(K. Song and Y. Yan, Northeastern University) and is not covered by this license.