steel-defect-resnet18

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)

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.