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 cracks
inclusion — embedded foreign material
patches — surface deposits
pitted_surface — small pits / pitting corrosion
rolled-in_scale — mill scale rolled into the surface
scratches — 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).

Downloads last month: 311