Risk-Aware Surface Defect Segmentation and Automated Decision System

Deep learning pipeline for industrial surface defect segmentation with uncertainty-aware risk scoring and automated accept/reject decision-making.

---

🧠 Overview

This project implements a production-grade, end-to-end system for detecting and segmenting surface defects in steel manufacturing images, specifically targeting the Severstal Steel Defect Detection dataset. Beyond simple segmentation, it layers Monte Carlo Dropout uncertainty estimation and a risk engine on top of model predictions, culminating in an automated decision system that can classify each sample as ACCEPT, AUTO_REJECT, or HUMAN_REVIEW.

Key Highlights

• 🔬 Multi-class Segmentation – Simultaneously detects 4 defect classes using pixel-wise sigmoid outputs
• 🎲 Uncertainty Quantification – Monte Carlo Dropout with configurable forward passes quantifies prediction uncertainty
• ⚖️ Risk Engine – Combines defect area, class severity, and uncertainty into a composite risk score
• 🤖 Automated Decisions – Rule-based decision logic classifies samples without human intervention
• 🏗️ Dual Architecture – Supports both U-Net (ResNet-34 encoder) and SegFormer (MiT-B2 backbone)
• 📦 Kaggle-Ready – Auto-detects Kaggle environment, adapts paths and workers accordingly
• 📊 Analytics Suite – Calibration analysis, uncertainty diagnostics, and decision reporting

---

🗂️ Project Structure

├── main.py                  # CLI entry point (train / evaluate / predict / analyze)
├── config.py                # Centralized dataclass-based configuration
├── kaggle_train.py          # Kaggle notebook-compatible training script
├── requirements.txt         # Python dependencies
│
├── data/
│   ├── dataset.py           # PyTorch Dataset with patch extraction and augmentation
│   ├── dataloader.py        # DataLoader factory with stratified splits
│   ├── patching.py          # Patch-based tiling and reconstruction utilities
│   └── rle.py               # Run-Length Encoding / decoding helpers
│
├── models/
│   ├── unet.py              # U-Net with segmentation-models-pytorch (ResNet encoder)
│   └── segformer.py         # SegFormer with HuggingFace Transformers
│
├── training/
│   ├── train.py             # Two-stage training loop (warmup → full)
│   ├── losses.py            # Dice-BCE and Dice-Focal combined losses
│   ├── metrics.py           # Per-class Dice, IoU, pixel accuracy
│   └── evaluate.py          # Validation and test evaluation harness
│
├── uncertainty/
│   └── mc_dropout.py        # Monte Carlo Dropout inference engine
│
├── decision/
│   ├── risk_engine.py       # Risk score computation (area + severity + uncertainty)
│   └── decision_logic.py    # accept / reject / human-review thresholding
│
├── analytics/
│   └── calibration.py       # Reliability diagrams and ECE/ACE calibration metrics
│
├── visualization/
│   └── visualize.py         # Overlay masks, uncertainty maps, and decision summaries
│
├── scripts/
│   ├── analyze_dataset.py   # EDA: class distribution, defect statistics, sample viz
│   ├── test_training.py     # Quick smoke test for the training pipeline
│   └── test_kaggle.py       # End-to-end Kaggle compatibility test
│
├── checkpoints/             # Saved model weights (git-ignored)
├── logs/                    # TensorBoard logs (git-ignored)
└── outputs/                 # Predictions, reports, visualizations (git-ignored)

---

🏛️ Architecture

Segmentation Models

Model	Encoder	Parameters	Notes
U-Net	ResNet-34	~25M	Fast, reliable baseline
SegFormer	MiT-B2	~25M	Transformer-based, stronger on texture

Both models output 4-channel sigmoid maps (one per defect class), enabling multi-label pixel classification.

Two-Stage Training

1. Warmup Stage (10 epochs) – Trains on a defect-heavy subset (90% defect : 10% clean) to teach the model defect patterns before class imbalance dominates
2. Full Training Stage (40 epochs) – Fine-tunes on the full dataset with cosine LR scheduling and early stopping

Uncertainty Estimation

Monte Carlo Dropout is applied at inference time by keeping dropout layers active over N stochastic forward passes. The pixel-wise variance across runs serves as an uncertainty map, surfacing regions where the model is unsure.

Uncertainty = σ²(predictions over N forward passes)

Risk Scoring

Each prediction is scored by a weighted combination of:

Factor	Weight	Description
Defect Area	0.4	Percentage of pixels flagged as defective
Uncertainty	0.3	Mean uncertainty across defect regions
Class Severity	0.3	Per-class severity multipliers [1.0, 2.0, 0.8, 1.5]

Decision Logic

Risk Score → ACCEPT         if risk < 0.20
Risk Score → AUTO_REJECT    if defect_area > 5% AND uncertainty < 0.30
Risk Score → HUMAN_REVIEW   if uncertainty > 0.30

---

⚙️ Configuration

All hyperparameters are centralized in config.py via Python dataclasses — no magic numbers anywhere else.

from config import get_config, ModelConfig

# Default config (auto-detects Kaggle vs local)
config = get_config()

# Switch to SegFormer
config = get_config(model=ModelConfig(model_type="segformer"))

Key config groups:

Config Class	Controls
DataConfig	Paths, patch size, augmentation, class balance
ModelConfig	Architecture, encoder, dropout
TrainingConfig	LR, scheduler, loss type, AMP, early stopping
UncertaintyConfig	MC passes, uncertainty threshold
RiskConfig	Area/severity/uncertainty weights
DecisionConfig	Accept/reject/review thresholds

Environment variable overrides are supported for all paths:

STEEL_DATA_ROOT=/path/to/data
STEEL_OUTPUT_DIR=/path/to/outputs
STEEL_CHECKPOINT_DIR=/path/to/checkpoints
STEEL_LOG_DIR=/path/to/logs

---

🚀 Quick Start

1. Install Dependencies

pip install -r requirements.txt

2. Prepare the Dataset

Download the Severstal Steel Defect Detection dataset and place it as:

data/raw/
  train.csv
  train_images/
  test_images/

Or set STEEL_DATA_ROOT to point to your data directory.

3. Run Dataset Analysis

python scripts/analyze_dataset.py

4. Train

# Train U-Net (both stages)
python main.py train --model unet --stage both

# Train SegFormer (full stage only)
python main.py train --model segformer --stage full

# Resume from checkpoint
python main.py train --model unet --resume checkpoints/best.pt

5. Evaluate

python main.py evaluate --model unet --checkpoint checkpoints/best.pt

6. Predict with Risk Assessment

python main.py predict --model unet --checkpoint checkpoints/best.pt --mc-passes 20

7. Generate Analytics

python main.py analyze --checkpoint checkpoints/best.pt

---

🧪 Kaggle Usage

The system auto-detects the Kaggle environment. Simply run kaggle_train.py as your kernel script — paths and worker counts are adjusted automatically.

# In Kaggle notebook
!python kaggle_train.py

Smoke test before full training:

python scripts/test_kaggle.py

---

📦 Dependencies

Package	Purpose
torch >= 2.0	Core deep learning framework
segmentation-models-pytorch	U-Net encoder-decoder
transformers >= 4.30	SegFormer backbone
albumentations	Image augmentation
scikit-learn	Stratified splitting, metrics
numpy, pandas	Data manipulation
opencv-python	Image I/O and processing
matplotlib, seaborn	Visualization
tensorboard	Training monitoring

---

📈 Results

Note: This is an active research project. Results will be updated as training progresses.

Evaluation metrics tracked per defect class:

• Dice Coefficient (primary metric)
• IoU (Jaccard Score)
• Pixel Accuracy

---

📁 Dataset

Severstal Steel Defect Detection — Kaggle Competition
Images: 1600×256 grayscale steel strip scans
Labels: 4 defect classes encoded as Run-Length Encoding (RLE) masks
Train set: ~12,568 images | Test set: ~1,801 images

---

📝 License

This project is open-source and released under the MIT License.

---

👤 Author

sharksurfauto-byte
GitHub: @sharksurfauto-byte