HeadCount

A semantic segmentation model for counting wheat heads in field images, designed for yield estimation and field maturity assessment.

Interactive Demo →

GitHub Repo →

Model Details

• Architecture: DeepLabV3+ with ResNet50 encoder
• Framework: PyTorch with segmentation-models-pytorch
• Input: RGB images (resized to 512×512)
• Output: 4-class segmentation (Background, Leaf, Stem, Head)
• Loss Function: Dice loss with inverse frequency weighting (1.5× stem class weight)
• Optimisation: Adam with CosineAnnealingLR learning-rate scheduling
• Counting: Distance transform + peak detection on head mask

Performance

Class	F1
Head	0.897
Leaf	0.889
Background	0.858
Stem	0.535

Example Usage

from inference import GWFSSModel
from PIL import Image

# Load model
model = GWFSSModel("model.pth")

# Process image
image = Image.open("input.jpg")
predictions = model.predict(image)

# Count heads
num_heads = model.count_heads(predictions)
print(f"Detected {num_heads} heads")

# Save overlay
overlay = model.overlay_mask(image, predictions, alpha=0.5, heads_only=True)
overlay.save("output.png")

Limitations

Best performance is achieved with overhead imagery under diffuse lighting. Known challenges include:

• Lighting Sensitivity: Harsh lighting can cause over-segmentation, splitting heads into multiple detections
• Overlapping Heads: Dense clusters with significant overlap are challenging to separate accurately
• Colour Dependency: Performance is worse on senesced plants due to limited training examples

Training Data

This model is trained on GWFSS_v1.0_labelled from the Global Wheat Full Semantic Organ Segmentation dataset.