A semantic segmentation model for counting wheat heads in field images, designed for yield estimation and field maturity assessment.
Manual wheat head counting is time-consuming and impractical at scale. Traditional computer vision methods struggle with overlapping heads and varying lighting conditions. HeadCount is a fine-tuned deep learning model that uses semantic segmentation to distinguish wheat heads from other plant organs, then applies distance transforms and peak detection to resolve individual heads, enabling automated and accurate counting from field images.
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")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
The model is trained on GWFSS_v1.0_labelled from the Global Wheat Full Semantic Organ Segmentation dataset.
This project is licensed under the MIT License - see the LICENSE file for details.