File size: 2,312 Bytes
b754fe4
78f1163
b754fe4
 
 
 
 
 
 
 
 
 
78f1163
b754fe4
78f1163
b754fe4
 
78f1163
 
 
62656f3
78f1163
4662b34
78f1163
5e5f960
 
 
e9ca6a3
78f1163
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ac9eaf5
78f1163
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4a10d47
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
---
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.