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