pranayr710/Self-Supervised-Defect-Detection

Self-Supervised Defect Detection (PatchCore-Lite)

An end-to-end "Label-Free" computer vision pipeline for detecting unknown surface defects, training exclusively on images of perfect parts.

🧠 Core Concept

The model never sees a single defect during training. Instead, it memorizes the spatial feature distribution of normal surfaces. At test time, any region whose features deviate significantly from this learned distribution is flagged as anomalous.

In short: Defects = deviations in feature space, not a classification problem.

🔬 Why PatchCore-Lite Works

📁 Project Structure

project/
├── config.py          # Device, paths, backbone, hyperparameters
├── dataset.py         # MVTec AD PyTorch Dataset + augmentations
├── model.py           # Frozen feature extractor (ResNet18 / DINOv2)
├── train.py           # Builds and saves the coreset memory bank
├── inference.py       # Single-image / folder heatmap generation + CSV export
├── evaluate.py        # Image-level & pixel-level AUROC evaluation
├── calibrate.py       # Optimal threshold via Youden J statistic
├── app.py             # Gradio web UI for interactive inspection
├── utils.py           # Distance computation, heatmap rendering
├── requirements.txt   # Python dependencies
├── results/           # Generated heatmaps & score distribution plot
└── README.md

⚙️ Setup

1. Install Dependencies

pip install -r requirements.txt

2. Download Dataset

Download the MVTec AD Dataset and extract the hazelnut category so the structure looks like:

dataset/
└── hazelnut/
    ├── train/
    │   └── good/          ← only normal images
    └── test/
        ├── good/          ← normal test images
        ├── crack/         ← defect type 1
        ├── cut/           ← defect type 2
        ├── hole/          ← defect type 3
        └── print/         ← defect type 4

🚀 Usage

Step 1 — Train (Build Memory Bank)

python train.py

Extracts patch features from all good images and saves memory_bank.pt.

Step 2 — Evaluate (Compute AUROC)

python evaluate.py

Runs all test images and prints the Image-Level AUROC score.

Step 3 — Inference (Single Image)

python inference.py --image dataset/hazelnut/test/crack/000.png
python inference.py --image dataset/hazelnut/test/crack/000.png --output my_heatmap.png

Outputs:

• —Anomaly score (printed to console)
• —Heatmap image saved to output_heatmap.png (or custom path)

📊 Performance

🖼️ Example Results

Generate these files first by running: `` python inference.py --image dataset/hazelnut/test/good/000.png --output results/heatmap_good.png python inference.py --image dataset/hazelnut/test/crack/000.png --output results/heatmap_crack.png python inference.py --image dataset/hazelnut/test/hole/000.png --output results/heatmap_hole.png ``

Good Image — Low Activation (No Defect)

Input: test/good/000.png
Anomaly Score: 0.7778
Result: Uniform cool blue — no anomaly detected.

[image] Normal hazelnut: Low-intensity blue heatmap (no red regions)

Defect Image — Crack Detected

Input: test/crack/000.png
Anomaly Score: 0.8465
Result: Bright red concentrated on the crack location.

[image] Hazelnut with crack: Intense red overlay precisely on the crack

Defect Image — Hole Detected

Input: test/hole/000.png
Anomaly Score: 0.8681
Result: Sharp red hotspot on the hole area.

[image] Hazelnut with hole: Focused red highlight on the hole

🛡️ Key Implementation Details

• —Feature Normalization: L2 normalization applied before memory bank storage and before distance computation to ensure scale-invariant matching.
• —Balanced Feature Fusion: layer3 features are scaled by 0.5 before concatenation with layer2, preventing deeper semantic features from dominating fine-grained texture signals.
• —Batched Distance Computation: torch.cdist is run in configurable batches (default 1024) to prevent OOM on standard GPUs.
• —Safe Heatmap Normalization: Division-by-zero guarded with 1e-8 epsilon.
• —Cross-Device Compatibility: Memory bank is saved on CPU and loaded to the correct device at runtime.
• —Reproducibility: Fixed seeds (torch.manual_seed(42), np.random.seed(42)).

📜 License

This project is for educational and research purposes.