Structured-DGM

September 14, 2026 · View on GitHub

Interactive inference demo for the paper Structured guided diffusion models for industrial defect image generation.
Paint a region on a defect-free product image, choose a defect type (and brush color), and synthesize an industrial defect with a guided diffusion model. Dual defects (two masks / two classes) are supported.

This repository is demo / inference only. Training code and full experimental pipelines are not released here (see Scope).

Demo

Full screen recording: assets/demo.mp4

Paper

Structured guided diffusion models for industrial defect image generation

Yulai Xie, Xiaoning Pi, Yang Zhang, Fang Ren
Knowledge-Based Systems, 2025, Volume 330, Article 114642

DOI: 10.1016/j.knosys.2025.114642
ScienceDirect: https://www.sciencedirect.com/science/article/abs/pii/S0950705125016818

Naming. The paper method is Structured-GDM (Structured Guided Diffusion Model). This repository is named Structured-DGM; it is the same project.

What the method controls

Industrial defect images differ from natural images: severe class imbalance, structured similarity, and the need for controllable diversity. Structured-GDM guides denoising diffusion so that three factors can be controlled independently:

Normal background — outline / reference guidance from a defect-free image
Defect class — classification guidance (e.g. crack / cut / hole / print on hazelnut)
Defect shape — segmentation / mask guidance over the painted region

The demo in this repo exposes class + shape (mask) control in the browser, with optional two-defect superposition (same idea as double guidance in the paper).

Citation

@article{XIE2025114642,
  title   = {Structured guided diffusion models for industrial defect image generation},
  author  = {Yulai Xie and Xiaoning Pi and Yang Zhang and Fang Ren},
  journal = {Knowledge-Based Systems},
  volume  = {330},
  pages   = {114642},
  year    = {2025},
  issn    = {0950-7051},
  doi     = {https://doi.org/10.1016/j.knosys.2025.114642},
  url     = {https://www.sciencedirect.com/science/article/pii/S0950705125016818}
}

Features (this repo)

Local Gradio UI: reference image picker, brush mask(s), defect type, color swatches
Single-defect and stacked two-defect generation
Device-portable sampling (cuda / cpu); no lab GPU whitelist, no MPI required for the demo path
Example checkpoints for MVTec AD hazelnut (crack / cut / hole / print)
Optional CLI wrappers (sample.py, RUN.sh) sharing the same sampler

Requirements

Python 3.10+ recommended (demo was validated on 3.10)
PyTorch matching your machine (CUDA build or CPU build) — install this before requirements.txt
Other deps: see requirements.txt (gradio, OpenCV, Pillow, NumPy, …)

# Example: install PyTorch yourself first, then:
pip install -r requirements.txt

Weights (download)

Checkpoint files are not stored in Git (large binaries). Download from Baidu Netdisk (Baidu Pan):

Item	Value
Shared folder name	`20250429_strucuted_diffusion` (keep the original spelling)
Link	https://pan.baidu.com/s/1ZGRRtwIjiTICug-vNX_yfQ?pwd=fj5t
Extraction code	`fj5t`

After download, place the three .pt files exactly here:

weights/dif_mvtec_hazelnut_c128_s256_new/model335000.pt          # diffusion
weights/cls_train_256x256_nazelnut/model100000.pt                # classifier
weights/seg_mvtec_nazelhut_all/case2_fcn_resnet50_T600_best_iou_62.2.pt  # segmentation guide

See also weights/README.md. If files are missing, python demo/app.py lists the missing relative paths on the page.

Note for overseas users: Baidu Netdisk is China-oriented. If the share link is hard to access from abroad, please open an issue and we may provide an alternate mirror later.

Quick start (Gradio demo)

python demo/app.py

Open http://127.0.0.1:7860 in a browser.

Typical workflow

Pick a hazelnut reference (good) image.
Paint the region where the defect should appear. Brush color is used as the fill / appearance cue.
Choose defect type: hole / print / crack / cut.
Click Generate. On a recent GPU this is about 1–2 minutes per image.
Optional: check Superimpose two defects to open a second canvas; paint type-1 on the left and type-2 on the right.

Brush color defaults

Type	Default brush color
crack	bright green
hole	black
print	white
cut	yellow

Palette also includes gray (and the colors above) so you can override the default.

Device options

CUDA_VISIBLE_DEVICES=0 python demo/app.py          # pick a GPU via env
python demo/app.py --device cuda:0
python demo/app.py --device cpu                    # works, but much slower
python demo/app.py --host 0.0.0.0 --port 7860      # LAN access

Unit tests (no GPU required):

python -m pytest tests/

Repository layout

Path	Role
`demo/app.py`	Gradio UI (paint + generate)
`demo/mask.py` / `demo/device.py`	Mask extraction, device resolution
`sdg/sampler.py`	Load models once, run one guided sample
`sdg/preset.py`	Hazelnut paths, labels, brush → `fill_color`
`sample.py` / `SAMPLE.sh` / `RUN.sh`	Optional CLI (same sampler; demo is the main path)
`assets/demo.gif` / `assets/demo.mp4`	Homepage GIF / full recording
`weights/`	Place downloaded checkpoints (see above)
`reference_images/`	Demo reference images (when present)
`tests/`	Lightweight tests without GPU

Command-line sampling (optional)

bash RUN.sh

RUN.sh calls SAMPLE.sh → sample.py. Sampling itself lives in sdg/sampler.py (demo path does not depend on MPI). Prefer CUDA_VISIBLE_DEVICES the same way as for the Gradio app.

Data & license notes

Example product images follow MVTec AD (hazelnut). Obtain the full dataset from MVTec under their terms.
Released code in this repository is for academic demonstration of the published method.
Weights are provided via Baidu Netdisk for reproducing the demo; redistribution of weights/data must respect MVTec and any institutional rules.

Scope / training code

This repository is an inference and interaction demo aligned with the paper’s controllable defect generation.

Training code is not published because of internal institutional regulations. We do not provide training scripts, full multi-class training recipes, or end-to-end experiment reproduction packages in this repository. Please use the paper for method details, and this demo to try guided sampling with the released hazelnut checkpoints.