---
license: cc-by-4.0
datasets:
- GlobalWheat/GWFSS_v1.0
metrics:
- f1
base_model:
- microsoft/resnet-50
pipeline_tag: image-segmentation
tags:
- computer-vision
- semantic-segmentation
- instance-counting
- phenotyping
- agriculture
- crop-analysis
- yield-estimation
- wheat
---

# 🌾 HeadCount

A semantic segmentation model for counting wheat heads in field images. Designed for yield estimation, flowering time detection, and field maturity assessment.

**[Interactive Demo →](https://huggingface.co/spaces/chmcbs/HeadCount)**

**[GitHub Repo →](https://github.com/chmcbs/HeadCount)**

## 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)
- **Counting Method:** Distance transform + peak detection on head mask
- **Loss Function:** Dice loss with inverse frequency weighting (1.5× stem boost)
- **Optimizer:** Adam with CosineAnnealingLR scheduling

## Performance

| Class      | F1    |
|------------|-------|
| Background | 0.858 |
| Leaf       | 0.889 |
| Stem       | 0.535 |
| Head       | 0.897 |

## Example Usage

```python
from inference import GWFSSModel
from PIL import Image

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

# Load and process image
image = Image.open("input.jpg")
predictions = model.predict(image)

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

# Create visualisation
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:** Bright or harsh lighting can cause over-segmentation, splitting single heads into multiple detections
- **Overlapping Heads:** Dense clusters with significant overlap are challenging to separate accurately
- **Colour Dependency:** Performance is lower on senesced plants

## Training Data

This model is trained on [`GWFSS_v1.0_labelled`](https://huggingface.co/datasets/GlobalWheat/GWFSS_v1.0) from the [Global Wheat Full Semantic Organ Segmentation](https://doi.org/10.1016/j.plaphe.2025.100084) dataset.