WB-Former: A Hybrid Model of CNN and Transformer for Water Body Extraction in Complex Scenes

Extracting small, narrow, or irregularly shaped water bodies from high-resolution remote sensing images is challenging. We introduce a novel framework that merges convolutional neural networks (CNNs) with Vision Transformers (ViTs) to extract complex-shaped water bodies. Extract data from the encoder and decoder, and integrate the ViT into the final downsampling layer to handle long-range dependencies. To better extract boundary information, we designed a boundary-aware edge enhancement module with a dedicated edge attention mechanism. Next, we developed an adaptive information enhancement module (AIEM). To handle complex data flows, we introduced a redundancy attenuation module. Furthermore, by using external outputs to generate queries for the ViT, we integrated features from different network levels, enabling it to fuse both local and global contextual information. Finally, we proposed a novel loss function that incorporates Rolle’s theorem and adaptive distance, making the smoothness constraint more flexible and enhancing sensitivity to object boundaries, thereby effectively improving segmentation accuracy in boundary regions. Extensive experiments on the WD_dataset, Wuhan High-resolution Dense-Labeling Dataset (WHDLD), and Sentinel-2 datasets demonstrate that our model outperforms existing state-of-the-art methods, including those based on CNN and Transformer architectures. Extensive experiments on the WD_dataset, WHDLD, and Sentinel-2 datasets demonstrate that our model outperforms existing state-of-the-art methods, including those based on CNN and Transformer architectures. Specifically, WB-Former achieves a 3% improvement in intersection over union (IoU) and a 1% increase in average precision over the WaterDeep model on the WHDLD dataset. On the WD_dataset, it shows a 7% improvement in IoU and a 3% increase in average precision. On the Sentinel-2 dataset, it achieves a 6% improvement in IoU and a 3% increase in average precision.