MaskDriven-SD

Mask-Driven Synthetic Dataset Generation with Stable Diffusion, ControlNet, and TensorRT

MaskDriven-SD is an experimental synthetic data generation pipeline for creating paired defect images and segmentation masks using diffusion models.

The project procedurally generates crack masks and uses them as structural conditioning for diffusion models such as Stable Diffusion 3.5 + ControlNet, producing realistic synthetic images whose defect geometry follows the generated ground-truth mask.

The main use case explored in this repository is synthetic crack generation on marble surfaces, with the goal of generating scalable training data for computer vision tasks such as segmentation and defect detection.

---

Motivation

Collecting and manually annotating defect datasets can be expensive and time-consuming.

MaskDriven-SD reverses the usual annotation workflow:

1. Generate the annotation first.
2. Use that annotation as structural guidance for a generative model.
3. Generate a realistic image that follows the annotation.

This makes it possible to create image-mask pairs automatically:

Procedural Mask
      │
      ▼
Dynamic Prompt Generation
      │
      ▼
Stable Diffusion / ControlNet
      │
      ▼
Synthetic Defect Image
      │
      ▼
Paired Image + Ground-Truth Mask

Instead of generating an image and trying to recover its segmentation mask afterward, the desired geometry is defined before generation.

---

Features

• Procedural crack-mask generation
• Multiple crack geometry types
• Stable Diffusion 3.5 + ControlNet generation
• TensorRT-accelerated diffusion inference
• Dynamic prompt generation
• Deterministic prompt variation based on mask filenames
• Optional negative prompts
• Batch dataset generation
• Automatic image-mask filename pairing
• FLUX experimentation
• Configurable generation resolution and ControlNet strength

---

Pipeline

1. Procedural Ground-Truth Generation

generate_seg_masks/ground_truth_masks.py generates synthetic crack masks using OpenCV and NumPy.

Several crack patterns are implemented:

• hairline
• impact
• fragment
• edge_to_edge
• corner_chip
• pressure_curve
• serrated

The generator introduces randomized geometry, direction, thickness, branching, curvature, and edge irregularities.

Masks are also softened and morphologically processed before being used as conditioning inputs.

Example output filename:

crack_0042_hairline.png

The filename is later used by the prompt system to identify the crack type and generate a corresponding prompt.

---

2. Dynamic Prompt Generation

The prompt system is defined in:

prompt/
├── load_prompts.py
└── prompts.json

Rather than using a single prompt for every generated image, MaskDriven-SD builds prompts dynamically.

Prompt templates can vary attributes such as:

• marble type and color
• surface texture
• lighting
• camera scale
• vein pattern
• fracture appearance
• edge characteristics
• environmental context

For example, a hairline crack and an impact fracture receive different descriptions.

Wildcard values are selected deterministically using the numeric ID contained in the mask filename.

This means that regenerating:

crack_0042_hairline.png

will produce the same prompt configuration, making dataset generation more reproducible.

---

3. Mask-Conditioned Diffusion

The generated masks are passed to a ControlNet-conditioned diffusion pipeline.

The primary pipeline uses:

Stable Diffusion 3.5
        +
    ControlNet
        +
     TensorRT

The mask acts as a structural constraint while the text prompt controls the appearance, material, lighting, and photographic context.

Conceptually:

                    ┌──────────────────┐
                    │ Procedural Mask  │
                    └────────┬─────────┘
                             │
                             ▼
                    ┌──────────────────┐
                    │    ControlNet    │
                    └────────┬─────────┘
                             │
                             ▼
┌────────────────┐   ┌──────────────────┐
│ Dynamic Prompt │──▶│ Stable Diffusion │
└────────────────┘   │       3.5        │
                     └────────┬─────────┘
                              │
                              ▼
                     ┌──────────────────┐
                     │ Synthetic Image  │
                     └──────────────────┘

The resulting synthetic image is saved using the same base filename as its source mask.

For example:

ground_truth_masks/
└── crack_0042_hairline.png

synthetic_images/
└── crack_0042_hairline.png

This creates an automatically aligned image-mask dataset.

---

Repository Structure

MaskDriven-SD/
│
├── demo_diffusion/
│   ├── image/
│   ├── model/
│   ├── path/
│   ├── pipeline/
│   ├── utils_sd3/
│   ├── dd_argparse.py
│   ├── engine.py
│   └── ...
│
├── generate_seg_masks/
│   └── ground_truth_masks.py
│
├── prompt/
│   ├── __init__.py
│   ├── load_prompts.py
│   └── prompts.json
│
├── demo_controlnet_sd35.py
├── demo_img2img_flux.py
├── demo_flux_controlnet.ipynb
├── demo_sd35_controlnet.ipynb
│
└── LICENSE

---

Getting Started

Clone the repository

git clone https://github.com/Doga0/MaskDriven-SD.git
cd MaskDriven-SD

Generate Ground-Truth Masks

Masks can be generated directly from the Python module.

For example:

python -c "from generate_seg_masks.ground_truth_masks import generate_dataset; generate_dataset(start_n=0, total_n=1000, img_size=640, output_dir='ground_truth_masks')"

This generates masks such as:

ground_truth_masks/
├── crack_0000_hairline.png
├── crack_0001_impact.png
├── crack_0002_fragment.png
├── crack_0003_edge_to_edge.png
├── crack_0004_corner_chip.png
├── crack_0005_pressure_curve.png
└── crack_0006_serrated.png

Crack modes are rotated automatically during dataset generation.

---

Generate Synthetic Images with Stable Diffusion 3.5

The main Stable Diffusion pipeline is:

demo_controlnet_sd35.py

A dataset generation run can be configured around a directory containing the generated masks:

python demo_controlnet_sd35.py \
    --version 3.5-medium \
    --control-image ground_truth_masks \
    --controlnet-type canny \
    --controlnet-scale 1.0 \
    --prompts-file prompt/prompts.json \
    --output-dir synthetic_images \
    --height 640 \
    --width 640

Optional negative prompts can also be enabled:

--use-negative-prompts True

When a directory is provided through --control-image, the pipeline iterates over the masks and generates a corresponding synthetic image for each one.

Existing output files are skipped, which allows interrupted generation runs to continue without regenerating completed samples.

---

Prompt Configuration

Prompt behavior can be customized through:

prompt/prompts.json

The configuration contains global wildcard categories and crack-specific templates.

A simplified structure looks like:

{
  "wildcards": {
    "color": [
      "...",
      "..."
    ],
    "texture": [
      "...",
      "..."
    ],
    "lighting": [
      "...",
      "..."
    ]
  },

  "modes": {
    "hairline": {
      "template": "...",
      "negative": "..."
    },

    "impact": {
      "template": "...",
      "negative": "..."
    },

    "default": {
      "template": "...",
      "negative": "..."
    }
  }
}

Adding new defect domains therefore requires more than changing the text prompt: both the mask-generation logic and the prompt configuration can be extended to describe new structures and materials.

---

FLUX Experiments

The repository also contains experimental FLUX pipelines:

demo_img2img_flux.py
demo_flux_controlnet.ipynb

The TensorRT pipeline includes support for variants such as:

flux.1-dev
flux.1-schnell
flux.1-dev-canny
flux.1-dev-depth
flux.1-kontext-dev

These experiments explore alternative diffusion architectures and conditioning strategies for synthetic data generation.

---

Why Mask-Driven Generation?

Traditional text-to-image generation provides limited control over the exact spatial location of a defect.

For synthetic computer vision datasets, this is a problem because the ground-truth annotation must correspond precisely to the generated object.

MaskDriven-SD solves this by treating the annotation as an input to the generation process:

Annotation → Generation

rather than:

Generation → Annotation

This approach provides:

• explicit spatial control
• automatically available ground truth
• reproducible dataset generation
• scalable generation of rare defect patterns
• greater control over class distribution
• variation in appearance without losing structural labels

---

Potential Applications

Although the current implementation focuses on cracks and marble surfaces, the same idea can be generalized to other synthetic computer vision datasets.

Examples include:

• industrial defect detection
• surface inspection
• crack segmentation
• manufacturing quality control
• infrastructure inspection
• rare-class augmentation
• segmentation dataset generation

A procedural generator can define the desired ground truth, while a conditional diffusion model converts that abstract annotation into a realistic image.

---

Limitations

Synthetic data generated by diffusion models should not automatically be considered equivalent to real-world data.

Potential limitations include:

• diffusion artifacts
• imperfect adherence to conditioning masks
• unrealistic material properties
• domain gap between generated and real images
• dataset bias introduced by prompt templates
• dependency on model and ControlNet capabilities

Generated datasets should therefore be evaluated against real validation data before being used in production systems.

---

License

This repository is licensed under the Apache License 2.0.

See LICENSE for details.

Model weights and third-party dependencies may be distributed under their own licenses. Make sure to review the corresponding licenses before commercial use.

---

Acknowledgements

Parts of the accelerated diffusion inference stack are based on NVIDIA TensorRT diffusion tooling.

The project also builds on the broader open-source diffusion ecosystem, including Stable Diffusion, ControlNet, PyTorch, and related tooling.

---

Author

Doğa

GitHub: @Doga0

---

Project Summary

MaskDriven-SD generates synthetic computer vision datasets by first creating procedural ground-truth defect masks and then using those masks to structurally guide diffusion models. This produces realistic synthetic images together with automatically aligned segmentation annotations.