Steel Defect Classification — ResNet18 (NEU-DET)
Classifies hot-rolled steel surface defects into 6 classes (rolled-in scale, patches, crazing, pitted surface, inclusion, scratches) with a ResNet18 fine-tuned on the real NEU-DET database. Grad-CAM heatmaps included for explainability.
Results (held-out test, 264 images — stratified 70/15/15, seed 42)
| Metric | Value |
|---|---|
| Test accuracy | 0.985 |
| Test ROC-AUC (one-vs-rest) | 0.9999 |
| Macro F1 | 0.985 |
| Class | Precision | Recall | F1 |
|---|---|---|---|
| Rolled-in Scale | 1.000 | 1.000 | 1.000 |
| Patches | 1.000 | 0.977 | 0.989 |
| Crazing | 0.957 | 1.000 | 0.978 |
| Pitted Surface | 1.000 | 0.955 | 0.977 |
| Inclusion | 0.977 | 0.977 | 0.977 |
| Scratches | 0.978 | 1.000 | 0.989 |
Training
- Base: ResNet18 (ImageNet-pretrained), final FC replaced with 6-way head
- Data: NEU Surface Defect Database — ~1,770 images (200×200 grayscale, ~295/class; canonical 300/class; a few upstream images absent from the mirror)
- Augmentation (train only): horizontal/vertical flip, rotation ±10°, translate ±5%
- Optimization: Adam (lr 1e-3, wd 1e-4), ReduceLROnPlateau, weighted cross-entropy, early stopping (patience 8) with best-epoch restore — stopped at epoch 16, best epoch 7
- Hardware: CPU (Ryzen 9 7950X, 16 threads), ~35 min; seed 42
Usage
import json
import torch
from safetensors.torch import load_file
from torchvision import models
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
cfg = json.load(open("config.json"))
model = models.resnet18(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, len(cfg["class_names"]))
model.load_state_dict(load_file("model.safetensors"))
model.eval()
tf = Compose([Resize((cfg["img_size"], cfg["img_size"])), ToTensor(),
Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
img = tf(Image.open("patch.jpg").convert("RGB")).unsqueeze(0)
prob = torch.softmax(model(img), 1)[0]
label = cfg["class_names"][prob.argmax().item()]
print(f"{label}: {prob.max().item():.3f}")
Grad-CAM overlays for each class and misclassified samples:
figures/gradcam_all_classes.png, figures/gradcam_misclassified.png.
Notes
- Images are mirrored as JPEG from the canonical BMP distribution (identical content).
- Research/educational use — not a production inspection system; validate per line/lighting.