defect-vision-efficientnet-b2

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

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

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>

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

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

Full training script: 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.