Attention-enhanced multi-task learning for binary segmentation and fine-grained aquatic plant classification in UAV imagery

Accurate monitoring of aquatic vegetation from unmanned aerial vehicle (UAV) imagery remains challenging due to complex water backgrounds, severe inter-class similarity, and the lack of balanced, dual-annotated datasets. Existing studies primarily address segmentation or classification independently, limiting their effectiveness for integrated species-level analysis. To address these gaps, this study proposes a clearly defined attention-enhanced multi-task learning framework that simultaneously performs binary segmentation and 14-class species classification, enabling unified structural and semantic understanding. The model employs a shared encoder with attention-guided skip connections and a joint optimization strategy to enhance feature discrimination while reducing redundancy. Comprehensive ablation analysis demonstrates that attention improves both segmentation and classification performance, while joint learning with Gaussian blur achieves the best overall balance, confirming the complementary role of spatial and semantic features. On a newly collected UAV dataset from diverse wetlands in Bangladesh, the proposed model achieves a Dice coefficient of 0.7344, mIoU of 0.6904, and pixel accuracy of 0.8757 for segmentation, along with 98.77% classification accuracy and an F1-score of 0.9874, indicating strong performance across both tasks. In addition, computational complexity analysis shows that the proposed framework reduces parameters by (sim)50% (31.10M vs. 62.09M), lowers FLOPs (54.66 vs. 96.31 GFLOPs), and improves inference speed by (sim)48.6% compared to deploying separate single-task models for segmentation and classification, demonstrating its suitability for real-time UAV deployment. Furthermore, Gradient-weighted Class Activation Mapping (Grad-CAM) and Grad-CAM++ are employed to provide visual explanations of model predictions, improving interpretability and reliability. The results demonstrate robust performance in complex aquatic environments and highlight the framework’s suitability for large-scale biodiversity monitoring, invasive species detection, and data-driven freshwater ecosystem management.