U-Prithvi: Integrating a Foundation Model and U-Net for Enhanced Flood Inundation Mapping

Deep learning approaches have been implemented for a variety of remote sensing classification problems pertaining to land cover, agriculture, open water bodies, and floods. Floods present a unique problem in that they can be spectrally very complex depending on the geographic area and terrain that the flood is covering. Examples of deep learning in the literature for flood mapping include various applications of CNNs, Multi-Layer Perceptrons, and Recurrent Neural Networks applied to Sentinel-2, Sentinel-1, and CubeSat data [6, 22, 31], a full review of which can be found in [4]. Specific flood mapping architectures have also been developed, such as Siam-DWENet [38] , which takes advantage of transfer learning and an attention mechanism, and DeepFlood, which employs feature-level fusion and classification of optical and radar data [18]. An architecture that is increasingly common for flood-mapping applications is U-Net, which has been used to classify floods based on satellite imagery and Twitter data [30], and for flood segmentation using radar data in southeastern Mexico [27].

The U-Net architecture was first introduced by Ronneberger et. al 2015 [29] as an improvement of fully-convolutional neural networks for biomedical image segmentation applications. The U-Net consists of a series of contracting convolutional layers that gradually reduce spatial dimensionality and increase feature dimensionality. The feature-dense bottleneck is then upscaled with skip connections that help map from a condensed feature space using higher resolution information available in the encoder hidden layers.

The U-Net architecture has been extensively used for medical image segmentation [3]. However, models trained to perform natural image segmentation, like those in medical images or classic computer vision tasks, do not translate directly to remote sensing applications [33]. Recent advances in applying U-Net for remote sensing segmentation problems include combinations of DenseNet and Dilated Convolutions [33] and attention mechanisms [37], as well as specific applications to flood detection using SAR imagery [20].

Deep learning approaches require substantial amounts of ground truth data for training from scratch, which can be costly to obtain. This need for data, combined with the recent popularity of foundation models, large pre-trained models that have greatly impacted fields such as natural language processing and multimodal tasks, has driven research into foundation models specifically for the geospatial field. As evidence of this trend, research on transfer learning for geospatial tasks saw a ten-fold increase in published articles between 2017 and 2022 [21]. The use of foundation models in geospatial tasks can generally be divided into three categories: (1) models trained on natural image datasets, (2) models trained on geospatial datasets, and (3) hybrid models that integrate both approaches [23].

The first category typically involves models trained on datasets like ImageNet [10]. Although this dataset is quite different from geospatial datasets, studies have shown that this approach can still be effective for specific satellite image tasks, such as land-use classification, urban zone classification, or burnt area detection [25, 26]. In the second category, numerous foundation models have been developed specifically for remote sensing data. Some of these models are trained on single-time images but incorporate data from multiple remote sensors [12], while others are trained on multi-temporal datasets, such as the Prithvi model [17] and SatSwinMAE [24].

Several hybrid models combining natural and remote sensing image modalities have emerged in recent years. For instance, the GFM model [23] employs a two-stage pretraining process. It is initially trained on the ImageNet22k dataset and subsequently on the custom GeoPile remote sensing dataset. This approach aims to enhance performance compared to models trained solely on natural image datasets, while also mitigating the expense of training foundation models from scratch on remote sensing data.

Flood segmentation mapping using GeoAI foundation models has demonstrated effective generalization. However, these models often underperform on in-sample test datasets compared to U-Net models [19]. Let us formally define in-sample and out-of-sample test data. An in-sample test dataset consists of samples drawn from the same distribution as the training data (e.g., from the same regions). An out-of-sample test dataset contains samples from a different distribution but shares the same features (e.g., a different region), where the trained models have not seen any samples from that region. We hypothesize that combining Prithvi and U-Net will improve performance for both in-sample and out-of-sample use cases. Therefore, we developed a novel fusion model integrating Prithvi features into a U-Net architecture. We validated this hypothesis by comparing our model’s generalization and predictive capacity against existing approaches on both in-sample and out-of-sample test datasets. Our experiments show that this combined architecture balances these two performance metrics while requiring fewer training epochs. Finally, we discuss the potential extension of this approach to other foundation models and applications.

Prithvi is a Vision Transformer model trained using a Masked Autoencoder strategy. While its original decoder is designed for input reconstruction and may not be optimal for segmentation tasks [17], we utilize only Prithvi’s encoder and pair it with a custom decoder similar to that of U-Net.

Our implementation, illustrated in Figure 1, begins with Prithvi’s encoder, producing a 14x14 feature map with 768 filters. This feature map is processed through a single block without parameter modification, then upsampled using four transposed convolutional layers to achieve the original input shape. A final convolutional layer with Softmax activation serves as the classifier. Prithvi and U-Net serve as baselines in our experiments.

U-Prithvi, the proposed model, integrates the strengths of both the Prithvi model and the U-Net architecture, combining Prithvi’s capability for capturing global context with U-Net’s proficiency in fine-detail segmentation. The schematic representation of the U-Prithvi architecture is illustrated in Figure 2.

The input passes through both the U-Net and Prithvi encoders, generating two 14x14 feature maps with 768 filters each. These feature maps are aggregated and subsequently passed into the decoder component, which comprises four upsampling blocks with skip connections to the U-Net encoder.

To address the potential training imbalance between the pretrained Prithvi encoder and the untrained U-Net encoder, we introduce a novel RandomHalfMaskLayer (RHM layer). This layer, positioned after the concatenation step, probabilistically masks either Prithvi’s or U-Net’s feature maps – or leaves both unmasked – during training. During inference, the layer has no effect on the input. We anticipate that this approach will promote balanced training and ensure optimal contributions from both components in the final segmentation.

Despite the existence of several approaches that combine transformer and U-Net architectures, most of these train the architecture from scratch and do not account for using fixed pretrained architectures. The Prithvi foundation model, being a ViT-based architecture, does not include multiple stages with varying resolutions, making it incompatible with the methods presented in the related work section 2.2.1. A significant advantage of our approach is its flexibility: Prithvi can be substituted with any other foundation model without requiring substantial modifications to the architecture.

Using the Sen1Floods11 [1] dataset, we evaluated the performance of our U-Prithvi architecture and compared it against standalone U-Net and Prithvi-decoder models from [19]. The Sen1Floods11 dataset contains 446 image samples paired with hand-labeled masks identifying flooded (water and flood) and non-flooded areas across various regions, including Ghana, India, the Mekong River, Nigeria, Pakistan, Paraguay, Somalia, Spain, Sri Lanka, and the USA. Each sample has a resolution of 10 meters and a pixel size of $512\times 512$. Six Sentinel-2 bands (RGB, NIR, SWIR1-2) are used as input to our models to align with the Prithvi encoder’s input requirements.

Our preprocessing pipeline includes data normalization (mean 0, variance 1), random cropping to $224\times 224$ (to meet Prithvi encoder’s input requirements), and random horizontal and vertical flips to augment the data. For performance evaluation, we use two sets provided in the dataset: (1) an in-sample test set with samples from the same regions as the training data, and (2) an out-of-sample test set containing samples from Bolivia, which were not used during training. To compare model performance, we use Intersection over Union (IoU) as our primary metric and accuracy (Acc) as a secondary metric. For each metric, we calculate both the mean (mIoU, mAcc) across classes and individual values for each class (IoU${}_{\text{flood}}$, IoU${}_{\text{nonflood}}$, Acc${}_{\text{flood}}$, Acc${}_{\text{nonflood}}$). These values are derived from true positive (TP), false positive (FP), true negative (TN), and false negative (FN) counts as follows:

The first experiment is designed to answer our first research question. The baseline performance of Prithvi versus U-Net for this application of flood inundation mapping was established in the paper published by [19]. We utilized the same in-sample and out-of-sample test sets as well as computed accuracy metrics to maintain consistency and comparability. These results show that U-Net_Base outperforms Prithvi_Base across all accuracy metrics when evaluated on in-sample test set as shown in Table 1. However, Prithvi_Base outperforms U-Net_Base when evaluated on out-of-sample Bolivia set. These results will be used as our baseline to compare to when analyzing the results of our experiments.

Our second experiment is designed to address our second and most important research question. By integrating the U-Net and Prithvi models, we anticipate that the U-Prithvi model will perform effectively on both same-distribution and unseen data, while demonstrating fast convergence for both scenarios. Table 1 presents our performance results compared to U-Net_Base and Prithvi_Base, evaluated on both the in-sample test set and the out-of-sample Bolivia dataset. Our findings indicate that the U-Prithvi model outperforms both approaches on out-of-sample data. For in-sample data, the performance of U-Prithvi falls between the two models, with U-Net dominating in this context. This suggests that the U-Prithvi architecture combines the ability to achieve strong predictive performance on datasets similar to the training set with the capacity to leverage the pre-trained foundation model for superior generalization on unseen data, without requiring additional training cycles.

Our third research question is addressed in this section. The U-Net architecture is widely used in semantic segmentation [32] and is less computationally and data-intensive to train compared to U-Prithvi. Because of this, we expanded the architecture from [19] and found that increasing the size of the U-Net model can yield similar performance for the benchmark dataset.

Our U-Net implementation, shown in Figure 4, comprises three main components: Encoder, Bottleneck, and Decoder. Each component includes several blocks, each consisting of a convolutional layer, Batch Normalization, and a ReLU activation function. The stride and type of convolution (normal or transposed) determine whether the block reduces, maintains, or increases feature map resolution.

The encoder downscales the input into a compact feature representation, using four downsampling blocks (stride 2) that reduce the resolution from 224 to 14 while increasing filter depth from 6 bands to 768. A single bottleneck block preserves the resolution and feature count. The decoder then upscales the feature map back to the input image resolution through four upsampling blocks, each concatenated with the corresponding encoder block. A final convolutional layer provides the desired output shape. This yields a total of over 37 million trainable parameters, compared to the 29 million reported in [19].

The performance of our revised U-Net model is presented in Table 2 alongside the performance of U-Prithvi in both datasets. Our results show that for the Sen1Floods11 dataset, it is possible to achieve results comparable to the performance of U-Prithvi using the U-Net model with a higher count of parameters. This is true for both in-sample and out-of-sample test sets. In the case of the Bolivia dataset, we find that performance is practically identical compared to the U-Prithvi model. This implies that a larger model is still able to capture the dynamics of the data, even compared to a foundation model. Nevertheless, it is unclear if this result would also be evidenced using a different dataset.

This section addresses our fourth research question and describes several additional experiments conducted to investigate fine-tuning the architectures. We conducted a series of experiments with various parameter configurations to find the optimal performance of our model. The following sections provide a detailed analysis of these parameters and their impact on the model’s performance.