Skip to content

gabe-zhang/sugarbeet-weed-segmentation

Repository files navigation

Sugarbeet Weed Segmentation

License: MIT Python PyTorch Lightning W&B

A segmentation framework for agricultural weed detection in sugarbeet fields, supporting semantic, instance, and panoptic segmentation across three categories: soil (background), crop (sugarbeet), and weed.

Trained on the PhenoBench dataset with CodaLab benchmark submission support.

Models

Task Model Framework
Semantic segmentation ERFNet, DeepLabV3+ PyTorch Lightning
Instance / Panoptic segmentation YOLO26 Ultralytics

Branches

Branch Platform Description
main Desktop / Cloud GPU Training, inference, and submission scripts
brain NVIDIA Jetson Xavier NX TensorRT edge deployment

Setup

git clone https://github.com/gabe-zhang/sugarbeet-weed-segmentation.git
cd sugarbeet-weed-segmentation
uv sync

Download PhenoBench and update the data path in the YAML config files under config/.

Usage

Semantic Segmentation (ERFNet / DeepLabV3+)

# Train
uv run src/train.py --config config/<config>.yaml --export_dir runs

# Validate / submit
uv run src/submit.py \
    --config config/<config>.yaml \
    --ckpt_path runs/<run>/checkpoints/<best>.ckpt \
    --export_dir runs --split val --tta

Supports ensemble (--config/--ckpt_path repeated), TTA (--tta, --tta_scales), and morphological postprocessing (--opening).

Instance / Panoptic Segmentation (YOLO26)

# Convert PhenoBench to YOLO format
uv run tools/convert_phenobench_to_yolo.py

# Train
uv run src/train_yolo.py --config config/yolo26x-baseline.yaml

# Predict (instance -> semantic)
uv run src/predict_yolo.py \
    --weights runs/<run>/weights/best.pt --split val

# Panoptic submission
uv run src/submit_panoptic.py \
    --weights runs/<run>/weights/best.pt --split test

Validate Submission

uv run tools/validator.py --task semantics \
    --phenobench_dir data/PhenoBench --zipfile <submission>.zip

Classes

Class ID Color (RGB)
Soil (background) 0 (0, 0, 0)
Crop (sugarbeet) 1 (0, 255, 0)
Weed 2 (255, 0, 0)

Benchmark (Jetson Xavier NX)

Input: 1920x1080 | ERFNet model

Model Latency (ms) FPS Speedup
PyTorch 717.9 1.4 1.0x
TensorRT FP32 309.3 3.2 2.3x
TensorRT FP16 83.6 12.0 8.6x

License

This project is licensed under the MIT License.

Attribution

Code adapted from PRBonn/phenobench-baselines. Trained on the PhenoBench dataset (CC BY-SA 4.0).

If you use the PhenoBench dataset and/or models, please cite accordingly (see CITATION.md).

About

Agricultural weed segmentation using YOLO, ERFNet, and DeepLabV3+ on the PhenoBench dataset. Includes TensorRT for edge inference.

Topics

Resources

License

Stars

Watchers

Forks

Contributors