Urban Waterlogging Detection: A Challenging Benchmark and Large-Small Model Co-Adapter

The internal structure of the enhanced-image adapter module is presented in Fig. 3 (a), which mainly consists of a histogram equalization, a high-frequency filter and MLP blocks. Given that the features of water are not pronounced in most challenging scenarios, we first conduct histogram equalization operation to highlight the contrast and texture of input image. The enhanced image is then passed through a high-frequency filter to extract high-frequency information beneficial for segmentation, and converted into frequency patch embedding. The patch embedding of original input image is reduced in dimension by fully-connected layer (FC) and added to the frequency patch embedding. This fused feature is mapped by $N$ individual MLP blocks and one parameter-shared MLP, and then merged with the original features of each transformer block in the SAM image encoder.

SAM originally considered two sets of prompts, sparse (boxes or points) and dense (masks), both of which can provide spatial location information for the object to be segmented. We propose to generate such prompts via a spatial prompter utilizing the outputs of the small model. The masks $M_{\mathrm{small}}$ predicted by small model can be directly used as the dense prompts. Further processing of the masks $M_{\mathrm{small}}$ yields either boxes or points as sparse prompts. For box prompts, we take the bounding boxes of the regions composed of all pixels predicted as the foreground in the masks, represented by the coordinates of the top-left and bottom-right corners of the boxes. For point prompts, we divide the masks into multiple grid regions. In each grid area $G_{g\times{g}}$, all pixels are divided into a positive point set $I_{P}=\{(i,j)\mid{M_{\mathrm{small}}(i,j)}\geq{\tau}\}$ and a negative point set $I_{N}=\{(i,j)\mid{M_{\mathrm{small}}(i,j)}<{\tau}\}$, where $\tau$ is the preset threshold. If the set $I_{P}$ is not empty, we select one point $p\in{I_{P}}$ with the highest prediction confidence as the positive prompt and the prompt label is set to 1. Otherwise, we take one point $p\in{I_{N}}$ with the lowest prediction confidence as the negative prompt and the prompt label is set to 0. All points from the grids together form the grid point prompts, represented by their coordinates and labels. Despite providing three different types of prompts, considering the redundancy of information, we only select one prompt as the final spatial prompt fed into the prompt encoder and obtain spatial prompt $\boldsymbol{e}_{\mathrm{Spa}}$. It is noteworthy that the prompt encoder processes dense and sparse prompts differently. Dense prompts are embedded using convolution before adding to the image embedding and sparse prompts are encoded with positional encoding to generate the corresponding sparse embedding.

The image embedding of large model contains rich semantic information. Therefore, we propose a prototype learning-based semantic prompter, which leverages useful foreground features from large model to generate semantic prompts. The above process is detailed in Fig. 3 (b). A projector is first adopted to map the image embedding $\boldsymbol{Z}\in\mathbb{R}^{H\times{W}\times{D}}$ into the projected embedding $\boldsymbol{\bar{Z}}\in\mathbb{R}^{H\times{W}\times{D}}$. Inspired by [42], we randomly initialize a group of $C\cdot K$ prototypes, i.e., $\{\boldsymbol{p}_{c,k}\in\mathbb{R}^{D}\}^{C,K}_{c,k=1}$ in the embedding space, where $C$ is the number of categories and each class is represented by $K$ prototypes. For each pixel sample $\boldsymbol{\bar{z}}_{i,j}\in\mathbb{R}^{D}$, $i\in\{1\cdots W\}$, $j\in\{1\cdots H\}$ in the projected image embedding $\boldsymbol{\bar{Z}}$, we respectively calculate its cosine similarity with each prototype $\boldsymbol{p}_{c,k}$ to obtain a similarity vector $\boldsymbol{s}_{i,j}\in\mathbb{R}^{J}$ located at the ($i$, $j$) position in the similarity matrix $\boldsymbol{S}\in\mathbb{R}^{H\times{W}\times{J}}$, where $J=C\cdot K$. The category of the prototype corresponding to the maximum value in the similarity vector $\boldsymbol{s}_{i,j}$ is assigned to the pixel sample $\boldsymbol{\bar{z}}_{i,j}$ as the pseudo label $c_{i,j}^{*}$. The pseudo mask generation (PMG) process can be represented as follows:

where $\langle\cdot,\cdot\rangle$ denotes the cosine similarity operator. The pseudo mask $\boldsymbol{M}$ is then one-hot encoded and employed in conjunction with the original image embedding $\boldsymbol{Z}$ to compute the masked average pooling (MAP). This filters out irrelevant background features and preserves significant foreground features, and derives the semantic embedding as follows:

where $Concat(\cdot)$ denotes the concatenation operator and $\boldsymbol{e}_{\mathrm{Sem}}^{c}$ represents the semantic embedding of class $c$ ($c\in\{1\cdots{C}\}$), computed as follows:

where $\odot$ denotes the Hadamard product. The prototype $\boldsymbol{p}_{c,k}$ is momentum-updated after each training iteration according to the center of the $k$-th sub-cluster of the training samples assigned to the $c$-th class via online clustering. Meanwhile, a prototype loss $\mathcal{L}_{\mathrm{proto}}$ in [42] is utilized to optimize the large model.

We introduce a spectrum-based style prompter that extracts image-specific style embedding from the input image as the third type of prompt. The style of an image refers to features such as color and texture. In the context of urban waterlogging detection, these features to some extent can reflect information about illumination and the scene, where illumination is a critical factor causing difficulty. Specifically, We first perform a 2D Fast Fourier Transform (FFT) on the input image $f(x,y)$ to acquire its frequency spectrum $F(u,v)$, which can be represented as:

where $A(u,v)$ is the amplitude spectrum and $\Phi(u,v)$ is the phase spectrum. $u$ and $v$ are the frequency coordinates. Since the amplitude spectrum reflects the image style while the phase spectrum interprets the image content, we thus reconstruct the image solely by the amplitude spectrum using the 2D Inverse Fast Fourier Transform(iFFT):

The reconstructed amplitude-only image contains style information, which is then encoded as style embedding prompt $\boldsymbol{e}_{\mathrm{Sty}}$ by a convolutional block.

We introduce a straightforward one-stage training strategy, as depicted in Fig. 4 (a). The image encoder of the large model is frozen and the remaining parts are optimized together. We employ a combination of the focal loss $\mathcal{L}_{\mathrm{focal}}$ [19], cross-entropy loss $\mathcal{L}_{\mathrm{ce}}$, IoU loss $\mathcal{L}_{\mathrm{iou}}$ and the prototype loss $\mathcal{L}_{\mathrm{proto}}$ [42] for the large model:

The total loss is given as:

where $\mathcal{L}_{\mathrm{small}}$ is the original loss of the small model (the loss function depends on specific model selection, and Mask-RCNN [10], U2Net [26] and SiNet [6] are tested in experiments) and $\lambda$ is a hyper-parameter.

A two-stage training strategy is provided to mitigate issues related to synchronization difficulties and gradient conflicts that may arise during the joint optimization of large-small models, as shown in Fig. 4 (b).

In the first stage, the triple-s prompt adapter module, the dynamic prompt combiner and the prompt encoder are not involved. The image encoder remains frozen, while the remaining modules of the large model and small model are independently optimized by their own loss functions, i.e., $\mathcal{L}_{\mathrm{large}}^{\mathrm{s1}}$ and $\mathcal{L}_{\mathrm{small}}^{\mathrm{s1}}$, respectively. The loss function of small model is the same as Eq. 8. The training loss of the large model for the first stage is defined as:

In the second stage, we load the parameters of the modules (small model, HE-adapt and mask decoder) trained in the first stage, while integrating all the modules that were not considered previously for training (TSP-adapt, DPC and prompt encoder). With the parameters of image encoder, HE-adapt and the small model fixed, the optimization objective for the second stage is as follows:

For advancing urban waterlogging detection challenge, we develop a UW-Bench dataset containing a total of 7,677 images from various scenarios, such as waterlogging scenes, dry roads and hard cases, such as slippery roads and nighttime roads. Using keywords such as urban waterlogging, waterlogged roads, and monitoring viewpoints, we crawled and filtered relevant images from surveillance videos and handheld cameras. The training set includes 5,584 images, while the test set is provided by Huawei inc. and contains 2,093 images of urban scenes captured by surveillance cameras only. For the test set, we consider general-sample and hard-sample cases. Some examples from the training set and test set in our UW-Bench are described in Fig. 5, which indicates the difficulty of detecting waterloggings. In the labeling phase, we use EasyData to annotate the dataset with masks. The pixel-level annotation process can be divided into several stages: training, annotation, validation, and correction. We first create some annotation samples and train the annotators to understand the annotation standard. We also assign an inspector to verify the mask annotations. For failed annotations, the inspector gives an explanation and feedback to each annotator to further improve the annotation quality. The overall annotation process ensures the accuracy and reliability of masks in waterlogging regions.

The waterlogging detection can be viewed as a pixel-level binary classification task for segmentation, and the waterlogged region is our prospect of interest. Based on the ground truth masks as well as the predicted waterlogging masks, we exploit the commonly used segmentation metrics, such as Precision, Recall, F1-score, and Intersection over Union (IoU) to evaluate the detection performance.

For the large model, we choose the ViT-B version of the pre-trained SAM as the backbone, and the input image is resized to 1024*1024. In the training phase, the input images are randomly flipped horizontally, and the batch size is set as 2. AdamW optimizer is used with an initial learning rate of 0.0005, and cosine annealing decay is applied. In the testing phase, the final binary mask is obtained by a simple thresholding operation with the threshold set as 0.5. To evaluate our approach more comprehensively, we choose the classic Mask-RCNN [10] for semantic segmentation, U2Net [26] for salient object detection and SINet [6] for camouflaged object detection, respectively, as the our task-specific small models, but not limited to these ones. All experiments are implemented in PyTorch on an NVIDIA Tesla V100S GPU (32G memory). See more implementation details in appendix.

The proposed LSM-Adapter that employs the one-stage training strategy is significantly inferior to the two-stage training strategy by comparing their best performances. This demontrates the proposed two-stage training strategy is more stable in adaptation. We posit that the following factors may impede the effective implementation of the one-stage training strategy. During the early stages of training, the predicted output of the small model is characterized by low accuracy, and add complexity to the training of the large model results in subsequent slow convergence. Concurrently, the joint training of two networks with distinct architectures is highly contingent upon the selection of suitable hyper-parameters to achieve a synchronized optimization process. Otherwise, the optimization objectives may conflict each other, thereby hindering the joint model from attaining the optimal performance.

Under identical conditions concerning the small model and training strategy, we compare three types of spatial prompts. Except LSM-Adapter_U, the performance by using mask as spatial prompt consistently surpasses that of other two types of prompts (box and point) in all scenarios. Although the box prompt in LSM-Adapter_U is the best, the performance gap from the mask prompt is very small, with a mere 0.03$\%$ gap for both F1-score and IoU when employing two-stage training strategy. Moreover, the performance of models employing mask prompts predominantly exceeds that of both large models and their respective small models. A possible explanation for this observation may be that, in comparison to the sparse prompts of boxes and points, mask prompts furnish a more abundant referential information.

To investigate the impact of the scale of the annotated dataset, we randomly select training data with and without waterlogging according to the ratio $r$ and then merge them to form a subset for training. The ratio $r$ is set to 1, 0.5, and 0.25, respectively, where a value of 1 indicates training with the original fully annotated data. Tab. 5 shows that models trained with more annotated data perform better on downstream tasks, indicating that although foundation model possesses powerful generalization capabilities, additional annotated data is still needed to help the model adapt to downstream tasks.

We analyze the impact of two hyper-parameters on performance, the threshold $\tau$ for spatial prompt of point type and the coefficient $\lambda$ in the total loss (Eq. 8) of the one-stage training. As the threshold $\tau$ varies from 0.3 to 0.7, the overall performance of the model changes minimally, reaching optimal performance at a threshold of 0.4. The choice of coefficient $\lambda$ significantly affects the model’s performance. When SINet [6] is used as the small model and $\lambda$ is set to 1, the overall performance of LSM-Adapter_S is optimal. For other values, both the F1-score and IoU decline noticeably. Therefore, we recommend careful selection of the coefficient value when choosing other models as task-specific small models to achieve better performance. Generally, the coefficient should be chosen to keep the loss values of the large and small models within the same order of magnitude.