|
|
--- |
|
|
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. |
