rotsl/grayleafspot-segmentation

🔬 Gray Leaf Spot Segmentation Model

PyTorch U-Net models for gray leaf spot (Magnaporthe and related fungal) colony segmentation on 90 mm petri-dish images.

▶ Try the live demo — upload images, run inference, see overlays & growth charts in your browser.

---

Model Weights

File	Architecture	Params	Area-Consistency Weight	Description
grayleafspot.pt	smp.Unet (ResNet-34)	~24.4 M	—	Main encoder–decoder model
best_area_w_0.1.pt	SmallUNet	~250 K	0.1	Light area regularisation
best_area_w_0.3.pt	SmallUNet	~250 K	0.3	Moderate area regularisation
best_area_w_0.5.pt	SmallUNet	~250 K	0.5	Balanced BCE + area
best_area_w_0.7.pt	SmallUNet	~250 K	0.7	Strong area consistency (used by demo) ✅ recommended

All SmallUNet variants share the same architecture:

Input (3 × 256 × 256)
  │
  ├─ enc1: ConvBlock(3 → 16)              ─── skip s1
  ├─ enc2: MaxPool2d → ConvBlock(16 → 32)  ─── skip s2
  ├─ enc3: MaxPool2d → ConvBlock(32 → 64)  ─── skip s3
  ├─ enc4: MaxPool2d → ConvBlock(64 → 128) ─── skip s4
  │
  ├─ bottleneck: MaxPool2d → ConvBlock(128 → 256)
  │
  ├─ up4: Upsample + cat(s4) → ConvBlock(384 → 128)
  ├─ up3: Upsample + cat(s3) → ConvBlock(192 → 64)
  ├─ up2: Upsample + cat(s2) → ConvBlock(96 → 32)
  ├─ up1: Upsample + cat(s1) → ConvBlock(48 → 16)
  │
  └─ head: Conv2d(16 → 1) → Sigmoid

Each ConvBlock = Conv3×3 (no bias) → ReLU → Conv3×3 (no bias) → ReLU.

Property	Value
Input	256 × 256 RGB
Output	1-channel sigmoid probability mask
Training loss	BCE + area-consistency loss
CPU compatible	✅ Pure PyTorch — no custom CUDA kernels

---

Quick Start

Download & Inference (SmallUNet)

import torch
from huggingface_hub import hf_hub_download

# Download weights
path = hf_hub_download("rotsl/grayleafspot-segmentation", "best_area_w_0.7.pt")

# Load checkpoint
ckpt = torch.load(path, map_location="cpu", weights_only=False)

# Build model (SmallUNet architecture — see demo repo for full class definition)
# https://huggingface.co/rotsl/grayleafspot-segmentation-demo/blob/main/app.py
from model import SmallUNet  # or copy the class from the demo app.py

model = SmallUNet(in_channels=3, out_channels=1, base_channels=16)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# Run inference on a 256×256 RGB tensor
import numpy as np
from PIL import Image

img = Image.open("petri_dish.jpg").convert("RGB").resize((256, 256))
x = torch.from_numpy(np.array(img).transpose(2, 0, 1)).float() / 255.0
x = x.unsqueeze(0)

with torch.no_grad():
    prob = model(x)[0, 0].numpy()

mask = (prob > 0.5).astype(np.uint8) * 255
Image.fromarray(mask).save("colony_mask.png")

Download & Inference (Main U-Net)

import torch

path = hf_hub_download("rotsl/grayleafspot-segmentation", "grayleafspot.pt")
model = torch.load(path, map_location="cpu", weights_only=False)
model.eval()

---

Training & Inference Pipeline

Environment Setup (Apple Silicon recommended)

python3.10 -m venv trainenv
source trainenv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt

Dataset Preparation

Place raw images in raw/ and corresponding masks in masks/ (matching filenames). Expand the dataset with augmentations:

python src/build_augmented_dataset.py --copies-per-image 4 --clean

Training

Standard U-Net training:

python src/train.py \
    --image-dir augmented_dataset/raw \
    --mask-dir augmented_dataset/masks \
    --epochs 40 --batch-size 4 --lr 1e-4 \
    --image-size 256 --freeze-encoder-epochs 5

Area-consistency U-Net (LabelMe JSON polygons):

./trainenv/bin/python src/area_consistency/train_area.py

Inference (CLI)

Single image:

python src/predict.py --input raw/your_image.jpg --weights models/best_finetuned.pt --output-dir predictions

Folder:

python src/predict.py --input raw --weights models/best_finetuned.pt --output-dir predictions

Best Practices

• Keep trusted human-labelled masks unchanged.
• Use augmentations and area-consistency loss for improved generalisation.
• Inspect overlay outputs to verify mask quality.
• On Apple Silicon, MPS acceleration is used automatically if available.

---

Demo Spaces

✅ Recommended: rotsl/grayleafspot-segmentation-demo

Uses best_area_w_0.7.pt (SmallUNet with area-consistency loss). More accurate segmentation with better boundary adherence thanks to the area-consistency regularisation.

Features: dish detection → colony segmentation → crack & hyphae analysis → 16 morphometric measurements → time-series growth charts → CSV/JSON export.

Source code: rotsl/grayleafspot-segmentation-demo (model repo)

Legacy: rotsl/fungal-colony-input

Uses grayleafspot.pt (smp.Unet with ResNet-34 encoder). This is the earlier, larger model trained with standard BCE loss only — it is less accurate than the area-consistency variant above, particularly for colony boundary delineation and area estimation. Kept available for reference and backward compatibility.

---

Citation

@misc{rohan_r_2026,
  author       = {rohan r},
  title        = {grayleafspot-segmentation (Revision 0e85f71)},
  year         = 2026,
  url          = {https://huggingface.co/rotsl/grayleafspot-segmentation},
  doi          = {10.57967/hf/8416},
  publisher    = {Hugging Face}
}

License

Apache License 2.0 — see LICENSE for details.