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 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

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).

Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Model tree for pasqualedem/weeddiffusion

Finetuned
(598)
this model