gfichetdc/casting-defect-vit

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