Casting Defect Detection (ViT)
Fine-tuned ViT-B/16 (google/vit-base-patch16-224) on 6,633 casting surface images to classify submersible pump impeller parts as defective or ok.
Metrics
| Metric | Value |
|---|---|
| Accuracy | 99.6% |
| Macro F1 | 0.995 |
| Training images | 6,633 |
| Test images | 715 |
| Epochs | 3 |
| Training time | ~10 min on RTX 3060 |
Usage
from transformers import pipeline
classifier = pipeline("image-classification", model="gfichetdc/casting-defect-vit")
result = classifier("path/to/casting_image.jpeg")
# [{'label': 'ok', 'score': 0.998}, {'label': 'defective', 'score': 0.002}]
Training
- Base model:
google/vit-base-patch16-224(ImageNet pretrained) - Replaced the 1000-class head with a 2-class head (defective / ok)
- Fine-tuned full model for 3 epochs, lr=2e-5, batch size 16
- Experiment tracking with MLflow
- Best checkpoint selected by F1
Dataset
Kaggle — Casting Product Image Data for Quality Inspection: 7,348 grayscale images (512×512) of submersible pump impellers, split into defective (~57%) and ok classes.
Demo
Try the model live: HuggingFace Space