README.md

---
license: other
license_name: intel-challenge-dataset
license_link: LICENSE
tags:
  - pytorch
  - image-classification
  - efficientnet
  - defect-detection
  - defect-classification
  - semiconductor
  - wafer-inspection
  - manufacturing
  - few-shot-learning
  - small-sample-learning
  - computer-vision
library_name: pytorch
pipeline_tag: image-classification
metrics:
  - accuracy
  - f1
model-index:
  - name: defect-vision-efficientnet-b2
    results:
      - task:
          type: image-classification
        metrics:
          - type: accuracy
            value: 0.9556
            name: Test accuracy
          - type: accuracy
            value: 0.975
            name: Best validation accuracy
---

# Defect Vision: EfficientNet-B2 Semiconductor Defect Classifier

Fine-tuned EfficientNet-B2 for **small-sample wafer defect classification**, built for the **Intel
Semiconductor Solutions Challenge 2026, Problem A: Small-Sample Learning for Defect Classification**.

Classifies gray-scale wafer/die images into **8 defect classes + "no defect"** (9-way), trained on a
class-balanced, heavily-augmented small dataset rather than large-scale labeled data. The challenge's core
constraint is that production defect data is scarce and imbalanced.

- **Code, FastAPI service, React demo UI, training notebook:** https://github.com/Sehastrajit-S/defect-vision
- **Backbone:** `torchvision.models.efficientnet_b2` (ImageNet-pretrained), custom classifier head
- **Params:** ~9.2M
- **Input:** 260×260 RGB (gray-scale images converted to 3-channel), ImageNet normalization

## Results

| Metric | Target (challenge brief) | Achieved |
|---|---|---|
| Overall classification accuracy | ~85% | **95.6%** (test, 360 held-out images) |
| Best validation accuracy | n/a | **97.5%** |
| Inference latency | ~1s/image | ~40–500ms/image (GPU), ~0.1–1s (CPU) |

<details>
<summary>Full per-class classification report (test set)</summary>

```text
Test Loss : 0.6153  |  Test Accuracy : 0.9556

              precision    recall  f1-score   support

     defect1     0.9773    0.9556    0.9663        45
     defect2     0.9375    1.0000    0.9677        45
     defect3     1.0000    1.0000    1.0000        45
     defect4     1.0000    1.0000    1.0000        45
     defect5     0.9130    0.9333    0.9231        45
     defect8     0.8837    0.8444    0.8636        45
     defect9     0.9556    0.9556    0.9556        45
    defect10     0.9773    0.9556    0.9663        45
    new_good     0.0000    0.0000    0.0000         0

    accuracy                         0.9556       360
   macro avg     0.8494    0.8494    0.8492       360
weighted avg     0.9555    0.9556    0.9553       360
```

`new_good` (no defect) has zero held-out samples in this dataset revision. The 9th output neuron is reserved
for future "no defect found" imagery without requiring re-architecture.

</details>

![Confusion matrix](confusion_matrix_test.png)
![Training curves](training_curves.png)

## Handling class imbalance with few samples

- **Class-balanced dataset construction**: equal train/val/test counts per class (210/45/45) via augmentation,
  instead of naive minority oversampling or loss reweighting, so the model never learns a majority-class prior.
- **Aggressive augmentation**: random crop, flips, rotation, perspective warp, and color jitter multiply the
  small per-class sample count without duplicating exact pixels.
- **Label smoothing (0.1)** on cross-entropy keeps the model from over-committing on visually similar defect
  types.
- **OneCycleLR + early stopping** (patience 7) for fast, stable convergence on limited data. This checkpoint
  converged and early-stopped at epoch 16.

## Usage

```python
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
from huggingface_hub import hf_hub_download

CLASSES = ["defect1", "defect2", "defect3", "defect4", "defect5",
           "defect8", "defect9", "defect10", "new_good"]

def build_model(num_classes: int) -> nn.Module:
    model = models.efficientnet_b2(weights=None)
    in_f = model.classifier[1].in_features
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.4),
        nn.Linear(in_f, 512),
        nn.SiLU(inplace=True),
        nn.Dropout(p=0.3),
        nn.Linear(512, num_classes),
    )
    return model

weights_path = hf_hub_download(repo_id="Sehastrajit/defect-vision-efficientnet-b2", filename="best_model.pth")
model = build_model(len(CLASSES))
ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state"])
model.eval()

transform = transforms.Compose([
    transforms.Resize((260, 260)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

img = Image.open("wafer_sample.png").convert("RGB")
x = transform(img).unsqueeze(0)
with torch.no_grad():
    probs = torch.softmax(model(x), dim=1)[0]

pred = CLASSES[probs.argmax().item()]
print(pred, probs.max().item())
```

## Training setup

| | |
|---|---|
| GPU | NVIDIA RTX 3060 12GB (fp16 AMP) |
| Optimizer | AdamW, lr 2e-4, weight decay 1e-4 |
| Schedule | OneCycleLR, cosine anneal |
| Batch | 64 × 2 grad-accum steps (effective 128) |
| Split | 70% train / 15% val / 15% test |
| Epochs | early-stopped at 16 (patience 7) |

Full training script: [`h1.ipynb`](https://github.com/Sehastrajit-S/defect-vision/blob/main/src/app/h1.ipynb) in
the main repo.

## Intended use & limitations

Built as a challenge submission demonstrating small-sample defect classification technique, not validated for
production fab deployment. Trained on Intel-provided sample imagery for the Semiconductor Solutions Challenge
2026; `new_good` has no held-out evaluation samples in this dataset revision. Intel and the Intel logo are
trademarks of Intel Corporation or its subsidiaries. This is an independent student project, not an Intel
product.