TinyViT for Surface-Defect Classification (NEU)
A tiny Vision Transformer trained from scratch (no pretraining, no ImageNet) that classifies steel surface-defect images into six defect classes. The point of the project is to show a ViT can be built for manufacturing and trained on CPU: 64x64 image -> 64 patches of 8x8 -> 2-layer transformer encoder -> class, ~113k parameters, trained on CPU in minutes.
Classes
crazing— fine surface cracksinclusion— embedded foreign materialpatches— surface depositspitted_surface— small pits / pitting corrosionrolled-in_scale— mill scale rolled into the surfacescratches— linear abrasions
Results
| Model | Test accuracy | Macro F1 |
|---|---|---|
| TinyViT (this model) | 76.11% | 76.21% |
| Small CNN (same 64×64 input, honest baseline) | 85.00% | 84.99% |
Per-class F1 (TinyViT): crazing 85.71% · patches 91.23% · pitted_surface 75.86% · inclusion 68.85% · scratches 68.97% · rolled-in_scale 66.67%.
Honest reading: the small CNN wins by ~9 points — convolutional inductive bias beats a from-scratch transformer on a 1,440-image dataset. The ViT is still the point of the project: 113k parameters, no pretraining, CPU-trained in ~14 minutes, and 76% is far above the 16.7% chance level. On larger industrial datasets (or with a conv stem) the gap closes; see "What I'd try next" in the notebook.
Usage
import numpy as np
from PIL import Image
from model import TinyViT, load_model, predict_image, CLASS_NAMES
model = load_model("pytorch_model.bin") # CPU by default
img = Image.open("steel_surface.jpg") # any size, gray or RGB
cls_id, probs = predict_image(model, img)
print(CLASS_NAMES[cls_id], f"{probs[cls_id]:.1%} confidence")
Preprocessing contract (must match training): resize to 64x64 grayscale,
scale to [0, 1], normalise with train-set mean/std. predict_image does this
for you.
Training details
- Data: NEU surface defect database (Northeastern University), 1,800 grayscale 200x200 images, 300 per class; stratified 80/10/10 split.
- Augmentation: random horizontal/vertical flips, ±15° rotation, resized crop (0.85–1.0) — orientation-safe for surface textures.
- Optimiser: AdamW (lr 3e-4, cosine schedule), early stopping on validation.
- Full reproducible pipeline (data loader, training, CNN baseline, figures): see the project repository.
Limitations
- Trained on lab-curated, evenly-lit 200x200 images of one material family (hot-rolled steel strip). Real production lines have different illumination, scale, motion blur, and defect mixtures.
- Grayscale only; colour/anomaly information is discarded.
- 64x64 input loses fine detail — crazing hairlines near the resolution limit are the hardest class (see per-class F1 in the report).
Citation / dataset credit
- Northeastern University surface defect database: http://faculty.neu.edu.cn/yunhyan/NEU_surface_defect_database.html
- Mirror used for this project: https://github.com/siddhartamukherjee/NEU-DET-Steel-Surface-Defect-Detection
Built as a portfolio project for graduate applications in intelligent manufacturing / surface quality inspection (mechanical engineering × AI).