Krrrrish/industrial-surface-defect-segmentation

🏭 Industrial Surface Defect Segmentation

Pixel-level defect detection and multi-class segmentation for industrial steel and metal strip surfaces ($128 \times 800$ resolution), derived from the end-to-end training and calibration pipeline in dl-24f3005189-notebook-t32026.ipynb.

---

📌 Model Summary

Automated surface inspection is crucial in high-speed manufacturing to detect defects before steel coils undergo further processing. This repository provides:

• Pretrained architectures: Custom U-Net with ResNet-34 and EfficientNet-B3 single-channel backbones.
• Trained with BCEDiceLoss: Balanced positive-pixel weighting and soft Dice loss to handle severe class imbalance.
• Post-processing & Calibration: Per-class probability thresholds and 8-connected component minimum area filters tuned to the competition metric.
• 4-View Test-Time Augmentation (TTA): Multi-flip averaging for resilient boundary localization.

---

🏷️ Defect Classes

Class ID	Defect Category	Description	Calibrated Threshold	Min Component Area
defect_a	Pitted / Scratches	Elongated narrow scratches or localized marks	0.35	$25\text{ px}$
defect_b	Minor defects	Clustered small pits and surface anomalies	0.22	$25\text{ px}$
defect_c	Inclusions	Slag/oxide inclusions and dark irregular masses	0.42	$220\text{ px}$
defect_d	Surface patches	Large low-contrast scuffing or rolling patches	0.42	$350\text{ px}$

---

📊 Benchmark Validation Results

Evaluated with the official competition metric:

$$\text{Score}=0.70\times {\text{Dice}}_{\text{foreground}}+0.20\times {\text{Score}}_{\text{FP}}+0.10\times {\text{Acc}}_{\text{empty}}\backslash text\{Score\}\; =\; 0.70\; \backslash times\; \backslash text\{Dice\}\_\{\backslash text\{foreground\}\}\; +\; 0.20\; \backslash times\; \backslash text\{Score\}\_\{\backslash text\{FP\}\}\; +\; 0.10\; \backslash times\; \backslash text\{Acc\}\_\{\backslash text\{empty\}\}$$

Model Architecture	Competition Score	Foreground Dice	False Positive Score	Empty Accuracy
EfficientNet-B3 U-Net	0.7056	0.6081	0.9333	0.9333
ResNet-34 U-Net	0.6701	0.5654	0.9144	0.9144
Ensemble (Top-2)	0.7000	0.5900	0.9566	0.9566
Scratch U-Net (Baseline)	0.5739	0.4481	0.8675	0.8675

---

🚀 Quickstart: Loading & Inference

1. Installation

pip install torch torchvision opencv-python-headless numpy pillow

2. Run Inference in Python

from models_def import load_model_with_weights
from utils import preprocess_image, predict_with_tta, apply_postprocess, create_segmentation_overlay

# 1. Load model architecture
model, status = load_model_with_weights("resnet34_unet", device="cpu")

# 2. Preprocess input image (converts to grayscale, resizes to 128x800, normalizes)
tensor_x, display_rgb = preprocess_image("examples/sample_02_scratch_defect_a.png")

# 3. Predict with 4-view TTA
probs = predict_with_tta(model, tensor_x, tta=True, device="cpu")

# 4. Apply calibrated post-processing (thresholds & connected-component min area)
thresholds = [0.35, 0.22, 0.42, 0.42]
min_areas = [25, 25, 220, 350]
masks = apply_postprocess(probs, thresholds, min_areas)

# 5. Generate high-contrast segmentation overlay
overlay = create_segmentation_overlay(display_rgb, masks)

---

💻 Interactive Gradio App

An interactive web UI (app.py) is included. Run it locally:

pip install gradio
python app.py

---

📁 Repository Contents

• models_def.py: Complete PyTorch architecture definitions (ResNet34UNet, EffNetUNet, ScratchUNet).
• utils.py: Normalization, TTA, connected-component filtering, and Kaggle RLE encoder.
• app.py: Standalone Gradio web demo application.
• inference_demo.py: Ready-to-run inference script.
• notebook_industrial_defect_segmentation.ipynb: The full training and experiment tracking notebook.
• examples/: 5 realistic test strip images ($128 \times 800$).

---

👤 Attribution & Experiment Metadata

• Notebook: dl-24f3005189-notebook-t32026.ipynb
• Roll Number: 24f3005189
• Experiment Tracking: Weights & Biases (24f3005189-t32026)