JFenic/industrial-defect-segmentation-ksdd2

Industrial Defect Segmentation (KolektorSDD2)

Pretrained semantic segmentation models from the Master's Thesis:

Vision Transformers vs CNNs for Semantic Segmentation of Industrial Defects

• Author: José Fenic Peiteado Padilla
• Code repository: GitHub — TFM-Computer-Vision-Segmentation-ConvNext-Swin
• Intended use: Academic and research purposes

Models included

Folder	Architecture	Encoder	Decoder	Params
convnext/	ConvNeXtSeg	ConvNeXt-Tiny (timm)	UPerNet-style (PPM + FPN)	~33.7M
swin/	SwinSeg	Swin-Tiny (timm)	UPerNet-style (PPM + FPN)	~33.4M

Both models were trained on KolektorSDD2 with input size 224×640 (H×W), binary defect segmentation (num_classes=1), and ImageNet normalization.

Test set metrics (threshold = 0.5)

Model	Dice (%)	IoU (%)	Precision (%)	Recall (%)	Latency (ms/img)
ConvNeXt-Tiny	27.41	15.88	93.48	16.06	21.96
Swin-Tiny	33.47	20.10	93.04	20.41	32.19

Dataset (not included)

These weights were trained on KolektorSDD2 / KSDD2, licensed under CC BY-NC-SA 4.0. The dataset is not redistributed in this Hub repository.

Please download the dataset from the official source and cite:

@article{tabernik2020segmentation,
  title={Segmentation-Based Deep-Learning Approach for Surface-Defect Detection},
  author={Tabernik, Domen and {\v{S}}tern, Danijel and Sko{\v{c}}aj, Danijel},
  journal={Journal of Intelligent Manufacturing},
  year={2020},
  publisher={Springer}
}

Loading weights

Clone the GitHub repository for model definitions in src/models.py, then download checkpoints from this Hub repo. Both safetensors (recommended) and standard PyTorch .pth are available:

from pathlib import Path
import sys
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import torch

PROJECT_ROOT = Path(".").resolve()
sys.path.insert(0, str(PROJECT_ROOT / "src"))
from models import ConvNeXtSeg, SwinSeg

def load_model(model_cls, subfolder: str, device="cpu", use_safetensors: bool = True):
    if use_safetensors:
        weights = hf_hub_download(
            repo_id="JFenic/industrial-defect-segmentation-ksdd2",
            filename=f"{subfolder}/model.safetensors",
        )
        state = load_file(weights, device=device)
    else:
        weights = hf_hub_download(
            repo_id="JFenic/industrial-defect-segmentation-ksdd2",
            filename=f"{subfolder}/pytorch_model.pth",
        )
        state = torch.load(weights, map_location=device, weights_only=False)
        if isinstance(state, dict) and "model_state_dict" in state:
            state = state["model_state_dict"]

    model = model_cls(pretrained=False, num_classes=1)
    model.load_state_dict(state)
    model.eval()
    return model

convnext = load_model(ConvNeXtSeg, "convnext")
swin = load_model(lambda **kw: SwinSeg(**kw, img_size=(224, 640)), "swin")

Inference notes

• Resize input to 224×640
• Normalize with ImageNet mean (0.485, 0.456, 0.406) and std (0.229, 0.224, 0.225)
• Apply sigmoid(logits) > threshold (default threshold: 0.5; validated alternative: 0.1)

See the Gradio demo in the GitHub repository.

Repository layout

convnext/
  model.safetensors
  pytorch_model.pth
  config.yaml
swin/
  model.safetensors
  pytorch_model.pth
  config.yaml
model_index.json
LICENSE
README.md

License

• Model weights and this Hub card: shared for academic/research use alongside the MIT-licensed code repository.
• KolektorSDD2 dataset: CC BY-NC-SA 4.0 — non-commercial terms apply to dataset usage.

Commercial use of the dataset or derived materials may require permission from the dataset authors.