Grain detection and yield prediction in wheat using hyperspectral imaging and machine learning

Accurate yield prediction in wheat is crucial for breeding, resource planning, and food security. Traditional image-based methods often rely on canopy-level features or vegetation indices that do not directly represent the grain-bearing parts of the plant. This study presents a hyperspectral imaging (HSI)-based approach for pixel-level grain detection and yield prediction in wheat spikes. Near-infrared HSI data (900–1700 nm) were collected from 30 spikes before and after grain removal to establish accurate ground truth. Two spectral libraries were developed, one comprising average spectra per spike and another containing pixel-wise spectral signatures. Machine learning models, including Partial Least Squares Discriminant Analysis (PLS-DA), Random Forest (RF), Support Vector Machine (SVM), Gradient Boosting (GB), and Decision Tree (DT), were trained to classify grain and non-grain regions. The RF and DT models achieved perfect performance with 100% accuracy and F1-score. Spectral regions within 950–975 nm, 1150–1250 nm, and 1350–1375 nm contributed most to model performance, suggesting that the full spectral range is not always required. Yield was predicted using the grain-refined spike area obtained from the RF model, achieving a strong correlation with measured yield (R2 = 0.800 ± 0.046). The results highlight the potential of integrating pixel-level HSI data with machine learning for precise, non-destructive grain detection and robust yield prediction in wheat.