File size: 3,461 Bytes
bce4c63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | ---
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).
|