---
license: mit
tags:
  - stable-diffusion
  - dreambooth
  - diffusers
  - semantic-segmentation
  - agriculture
  - weed-mapping
base_model: runwayml/stable-diffusion-v1-5
library_name: diffusers
pipeline_tag: text-to-image
---

# WeedDiffusion — Model Checkpoints

Model checkpoints for **WeedDiffusion: A Dual-Branch Synthetic Augmentation Framework for Weed Mapping**.

Code: https://github.com/pasqualedem/WeedDiffusion

This repository hosts the trained checkpoints that are **not** shipped with the GitHub repository due to their size:

1. Two class-specific **DreamBooth** Stable Diffusion v1.5 models used for synthetic data augmentation.
2. The best **semantic-segmentation** checkpoints used to produce the paper's result tables.

## Repository layout

```
crops/           # DreamBooth SD-1.5 pipeline (diffusers format) — crop inpainting branch
weeds/           # DreamBooth SD-1.5 pipeline (diffusers format) — weed generation branch
segmentation/    # Best-mIoU segmentation checkpoints, grouped by dataset/experiment
  ├── images_37/     # real-only baseline experiments
  └── weeddiff_37/   # experiments trained with WeedDiffusion-augmented data
```

## DreamBooth generative models

Both are [Stable Diffusion v1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) fine-tuned with DreamBooth (`resolution=512`, `learning_rate=1e-6`, `max_train_steps=4000`, text encoder trained).

| Branch  | Instance prompt | Used for |
|---------|-----------------|----------|
| `crops` | `sks crop`      | Inpainting masked regions of real field images |
| `weeds` | `sks weed`      | Generating individual synthetic weeds (then segmented with SAM and inserted) |

Only inference weights are included (the full diffusers pipeline: `unet`, `vae`, `text_encoder`, `safety_checker`, `tokenizer`, `scheduler`, `feature_extractor`). Optimizer / training-resume state is omitted.

### Usage

```python
from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
from huggingface_hub import snapshot_download

# Download a single branch
crops_dir = snapshot_download("pasqualedem/weeddiffusion", allow_patterns="crops/*")
weeds_dir = snapshot_download("pasqualedem/weeddiffusion", allow_patterns="weeds/*")

# Weed generation branch
pipe = StableDiffusionPipeline.from_pretrained(f"{weeds_dir}/weeds").to("cuda")
image = pipe("sks weed").images[0]

# Crop inpainting branch
inpaint = StableDiffusionInpaintPipeline.from_pretrained(f"{crops_dir}/crops").to("cuda")
```

In the GitHub pipeline these correspond to the local paths `out/model_crops` and `out/model_weeds`.

## Segmentation checkpoints

For each experiment, we provide the single checkpoint selected by the evaluation code
(`main.py`): the highest-`train_mIoU` checkpoint of the final training run. Filenames encode the
epoch and the training mIoU, e.g. `phenobench_epoch=185_train_mIoU=0.8644.ckpt`.

- `segmentation/images_37/<arch>_<variant>` — models trained on real data only.
- `segmentation/weeddiff_37/<arch>_<variant>` — models trained with WeedDiffusion augmentation.

Architectures: `deeplab`, `unet`, `erfnet`. Variants: `base`, `color`, `geo`, `geocolor`
(and numbered re-runs). These are PyTorch-Lightning checkpoints; load them with the matching
config via `main.py test` (see the GitHub README).

## License

MIT — see the code repository.

## Citation

If you use these models, please cite WeedDiffusion (see the GitHub repository).