metadata
license: mit
tags:
- image-classification
- vision
- vit
- casting
- defect-detection
- quality-inspection
datasets:
- ravirajsinh45/real-life-industrial-dataset-of-casting-product
metrics:
- accuracy
- f1
model-index:
- name: casting-defect-vit
results:
- task:
type: image-classification
metrics:
- name: Accuracy
type: accuracy
value: 0.996
- name: Macro F1
type: f1
value: 0.995
pipeline_tag: image-classification
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