---
license: mit
tags:
  - pytorch
  - image-classification
  - defect-detection
  - semiconductor
  - wafer-map
  - quality-control
  - computer-vision
datasets:
  - wm811k
metrics:
  - accuracy
  - f1
model-index:
  - name: radai-wm811k-defect-detection
    results:
      - task:
          type: image-classification
          name: Wafer Defect Classification
        dataset:
          name: WM-811K
          type: wm811k
        metrics:
          - type: accuracy
            value: 95.25
          - type: f1
            value: 95.28
---

# RadAI WM-811K Wafer Defect Detection

A ResNet-based model for semiconductor wafer map defect classification, trained on the **full** WM-811K dataset.

## Model Description

This model classifies semiconductor wafer map defects into 8 categories. Unlike many published results that use small cherry-picked subsets, this model was trained on the complete WM-811K dataset with real-world class imbalance.

**Key differentiator:** Trained on 172,000+ samples (100% of labeled data), not 902 cherry-picked samples (0.5%) as seen in some publications.

## Training Results

| Metric | Value |
|--------|-------|
| **Validation Accuracy** | 95.25% |
| **Validation F1 Score** | 95.28% |
| **Training Epochs** | 76 (early stopping) |
| **Training Time** | ~5 hours (CPU) |

## Dataset

**WM-811K (LSWMD.pkl)**
- Total wafer maps: 811,457
- Labeled defect samples: ~172,000
- Classes: 8 defect types

### Class Distribution

| Class | Samples | Percentage |
|-------|---------|------------|
| Center | 4,294 | 2.5% |
| Donut | 555 | 0.3% |
| Edge-Loc | 5,189 | 3.0% |
| Edge-Ring | 9,680 | 5.6% |
| Loc | 3,593 | 2.1% |
| Random | 866 | 0.5% |
| Scratch | 1,193 | 0.7% |
| Near-full | 149 | 0.1% |

## Model Architecture

- **Base:** ResNet34 (pretrained on ImageNet)
- **Input:** 64x64 grayscale wafer maps
- **Output:** 8 defect classes
- **Modifications:** 
  - First conv layer adapted for 1-channel input
  - Custom classifier head with dropout

## Usage

```python
import torch
import torch.nn as nn
from torchvision import models
import numpy as np
from PIL import Image
from scipy.ndimage import zoom

# Define model architecture
class RadAI_ResNet(nn.Module):
    def __init__(self, num_classes=8):
        super().__init__()
        self.base = models.resnet34(weights=None)
        self.base.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.base.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(self.base.fc.in_features, num_classes)
        )
    
    def forward(self, x):
        return self.base(x)

# Load model
model = RadAI_ResNet(num_classes=8)
checkpoint = torch.load('best_radai_resnet.pt', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Class labels
CLASSES = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Random', 'Scratch', 'Near-full']

# Predict on a wafer map
def predict(wafer_map):
    """
    wafer_map: 2D numpy array (any size)
    returns: predicted class name, confidence
    """
    # Resize to 64x64
    h, w = wafer_map.shape
    resized = zoom(wafer_map, (64/h, 64/w), order=1)[:64, :64]
    
    # Normalize
    if resized.max() > 0:
        resized = resized / resized.max()
    
    # To tensor
    tensor = torch.FloatTensor(resized).unsqueeze(0).unsqueeze(0)
    
    # Predict
    with torch.no_grad():
        output = model(tensor)
        probs = torch.softmax(output, dim=1)
        pred_idx = output.argmax(1).item()
        confidence = probs[0, pred_idx].item()
    
    return CLASSES[pred_idx], confidence

# Example
# prediction, conf = predict(your_wafer_map)
# print(f"Predicted: {prediction} ({conf*100:.1f}%)")
```

## Training Configuration

| Parameter | Value |
|-----------|-------|
| Optimizer | AdamW |
| Learning Rate | 1e-3 |
| Weight Decay | 1e-4 |
| Batch Size | 64 |
| Scheduler | OneCycleLR |
| Loss | CrossEntropyLoss (weighted) |
| Augmentation | Flip, Rotate, Noise |
| Early Stopping | Patience 15 |

## Comparison with Published Results

| Method | Dataset Size | Accuracy |
|--------|--------------|----------|
| Published Paper (2024) | 902 (cherry-picked) | 91.7% |
| **RadAI (this model)** | **172,000+ (full)** | **95.25%** |
| Nakazawa (2018) | Full | 98.2% |
| Shawon (2019) | Full | 99.29% |

## Intended Use

- Semiconductor wafer defect classification
- Quality control automation
- Research and benchmarking
- Transfer learning base for similar domains (CdTe, CZT radiation detectors)

## Limitations

- Trained on silicon wafer data (WM-811K)
- May require fine-tuning for other materials
- CPU-trained (GPU training may improve results)
- Class imbalance affects rare defect detection

## Future Work

- GPU-accelerated training for higher accuracy
- Transfer learning to CdTe/CZT radiation detector defects
- Real-time inference optimization

## Citation

If you use this model, please cite:

```bibtex
@misc{radai-wm811k-2024,
  author = {RadAI},
  title = {WM-811K Wafer Defect Detection Model},
  year = {2024},
  publisher = {Hugging Face},
  url = {https://huggingface.co/radai/wm811k-defect-detection}
}
```

## Links

- **Website:** [rad-ai.org](https://rad-ai.org)
- **Dataset:** [WM-811K on Kaggle](https://www.kaggle.com/datasets/qingyi/wm811k-wafer-map)

## License

MIT License

---

*Built with passion, evenings & weekends. RadAI - Real AI for Real Problems.*