Uncertainty and Feature-Based Weighted Loss for 3D Wheat Part Segmentation

Deep learning techniques and point clouds have proved their efficacy in 3D segmentation tasks of objects. Nevertheless, the accurate plant organ segmentation is a formidable challenge due to their complex structure and variability. Furthermore, presence of over-represented and under-represented parts, occlusion, and uneven distribution complicates the 3D part segmentation tasks. Even though deep learning techniques often exhibit exceptional performance, they also face challenges in applications where accurate trait estimation is required. To handle these issues, we propose a novel uncertainty and feature based weighted loss that incorporates uncertainty metrics and features of the plant or crop. We use Gradient Attention Module (GAM) with PointNet++ baseline to validate our approach. By dynamically introducing uncertainty and feature scores into the training process, it promotes more balanced learning. Through comprehensive evaluation, we illustrate the advantages of UFL (Uncertainty and Feature based Loss) as compared to standard CE (Cross entropy loss) with our own constructed real Wheat dataset. The outcomes demonstrate consistent improvements in Accuracy (ranging from 0.9% to 4.2%) and Ear mIoU (ranging from 1.8% to 15.3%) over the standard Cross-Entropy (CE) loss function. As a result, our work contributes to the development of more robust and reliable segmentation models. This approach not only pushes forward the boundaries of precision agriculture but also has the potential to influence related areas where accurate segmentation is pivotal.