Flood-specific datasets were developed by processing fine-resolution imagery obtained through autonomous aerial systems. This dataset comprised 290 original images, which were expanded to 1160 images through data augmentation, as shown in Fig. 2. Augmentation procedures included controlled brightness adjustment (±20%), contrast shifting (±25%), horizontal and vertical flipping (applied with a 50% probability), and salt-and-pepper noise injection (applied with a 5% probability). These transformations were carefully chosen to simulate diverse environmental conditions observed in real-world flood imagery, such as varying lighting, reflections, and occlusions. The augmentation strategy enhances data diversity, mitigates overfitting, and promotes better generalization across challenging flood scenarios. The dataset was split into 70% for training with 812 images, 20% for validation with 232 images, and 10% for testing with 116 images, ensuring a balanced evaluation.

Prior to model training, a comprehensive series of pre-processing steps was applied to all images to ensure high-quality input for the segmentation models. First, each image was resized to a standard resolution of 512 × 512 pixels, aligning with the input specifications required by both the Vision Transformer (ViT) and U-Net architectures. The images were then converted to the YCbCr color space and normalized, which served to standardize pixel intensity distributions, reduce variability, and promote stable training behavior.

For processing through the Vision Transformer encoder, images were divided into 16 × 16 patches with a 4-pixel overlap, maintaining spatial continuity across patch boundaries and facilitating efficient patch-wise embedding. Finally, batch normalization was integrated into the training pipeline, stabilizing the learning process by reducing covariate shift in the feature maps and accelerating convergence. Collectively, these pre-processing and augmentation steps optimized the dataset for deep learning, improved model generalization, and contributed to the reliability of segmentation results across diverse and challenging flood scenarios. Cross-entropy loss was used for training, and evaluation metrics included precision, recall, accuracy, F1-score, and mean IoU. Model comparisons were statistically validated using the Kruskal-Wallis test.

The images were normalized, resized and converted to YCbCr color space before embedding into the ViT Encoder using patch-wise embedding. The uniform resolution matches the ViT and U-Net segmentation model for computational efficiency. The dataset is separated into three subsets for learning, tuning, and evaluation. These pre-processing steps prepare data for the ViT-UNet model to delineate flood-affected regions, ensuring reliable segmentation results as shown in Eq. (1). Additionally, batch normalization was applied during training to stabilize learning and improve convergence. (1)Inorm=I−μσ

Here I, μ, and σ represent the input image, mean, and standard deviation. This reduces input variability.

During normalization, images are divided into 16 × 16 pixel patches with 4-pixel overlap between adjacent patches for efficient Vision Transformer processing. The overlap maintains spatial continuity across boundaries. A 512 × 512 pixel input image generates 1024 patches.

Shown in Fig. 3, the ViT-UNet architecture starts by leveraging a Vision Transformer backbone to extract both global and local contextual features from input imagery. This is followed by a U-Net decoder, which reconstructs segmentation masks through a four-stage upsampling strategy. Each upsampling stage applies a transposed convolution succeeded by a 3 × 3 convolution, batch normalization, and ReLU activation to progressively restore spatial resolution and refine feature representations. Skip connections are employed to link corresponding encoder and decoder layers, effectively preserving fine-grained details and compensating for information lost during downsampling. Throughout these stages, the number of feature channels is systematically decreased (512, 256, 128, 64) to balance computational efficiency with representational power.

To address the persistent challenges of complex flood boundaries, such as reflections and blurred edges the final feature map produced by the decoder is processed by the proposed Flood-Aware Refinement Block (FARB). This specialized module integrates residual smoothing with spatial and channel attention to enhance boundary delineation and suppress noise. The network outputs a binary, pixel-wise segmentation mask, in which each pixel is classified as either flooded or non-flooded, enabling fine-grained flood extent delineation.

To address the ambiguity of flood boundaries in aerial imagery, we propose a flood-aware refinement block (FARB), a learnable module integrating residual refinement for preserving essential spatial features, spatial-channel attention for highlighting flood regions and suppressing noise, and a convolution-based smoothing layer to ensure boundary continuity. This yields sharper and more reliable flood masks.

Unlike standard smoothing operations, we propose a refinement block that combines spatial-channel attention and a residual learning path. The residual connection preserves essential spatial features by combining smoothed inputs while filtering out irrelevant regions, improving accuracy. This process can be mathematically formalized, as shown in Eq. (6).

Let X be the input feature map: (2)Xsmooth=ReLU(BN(Conv3×3(X))) (3)Ms=σ(Conv1×1(X))where σ is the sigmoid function. (4)Mc=σ(FC2(ReLU(FC1(GAP(X)))))where GAP is global average pooling, FC₁ reduces channels to C/r, FC₂ restores to C, and σ is sigmoid. (5)Xattn=X×Ms×Mc (6)Xrefined=Xsmooth+Xattn

The decoder output is further refined by the proposed Flood-Aware Refinement Block (FARB). A detailed, step-by-step schematic of the FARB architecture, including all operations and their mathematical formulations, is provided in Table 1 and Eqs. (2) to (6). As illustrated in Fig. 3, FARB integrates residual smoothing, spatial attention, and channel attention with fusion and residual addition, sharpening segmentation boundaries and suppressing irrelevant noise. Complete hyperparameter details are provided to ensure full reproducibility. The complete architecture of the Flood-Aware Refinement Block (FARB) is summarized in which lists each operation, its function, and the resulting output shape. This table provides a clear step-by-step view of how residual smoothing, spatial attention, and channel attention are integrated and fused, ultimately producing the refined feature map used by the segmentation head.

The post-processing block also contributes to real-time prediction capability, enabling accurate and actionable outputs in resource-constrained flood response environment.

The Vision Transformer encoder introduces O(n2d) complexity due to its self-attention mechanism, where n is the number of image patches (e.g., 1024 for a 512 × 512 input with 16 × 16 patches) and d is the embedding dimension (768). Although effective at modeling global dependencies, the quadratic scaling with sequence length can be computationally intensive. The U-Net decoder’s complexity is mainly from convolutional operations, approximated as O(hwc2k2) per layer, where h and w are spatial dimensions, c is the channel count, and k is the kernel size (typically 3). For greater efficiency, our model reduces channel dimensions progressively from 512 to 256 to 128 during upsampling. The Flood-Aware Refinement Block (FARB), a post-processing module, improves boundary delineation and suppresses noise by combining residual smoothing with spatial-channel attention mechanisms. Our model achieves a balanced trade-off, requiring only 1.29 h for training and 36.17 min for testing. It yields a segmentation accuracy of 93.09% and an mIoU of 0.8640.

We also report model complexity in terms of trainable parameters and FLOP ViT-UNet+FARB has 45.2 M parameters and 34.8 GFLOPs per 512 × 512 image during inference. For comparison, baseline U-Net Conv2D has 7.8 M parameters and 8.9 GFLOPs, while the Swin Transformer baseline has 63.1 M parameters and 42.5 GFLOPs. This demonstrates that ViT-UNet+FARB achieves a favorable trade-off between accuracy and computational cost, being more efficient than Swin and significantly outperforming traditional U-Net models. (7)Yout=W2⋅ReLU(W1⋅X+b1−μσ2+ϵ+β)+b2

As shown in Table 2, our method achieves a favorable trade-off between accuracy and computational cost, being more efficient than Swin Transformer and significantly outperforming traditional U-Net models.

To address the challenges of accurately segmenting ambiguous flood boundaries, we propose a Flood-Aware Refinement Block (FARB) as a learnable post-processing stage. This block integrates residual fine-tuning, spatial-channel attention, and convolutional smoothing to enhance mask predictions. The residual pathway preserves essential spatial features by learning an additive correction over the smoothed input, effectively reducing noise while maintaining edge sharpness. The spatial-channel attention mechanism emphasizes flood-relevant regions and suppresses background interference, allowing precise delineation of water extents. The convolution-based relative smoothing and attention module ensures boundary continuity and smooth transitions. These combined operations significantly sharpen the raw segmentation outputs while adding negligible computational overhead.

The model was trained using a low learning rate of 0.00001 to prevent oscillations in the loss function and ensure stable convergence. A batch size of 8 was used to balance GPU memory constraints with gradient accuracy. Training was conducted over 50 epochs, with validation loss stabilizing around epoch 45. To prevent overfitting on the small dataset, a dropout rate of 0.5 was applied, encouraging robust feature learning. The Adam optimizer was selected for its ability to dynamically adapt learning rates across parameters, proving more effective than SGD for our hybrid architecture. Furthermore, a smoothing module was incorporated to refine edge clarity and align the predicted features with the marginal distributions of flood and non-flood regions, thereby reducing noise and weak artifacts. Learnable parameters γc and βc further allow adaptive refinement of feature maps, correcting inconsistencies from earlier layers. Overall, the post-processed masks are both visually precise and quantitatively reliable, supporting practical deployment in real-world flood detection systems.

To assess the segmentation performance of the proposed ViT-UNet with Flood-Aware Refinement Block (FARB), a comprehensive set of metrics was employed, accounting for the class imbalance inherent in flood datasets where non-flooded areas often dominate. The primary focus is on metrics that provide a balanced evaluation of segmentation quality. Mean Intersection over Union (mIoU) serves as a key indicator, measuring the overlap between predicted and ground truth segmentation masks across both flooded and non-flooded classes. Precision and recall are included to evaluate positive predictions and coverage, respectively, while the Dice Score (also known as F1-Score) is introduced as a critical metric, balancing precision and recall to better reflect performance on imbalanced datasets. Accuracy, while reported, is de-emphasized due to its potential to mislead in the presence of class imbalance, as a model predicting the majority (non-flooded) class could achieve high scores without effective segmentation.

The proposed ViT-UNet was benchmarked against recent state-of-the-art segmentation models, including U-Net variants and the Swin Transformer. Our model achieved the highest performance with an mIoU of 0.8640 and test accuracy of 93.09%, demonstrating a significant improvement over conventional CNN-based approaches. This can be attributed to the hybrid architecture, which leverages the ViT encoder for global context modeling and the U-Net decoder for precise spatial recovery.

All baseline models U-Net variants and Swin Transformer were retrained and evaluated using our UAV dataset with identical data splits, augmentation strategies, and preprocessing steps to ensure fair comparison. Our proposed ViT-UNet leverages the complementary strengths of Vision Transformers for global context modeling and U-Net decoders for fine spatial recovery. This combination, along with the Flood-Aware Refinement Block (FARB), enables superior boundary delineation and noise suppression in complex flood scenarios, as supported by both quantitative and visual results.

While other state-of-the-art segmentation architectures such as DeepLabv3+ and SegFormer offer strong performance in general-purpose vision tasks, they were not included in this study due to their significantly higher computational demands when applied to 512 × 512 UAV imagery. These models often require larger GPU memory and longer inference times, making them less practical for real-time flood monitoring in resource-constrained environments. Our benchmark selection thus focused on models that are widely adopted, flood-specific, and aligned with real-world deployment requirements.

While the total training duration for 50 epochs was approximately 1.29 h on an NVIDIA RTX 4080 GPU, this was achievable due to the relatively small dataset size (1160 images) and the moderate input resolution (512 × 512). A batch size of 8 and early convergence around epoch 45 further reduced training time. Although the measured inference speed of 0.15 s per image on the same GPU suggests real-time capability in high-performance environments, this does not directly imply feasibility on UAVs or other edge platforms. Real-time performance in such contexts must account for hardware limitations, power constraints, and thermal efficiency.(8)H=12N(N+1)∑i=1kni(R¯i−R¯)2where N denotes the total sample size across all groups, k represents the number of comparison groups, and ni refers to the sample size within group i, R¯ is the group i average rank, and R¯ is average rank across all the groups. Applying the Kruskal-Wallis H-test to the model results yielded a test statistic of 16.73 and an associated p-value of 0.002. Since this p-value falls below the conventional threshold of 0.05 and we confirm that the differences in performance among models are statistically significant. This finding suggests that the improvement of the proposed ViT-UNet model is not due to a random data-dependent fluctuation, but a systematic methodological advancement over the state-of-the-art techniques.

To assess whether the observed performance differences were statistically significant, we conducted a Kruskal-Wallis H-test using mIoU scores from all models under identical data splits. The test yielded a statistically significant result (p = 0.0003), indicating meaningful performance differences. Pairwise comparisons confirmed that the proposed ViT-UNet significantly outperformed both U-Net (p = 0.0032) and Swin Transformer (p = 0.0081), while the full three-way test also highlighted the contribution of FARB (p = 0.0017), as shown in Table 3.

To ensure a fair and consistent comparison, we re-implemented and trained all baseline models, including U-Net Conv2D, Residual U-Net, Attention U-Net, Multi-Residual U-Net, and Swin Transformer, on our augmented UAV dataset using the same experimental settings. These models were selected based on their frequent usage in recent flood segmentation and general semantic segmentation literature. Table 4 presents the implementation of various segmentation models with data augmentation. Results show that data augmentation increases accuracy, recall and precision across models. Classification error rates dropped significantly for models with poor coverage of the original dataset. UNet variants and Transformer architectures showed improved feature extraction and segmentation with augmented data. The Swin Transformer outperformed other approaches, demonstrating improvements in class imbalance as found previously and highlighting needs in data methodology and mining techniques. This confirms augmentation’s importance for segmentation accuracy and model robustness.

Table 5 presents a comparative summary of segmentation performance and computational cost. Traditional CNN models like U-Net Conv2D achieved faster training (01:03 h) but yielded lower mIoU (0.6952), while Swin Transformer demonstrated better context modeling (mIoU = 0.7962) but still struggled with boundary delineation, especially in reflective or ambiguous regions.

The total training time for 50 epochs was approximately 1.29 h on an NVIDIA RTX 4080 GPU. This efficiency stems from the input resolution of 512 × 512 pixels. With a batch size of 8, each epoch involved approximately 145 steps, resulting in 7250 steps in total. Given this configuration, the average processing time per image per step was approximately 0.62 s, which is reasonable given the GPU capabilities and the streamlined hybrid architecture. Additionally, early convergence was observed around the 45th epoch, further reducing unnecessary iterations. While this training speed is feasible under these controlled settings, we acknowledge that training times would increase substantially for larger datasets or higher-resolution inputs due to the quadratic complexity of self-attention in ViT. Therefore, this timing should be interpreted within the context of the current experimental scale.

The multi-residual U-Net, despite improved architecture, required 4.06 h for training and 38.64 min for testing, while yielding an mIoU of only 0.7247—indicating diminishing returns with added complexity. Our model’s architecture, which integrates the ViT’s long-range contextual modeling with the U-Net’s fine-grained spatial recovery, results in a superior segmentation framework suitable for real-time or near-real-time flood monitoring applications. We compare flood segmentation effects from various models including UNet-based architectures, Swin Transformer and our proposed ViT-UNet model, as shown in Fig. 4. The ViT-UNet model yields the highest MIoU of 0.8640, indicating best segmentation performance. UNet Convolutional 2D and ResNet architectures show lower MIoU scores and misclassify larger areas along complex flood boundaries. The Swin Transformer performs competitively due to self-attention mechanisms, though inferior to our proposed model. These results confirm the ViT-UNet hybrid approach better captures local and global flood features and improves segmentation compared to ViT and UNet. The model achieved strong metrics, reaching of 0.1448 and 0.1906 in training and validation loss, demonstrating good generalization.

As shown in Table 6, several recent models have reported strong flood segmentation performance on similar public datasets. For example, the Graph Attention Convolutional U-Net achieved an accuracy of 91% and mIoU of 0.7600, while our proposed ViT-UNet attained a higher accuracy of 93.09% and mIoU of 0.8640. Unlike [33], which focuses on graph-based spatial modeling, our approach explicitly combines global attention from Vision Transformers with the spatial precision of U-Net and further enhances boundary detection using the refinement block. Moreover, compared to [10], which used a color-space segmentation approach with 87.4% accuracy, and [32], which used SegNet with 89.46% accuracy, our model consistently outperformed both in all metrics. These comparisons confirm that our method not only exceeds recent SOTA results in segmentation accuracy but also provides a tailored architecture for flood-specific challenges, validated using a publicly available benchmark.

In Fig. 5, specifically includes challenging and failure cases (bottom row), such as strong reflections or complex terrain, where all models, including our proposed ViT-UNet, demonstrate degraded boundary segmentation. Common failure patterns include missed small flood patches and confusion with urban surfaces or debris. These highlight the need for further data enrichment and targeted attention mechanisms.

The flooded areas are also presented on the segmented regions. Column (b) both the baseline and our model fail to segment regions with intense surface reflections (Image 2), while debris occlusion causes false positives in Image 3, while column (c) is our proposed method. For example in Image 1 the baseline is missing the water shown by the red circle, but ours is able to detect the water. The baseline does not recognize water bodies in the red-circled locations in Image 2, while our method successfully detects them. In Image 3, the baseline predicts water, while ours establishes water from non-water correctly. These outcomes show the improvement of our ViT-UNet model in capturing flood boundaries.

We present a case study comparing segmentation outputs between the baseline models and our proposed ViT-UNet+FARB using test images. The top panel highlights successful cases where the proposed model captures flood boundaries more accurately and with better visual clarity. The bottom panel showcases difficult scenes with water reflections, debris, and occlusions where baseline models tend to fail, either missing flood regions or producing blurred boundaries. These visual comparisons demonstrate that the ViT-UNet+FARB model provides more consistent and refined segmentation across a variety of challenging conditions.

The ViT encoder uses multi-head self-attention (MHSA) to extract global context. The encoder consists of 12 attention layers, a configuration chosen based on extensive experimentation to balance performance and computational efficiency. A comparative analysis of different configurations (6, 8, and 12 attention layers) was conducted, and 12 layers provided the best trade-off between accuracy and training stability. The decision to use 12 layers was based on empirical results, as shown in Table 7, where this configuration demonstrated the highest accuracy of segmentation and Mean Intersection over Union (MIoU). This structure allows the model to effectively capture long-range dependences in flood imagery while maintaining feasible training times.