Casting Defect Detection (ViT)

Fine-tuned a Vision Transformer on casting surface images to classify defective vs. ok parts. Model hosted on HuggingFace Hub.

githubhuggingface

Goal

Detect surface defects on metal casting parts using a fine-tuned Vision Transformer, with experiment tracking and a serving endpoint.

Approach

  1. Load the Kaggle casting inspection dataset (6,633 train / 715 test images of pump impellers, ~57% defective).
  2. Preprocess with the ViT image processor: resize to 224×224, normalize to ImageNet statistics.
  3. Fine-tune google/vit-base-patch16-224 for 3 epochs on an RTX 3060, tracking everything with MLflow.
  4. Evaluate on the held-out test set. Serve predictions via FastAPI.

Guiding questions

  • Can a pretrained ViT reach near-perfect accuracy on a small industrial dataset with minimal tuning?
  • How fast does the model converge — is 3 epochs enough, or does it overfit?
  • Is the model confident on both classes, or does it hedge on borderline samples?

Data

Headline numbers

Training images6,633
Test images715
Macro F10.995
Accuracy99.6%
Base modelViT-B/16224px
Epochs3~10 min on RTX 3060

Model architecture

Image512×512
→
Resize224×224
→
Patch16×16 ×196
→
ViT Encoder12 layers
→
CLS Head2 classes
→
Outputok / defect
1
LoadHuggingFace datasets

Load 6,633 train + 715 test images from folder structure. Labels assigned from directory names (ok_front / def_front).

2
PreprocessViTImageProcessor

Resize 512×512 casting photos to 224×224, normalize pixel values to ViT input distribution.

3
Fine-tuneHuggingFace Trainer

Replace ImageNet classification head (1000 → 2 classes). Fine-tune full model for 3 epochs, lr=2e-5, batch size 16.

4
TrackMLflow

Log hyperparameters, loss curves, and per-epoch F1/accuracy. Compare runs to pick the best checkpoint.

5
ServeFastAPI

REST endpoint accepts image upload, returns defect/ok prediction with confidence score.

6
PublishHuggingFace Hub

Push fine-tuned weights to Hub. Anyone can load the model with one line of code.

Live demo

Upload a casting image or use the examples — the model runs on HuggingFace Spaces.

Conclusion

  • A pretrained ViT-B/16 reaches 99.5% macro F1 after just 3 epochs of fine-tuning on ~6,600 images — transfer learning makes small industrial datasets practical.
  • The full pipeline covers fine-tuning (HuggingFace Trainer), experiment tracking (MLflow), model serving (FastAPI), and model sharing (HuggingFace Hub).
PythonPyTorchHuggingFaceVision TransformerMLflowFastAPIComputer Vision