Skip to content

DaRL-GenAI/PlantShade-code

Repository files navigation

PlantShade

Predicting Plant Shadows for Lighting-Aware Robotic Agricultural Operation (IROS 2026)

Project Page Dataset Model Data Simulator

PlantShade is a plant shade dataset and a ControlNet-conditioned diffusion model that generates realistic, time-varying plant shadow maps from a fixed top-down camera. Given a plant image and a text prompt describing the supplementary light position, the model predicts the corresponding shadow map without full 3D rendering, enabling in silico evaluation of supplemental lighting and its effect on canopy photosynthesis.

This repository is the minimal, runnable pipeline:

zip (raw capture) ─▶ 1 extract shadow ─▶ 2 build split ─▶ 3 train ─▶ 4 infer ─▶ 5 photosynthesis

Repository layout

.
├── environment.yaml         # conda env (recommended)
├── requirements.txt         # pip fallback
├── 1_extract_shadow.py      # raw renders -> shadow targets (target/ + targetRaw/)
├── 2_build_split.py         # assemble (source, target, prompt) train/test JSON
├── 3_train.py               # train the ControlNet shade model
├── 4_infer.py               # generate shadow maps + metrics + compare images
├── 5_photosynthesis.py      # shadow maps -> canopy photosynthesis (NRH)
├── cldm/ ldm/ annotator/    # ControlNet + Stable Diffusion backbone
├── share.py  config.py
├── models/
│   ├── cldm_v21.yaml        # model config (SD 2.1)
│   └── cldm_v15.yaml
├── data/                    # empty; download scene zips from HuggingFace here
├── outputs/                 # empty; created by steps 4-5
├── results/                 # example predictions + reproduced paper numbers
└── notebooks/
    └── pipeline.ipynb       # the whole pipeline, step by step

This is a code-only repository. Dataset archives and the trained checkpoint are hosted on HuggingFace (see below) and are not committed here.

Setup

Requires a CUDA GPU.

conda env create -f environment.yaml
conda activate control
# (pip fallback, if you already have torch+CUDA: pip install -r requirements.txt)

Data. No dataset is committed here. Download example scenes from HuggingFace into data/examples/ (Step 1 shows the command); the full dataset (all species, growth stages, layouts) lives on HuggingFace.

Dataset: https://huggingface.co/datasets/xiao0o0o/PlantShade Data-collection simulator (Helios + UE5): https://github.com/ARLabXiang/AgriRoboSimUE5/releases/tag/plantshade

Checkpoint.

The only PlantShade model is plantshade_epoch51.ckpt — the trained checkpoint, used for inference and to reproduce the results. Download it once:

hf download xiao0o0o/PlantShade-ControlNet plantshade_epoch51.ckpt --local-dir models/

That is all you need for Steps 4–5. The two files below are only for training from scratch (Step 3), and are standard third-party weights, not part of PlantShade:

file what it is needed for
models/v2-1_512-ema-pruned.ckpt public Stable Diffusion v2.1 base (from stabilityai/stable-diffusion-2-1-base) training only
models/control_sd21_ini.ckpt SD2.1 + zero-init ControlNet, derived from the base training only
# only if training from scratch: build the ControlNet-initialised checkpoint once
python tool_add_control_sd21.py models/v2-1_512-ema-pruned.ckpt models/control_sd21_ini.ckpt

Step 1 — Extract shadow targets

Unzip a scene, then threshold the shadow render and composite the color-coded training target (background #a0a0a0, skeleton #ffffff, shadow #070707).

# download one example scene from HuggingFace
hf download xiao0o0o/PlantShade tomato_seed2_14_7_77.zip \
    --repo-type dataset --local-dir data/examples/

mkdir -p data/PlantShadeUnzip
unzip data/examples/tomato_seed2_14_7_77.zip -d data/PlantShadeUnzip/
# -> data/PlantShadeUnzip/tomato_seed2_14_7_77/{14, 14_shadow, 21, 21_shadow, ...}

python 1_extract_shadow.py \
    --base_root data/PlantShadeUnzip \
    --out_root  data/shadow
# writes data/shadow/target/<scene>/<day>/seg_*.png  and  data/shadow/targetRaw/...

Step 2 — Build the train/test split

Pairs each structural source (skeleton) with its target (shadow map) and the prompt (light position on the circular trajectory); splits 80/20 (seed 42).

python 2_build_split.py \
    --source_root data/PlantShadeUnzip \
    --target_root data/shadow/target \
    --out_dir     data
# -> data/train_split.json , data/test_split.json

Step 3 — Train

python 3_train.py
# override defaults with env vars, e.g.:
#   PS_TRAIN_JSON=data/train_split.json PS_BATCH=16 PS_EPOCHS=501 python 3_train.py
# checkpoints -> outputs/train/<timestamp>/best/ and periodic/

Backbone: Stable Diffusion v2.1 (frozen) + ControlNet branch. Defaults: batch 16, lr 1e-5, 512×512.

Step 4 — Inference (use the provided trained checkpoint)

No training needed to try the model — point PS_CKPT at the trained checkpoint:

# checkpoint downloaded in Setup: models/plantshade_epoch51.ckpt
PS_CKPT=models/plantshade_epoch51.ckpt \
PS_TEST_JSON=data/test_split.json \
python 4_infer.py
# -> outputs/predictions/generated_*.png , compare_*.png , result.json , ave_result.json

4_infer.py writes generated_*.png (the pure predicted shadow map) and compare_*.png (left: prediction, right: ground truth, with a SSIM / MSE / mIoU / B-IoU / LPIPS header) into outputs/predictions/. Example predictions and reproduced numbers are in results/.

Step 5 — Photosynthetic gain

Converts predicted shadow maps into canopy photosynthesis with the non-rectangular hyperbola (NRH) model, using rendered depth maps for leaf area.

python 5_photosynthesis.py \
    --split      data/test_split.json \
    --pred_dir   outputs/predictions \
    --output_csv outputs/photosynthesis.csv
# CSV columns: scene, day, seg, species, gt_P_canopy, pred_P_canopy, ...

Reproducing paper results

results/ holds pure predicted shadow maps and one quantitative group computed from those prediction images by 5_photosynthesis.py. For tomato · 1×1 · day 63, canopy photosynthesis (µmol CO₂ m⁻² s⁻¹) over the sampled supplementary-light positions:

mean optimal placement
ground-truth shadow 14.21 14.92
predicted shadow 13.77 14.98

Running the full test set reproduces the paper's complete finding (soybean largest gain, strawberry smallest). See results/README.md and results/predictions/.

Notebook

notebooks/pipeline.ipynb: Steps 1–2 and 4 are runnable but need the dataset / checkpoint downloads and a GPU; Step 3 (training) is documentation only; Step 5 runs directly on the bundled example (results/example_run/, tomato · 1×1 · day 63) with no download and no GPU, and prints the reproduced photosynthesis numbers.

Citation

@inproceedings{da2026plantshade,
  title     = {PlantShade: Predicting Plant Shadows for Lighting-Aware Robotic Agricultural Operation},
  author    = {Da, Longchao and Liu, Xiaoou and Li, Xingjian and Xiang, Lirong and Wei, Hua},
  booktitle = {IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
  year      = {2026}
}

The ControlNet / Stable Diffusion backbone code is adapted from lllyasviel/ControlNet (see LICENSE).

About

No description, website, or topics provided.

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages