| --- |
| 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 |
|
|
| ```python |
| 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](https://www.kaggle.com/datasets/ravirajsinh45/real-life-industrial-dataset-of-casting-product): 7,348 grayscale images (512×512) of submersible pump impellers, split into defective (~57%) and ok classes. |
|
|
| ## Demo |
|
|
| Try the model live: [HuggingFace Space](https://huggingface.co/spaces/gfichetdc/casting-defect-vit) |
|
|
