---
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)