Squeezing Context into Patches: Towards Memory-Efficient Ultra-High Resolution Semantic Segmentation

1 Introduction

Semantic segmentation is a foundational challenge in computer vision with applications ranging from land cover mapping in remote sensing to object extraction in medical imaging and scene parsing in autonomous driving. Deep learning approaches have shown impressive performance in recent years [Azad et al., 2024; Badrinarayanan et al., 2017]. However, the increasing resolution of images captured by modern sensors presents a significant obstacle to semantic segmenting ultra-high-resolution images (UHRSS) on memoryconstrained edge devices [Zhu et al., 2024]. Our focus is on developing an accurate and memory-efficient UHRSS method tailored to the limitations of edge equipment, offering a promising solution for this challenging scenario.

∗Corresponding Author

To address the above challenge, the downscaling paradigm is employed, which involves resizing the UHR images to a more manageable size before feeding them into the segmenter. However, this approach leads to coarse predictions due to the loss of fine-grained detail [Minaee et al., 2022]. An alternative is the stride inference paradigm, which works by cropping image patches from the original UHR image and processing these patches individually. While this method preserves more detail, it suffers from insufficient contextual information. Ultimately, both strategies fail to strike a balance between detail preservation and context utilization.

In recent years, various advanced methods have been proposed to address the trade-off between local detail preservation and the need for global context understanding in parsing UHR images. These methods can be categorized into detail refinement [Huynh et al., 2021; Cheng et al., 2020], shallow and deep fusion [Guo et al., 2022; Ji et al., 2023b; Ji et al., 2023a], and cropped global and local fusion [Li et al., 2021; Ding et al., 2022; Zhang et al., 2024; Zhu et al., 2024], as depicted in Figure 1 (a)-(d). In the detail refinement paradigm, a refining module is introduced to correct the coarse predictions obtained from down-scaled images based on either the detailed output [Huynh et al., 2021] or global predictions [Cheng et al., 2020]. However, this approach requires pre-processing the global predictions, leading to a high computational cost. The shallow and deep fusion paradigm utilizes a shallow branch to process the original UHR images and a deep segmenter to parse the down-scaled images. The cropped global and local fusion paradigm parses global patches by a shallow branch and the local patches by a deep branch. Despite these paradigms achieving better performance than generic ones, they still suffer from high memory consumption due to the multi-branch architecture.

In this study, we introduce a novel single-branch UHRSS approach that incorporates context information into patches, illustrated in Figure 1 (d). Our method processes images in a single branch akin to slide inference, compressing redundant context surrounding local patches using a Context Squeezing Module (CSM). By slightly downsizing local patches and squeezing in context data efficiently, our CSM achieves a balance between global context awareness and local detail retention. Additionally, we propose an auxiliary super-resolution decoder (SRD) assisted local feature alignment loss (LFA) to enhance segmenter detail features, aiding in preserving local intricacies when analyzing the squeezed global patch. This approach enhances context and detail information without imposing excessive memory demands during testing, elevating segmentation accuracy for UHR image processing compared to conventional methods.

The key contributions of this work are as follows:

• We introduce SCPSeg, a method for semantic segmentation of UHR images that efficiently incorporates context into local patches within a single branch.
• Specifically, we propose a context squeezing module (CSM) to merge global context understanding with local detail retention effectively.
• Furthermore, we employ a multi-task learning strategy and propose a local feature alignment loss to enhance detail preservation within the segmenter.
• Through extensive experiments, SCPSeg demonstrates an exceptional balance between memory efficiency and segmentation accuracy, outperforming conventional semantic segmentation methods without increasing memory demands during testing.

2 RelatedWorks

2.1 Generic Semantic Segmentation

Semantic segmentation has made notable advancements in recent years with the development of convolutional neural networks [He et al., 2016] and vision transformers [Xie et al., 2021]. Most of these methods follow an encoderdecoder architecture [Ronneberger et al., 2015] inspired by the first fully convolutional networks for semantic segmentation [Long et al., 2015]. To enhance segmentation accuracy, several strategies are commonly adopted, including increasing the receptive field [Zhao et al., 2017; Chen et al., 2018], integrating global and local features [Lin et al., 2017; Fu et al., 2019], and maintaining high-resolution branches [Yu et al., 2018; Sun et al., 2019; Zhao et al., 2018]. In efforts to reduce computational costs, numerous real-time semantic segmentation methods have been developed [Pan et al., 2023; Xu et al., 2023]. These methods typically feature a lightweight encoder [Fan et al., 2021; Sandler et al., 2018], an efficient segmentation head [Li et al., 2020; Xie et al., 2021], and sometimes an auxiliary training decoder for detailed prediction [Fan et al., 2021; Zhao et al., 2018]. However, these generic semantic segmentation methods face challenges related to high memory consumption when processing UHR images.

2.2 Semantic Segmentation for UHR Images

Despite the potential of employing slide inference and image down-scaling techniques to alleviate the substantial memory overhead associated with parsing Ultra-High-Resolution (UHR) images, these methods often suffer from decreased accuracy due to the loss of contextual information and detail blurring. To address the challenge of achieving both high accuracy and efficient memory utilization for UHR image segmentation, a set of specialized methods has been developed, primarily falling into two categories: detail refinement [Xia et al., 2016; Cheng et al., 2020; Huynh et al., 2021; Kirillov et al., 2020] and global-local fusion [Zhao et al., 2018; Chen et al., 2019; Li et al., 2021; Guo et al., 2022; Ji et al., 2023b; Ji et al., 2023a]. Detail refinement approaches typically involve segmenting high-resolution scenes in multiple stages. For example, the Cascade Segmentation Refinement model [Cheng et al., 2020] generates detailed results by parsing coarse predictions with high-resolution images. Likewise, MagNet [Huynh et al., 2021] parses high-resolution image patches by integrating low-resolution outputs from preceding stages. However, this iterative refinement paradigm makes them computationally intensive and impractical for real-time applications. Global-local fusion often segments high-resolution images using a multi-branch encoder, where a local branch extracts detailed information, and global patches collect contextual details. For example, recent approaches such as ISDNet [Guo et al., 2022], FCtL [Li et al., 2021], and GINet [Zhu et al., 2024] use lightweight shallow CNNs to extract detailed features in the high-resolution branch and an efficient image encoder for context features in the lowresolution branch. Nevertheless, the incorporation of multibranches introduces additional computational and memory costs, posing challenges for deployment on edge devices.

3 Method

This section outlines the architecture and functionalities of the proposed SCPSeg and its components in detail.

3.1 Overview of the ProposedMethod

This work presents a novel ultra-high-resolution segmentation approach, named SCPSeg, as illustrated in Figure 2. SCPSeg effectively integrates global context and local details within a single-branch encoder fE, achieving high accuracy with limited memory consumption. The fundamental concept behind SCPSeg involves compressing global context information into local patches, enabling the segmentation network to sense a broader image field. To facilitate this, we employ a context squeezing module (CSM) that strategically aggregates context information into local patches, ensuring effective utilization of both macro-level context and microlevel details. Furthermore, to enhance the preservation of local details, we introduce a super-resolution guided local feature alignment technique (LFA), that helps the segmentation decoder yield detailed predictions. Additionally, the singlebranch encoder processes ultra-high-resolution (UHR) images in a sliding window manner. A series of global patches XG are cropped from the origin UHR image XUHR and fed into the encoder. The overall architecture is optimized using a multi-task learning strategy, which combines segmentation loss LCE, super-resolution loss LSR, and a novel local feature alignment loss LLFA to ensure comprehensive training and achieve optimal performance.

3.2 Context SqueezingModule (CSM)

The CSM is the cornerstone of our proposed SCPSeg framework. Its primary function is to compress and integrate global context information into local patches through a series of interpolations. The CSM operates on the principle of context compression, where redundant pixels in the global context are removed. This process involves the following steps:

Given a square-like global patch XG with size of G where its centric is the local patch XL with size of L. We first resize the local patch to a down-scaled local patch XD:

where ∇() represents the resize operation. Secondly, we squeeze the context information (the blue part in Figure 2 (b)):

where stack () indicates the tensor stack operation. Finally, we concatenate the squeezed context Xp1 C , Xp2 C , and the scaled local patch Xcs L to get the final squeezed patch XCS:

where cat () indicates the tensor concatenate operation. The squeezed context and local patch XCS has the same size as the origin local patch XL.

The final squeezed patches contain both highly squeezed global context and slightly squeezed local details, providing a comprehensive representation of the image. The effectiveness of the CSM lies in its ability to compress context information without losing critical details. In this way, the CSM ensures that the segmenter can perceive a larger image field while maintaining high accuracy. This approach is particularly beneficial for UHR images, where both global context and local details are essential for accurate segmentation.

3.3 Local Feature Alignment (LFA)

To further enhance the preservation of local details, we introduce a super-resolution guided local feature alignment technique. This technique leverages the high-resolution details features in the super-resolution decoder (SRD) hSR to guide the alignment of local features in the segmentation decoder hSS. The primary goal is to ensure precise alignment of local features, leading to more accurate segmentation outputs.

Super-resolution Decoding

For the auxiliary super-resolution task, we utilized several convolutions and de-convolutions to construct the SRD hSR. It decodes the encoded local features and recovers the original high-resolution local patch:

where fE is the image encoder. Then, a super-resolution loss is computed between the predicted high-resolution local patch PSR and the real high-resolution local patch XL:

where ∥·∥p=1 is the L1 norm. ccrop () indicates the centercrop function for cut out the features of XL.

Local Feature Aligning

The features of the SRD contain rich structural information of the original image, although they do not explicitly map the categories. Since computing similarity over all the feature pixels will lead to heavy memory and computational cost, we can effectively model the pixel relationship within sliding windows with a size of k×k. These relationships can implicitly deliver semantic information, thus benefiting the task of semantic segmentation.

Firstly, the segmentation features PSS and super-resolution features PSR are cropped by the slide windows. Then, we compute the self-similarity matrix within the sliding window:

where Si,j indicates the similarity between pixel i and pixel j within a sliding window. Besides, PSS = hSS (fE (XCS)). The decoded features P can be crop(PSS) and crop(PSR). At last, to optimize the alignment process, we introduce a local feature alignment loss to align the features by minimizing the similarity matrixes between slide windows at the same location:

3.4 Optimization

The final objective function is computed as follows:

where LCE is the cross-entropy loss for segmentation. LSR is the super-reolution loss defined in Equation.6. LLFA is the local feature alignment loss defined in Equation.8. w1 and w2 are set as 0.5 and 0.5, making the gradients of the loss components ranges comparable.

4 Experiments

4.1 Datasets

We conduct experiments on four UHRSS datasets to comprehensively evaluate the effectiveness of our proposed method.

ISPRS Potsdam. This dataset is a land cover mapping dataset collected in an urban area. It contains 38 tiles with the size of 6000 × 6000. The train, val, and test sets are split to 18, 6, and 14 tiles following [Zhang et al., 2024]. The bands of this dataset include red, green, blue, near-infrared, and DSM information.

BLU. This dataset is collected in urban and rural areas. It contains 4 tiles with the size of 15680 × 15680. We crop the original URH data into nonoverlapped tiles with resolution 2048×2048 following [Ding et al., 2022]. It is split into 192, 28, and 32 tiles for training, validating, and testing, respectively.

DeepGlobe. This is a land cover mapping dataset collected in both urban and rural areas. It contains 803 tiles with the size of 2448×2448. We split it into 455, 142, and 206 subsets for training, validating, and testing following previous work [Chen et al., 2019].

Inria Aerial. Inria Aerial is a building extraction dataset collected in urban areas. It contains 180 UHR images with 5000 × 5000 pixels. We split it into 126, 27, and 27 subsets for training, validating, and testing following previous work [Chen et al., 2019].

4.2 EvaluationMetrics

We utilize the mean intersection-of-union (mIoU), mean F1 (mF1), and accuracy (Acc) to assess the effectiveness. The GPU memory cost is monitored by the "gpustat" tool.

4.3 Implementation Details

We utilize Deeplabv3+ with a ResNet-18-d8 as the basic segmenter. All the experiments are conducted in a single Nvidia RTX4090. The encoders were pre-trained on the ImageNet1K dataset. SGD is used to optimize the neural networks. The learning rate is initially set to 0.01 and decayed by a cosine learning rate policy after each iteration. The training iteration number is set to 40000. The slide window size k in LFA is set to 7. For the ISPRS Potsdam, BLU, and Inria Aerial datasets, we set G = 512, L = 256, and D = 192, respectively. The training batch size is set to 16. For DeepGlobe datasets, we set G = 1024, L = 512, and D = 384, respectively. The training batch size is set to 8.

4.4 Comparing with Other Methods

In this section, we compare our method with other stateof-the-art ones including FCN-8s [Long et al., 2015], DeepLabv3+ [Zhao et al., 2017], PSPNet [Zhao et al., 2017], ST-UNet [He et al., 2022], BiTSRS [Liu et al., 2023a], UNetFormer [Libo et al., 2022], GLOTS [Liu et al., 2023b], CFNet [Peng et al., 2022], LANet [Ding et al., 2021], GLNet [Chen et al., 2019], WiCoNet [Ding et al., 2022], TCNet [Zhang et al., 2024], U-Net [Ronneberger et al., 2015], ICNet [Zhao et al., 2018], BiseNetv1[Yu et al., 2018], STDC [Fan et al., 2021], CascadePSP [Cheng et al., 2020], PointRend [Kirillov et al., 2020], MagNet [Huynh et al., 2021], ISDNet [Guo et al., 2022], FCtL [Li et al., 2021], WSDNet [Ji et al., 2023b], and GINet [Zhu et al., 2024] on ISPRS Potsdam, BLU, DeepGlobe, and Inria Aerial, in terms of mIoU (%), mF1 (%), OA (%), Mem (M). Some of these methods are generic semantic segmentation methods denoted as "Generic Methods" and the other methods are specially designed for UHR images, denoted as "UHR Methods". The generic methods tagged by "*" are tested in the slide inference mode.

Results on the ISPRS PotsdamDataset. We compare our SCPSeg with the aforementioned methods on the ISPRS Potsdam testing set. Due to the high inter-class similarity between buildings and impervious surfaces, this dataset poses significant challenges. Table 1 presents the performance of each semantic segmentation method. The experiments demonstrate that our method outperformed all others in terms of mean Intersection over Union (mIoU), mean F1 score (mF1), and overall accuracy (Acc). Specifically, we surpass the UHRSS methods (GLNet, WiCoNet, and TCNet) by substantial margins, clearly demonstrating the effectiveness of our segmentation approach and the improvements in performance. Additionally, we evaluate memory usage during the testing phase, and the experimental results indicate that our approach achieves the highest scores among all models in this regard. With such impressive performance, our method strikes an excellent balance between accuracy and memory cost.

Results on the BLU Dataset. Table 2 presents the performance metrics for both generic and UHRSS methods on the BLU dataset. In all evaluation scenarios, our approach demonstrated the highest accuracy coupled with the lowest memory usage. Overall, the UHRSS methods consistently outperformed the generic methods, highlighting the advantages of utilizing extensive contextual information in semantic segmentation tasks. Notably, our method surpassed the performance of the cropped global and local fusion approach, WiCoNet and TCNet, indicating the effectiveness of our context-squeezing strategy, which achieves superior results with lower memory requirements compared to this architecture.

Results on the DeepGlobe Dataset. As shown in Table 3, we first compare our method with the above-mentioned methods on the DeepGlobe testing set. The WSDNet achieves the best performance because it encodes the whole image which contains all the context information. Our method achieves comparable results to other UHRSS methods. With such impressive performance, our method is economical in the memory cost, attaining a satisfactory balance between effectiveness and memory efficiency.

Results on the Inria Aerial Dataset. We present a comparison using the Inria Aerial test dataset in Table 4. This dataset is particularly challenging, featuring images with approximately 25 million pixels, around four times that of the DeepGlobe dataset. The substantial size of this UHR dataset places significant demands on memory usage for semantic segmentation tasks. Our experimental results demonstrate that SCPSeg significantly outperforms other UHRSS methods across all accuracy metrics while maintaining superior memory efficiency.

4.5 Ablation Study

In this section, we investigate the proposed modules and demonstrate their effectiveness. All the ablation studies are performed on the ISPRS Potsdam validating set.

Effectiveness of the SCPSeg Components. We conduct experiments to verify the effectiveness of the proposed SCPSeg based on various basic segmenters, such as FCN-8s [Long et al., 2015], Deeplabv3+ [Chen et al., 2018], PSPNet [Zhao et al., 2017], SegFormer [Xie et al., 2021], TopFormer [Zhang et al., 2022], DDRNet [Pan et al., 2023], PIDNet [Xu et al., 2023]. The experiment results show that the segmentation accuracy (mIoU) improves around 1.2% after incorporating the proposed method, which indicates the effectiveness of our SCPSeg. We also verify the effectiveness of different components in SCPSeg, as shown in Table 6. Firstly, incorporating CSM leads to considerable performance gain, indicating the effectiveness of the context-squeezing paradigm. Secondly, the auxiliary SRD improves the performance of the segmentation task slightly, which proves that super-resolution guided multi-task learning is useful for parsing down-scaled local images. Thirdly, a considerable improvement is gained after the utilization of LFA because the local feature alignment loss helps transfer mutual knowledge from different decoders. Lastly, the memory footprint is not increased, indicating our method is memory-efficient in the testing stage.

Influence of Window Size. We conduct experiments to analyze the influence of the context window size G and downscaled local size D. As Figure 5 shows, the best performance gains when the context size is set to 512 × 512 and the down-scaled local size is set to 192 × 192. As context size increases, the performance decreases because squeezing too much context information into a fixed area results in context degradation. Moreover, leveraging a few context information improves performance only marginally. As the down-scaled local size increases, the squeezed context information decreases, which results in a decrease in performance.

As the down-scaled local size decreases, the local images become blurred, which leads to poor segmentation performance. Therefore, a moderate context window size and a properly down-scaled local size are important to balance the global information and local details.

Visualization of Features. To show the effectiveness of the proposed local feature alignment, we visualize the features of super-resolution and semantic segmentation branches in Figure 6. By comparing (b) and (c) in Figure 6, we can easily find that the super-resolution branches contain more complete structure information of objects. The locally aligned semantic features in Figure 6 (d) have clearer edges than the unaligned features. It indicates that the local-consistent relationship between super-resolution features can be effectively transferred to the semantic segmentation branch, thus benefiting the task of semantic segmentation.

5 Conclusion

We proposed a memory-efficient semantic segmentation method, SCPSeg, for the UHR images. It leverages both the global context and local fine structure effectively to enhance the segmentation in the scenario of ultra-high resolution without sacrificing the GPU memory usage. Our work provides a bright new paradigm for parsing UHR images where squeezing redundant context information is useful. Extensive experiments indicate that our method provides a good trade-off between accuracy and memory efficiency. In the future, we plan to extend SCPSeg to other high-resolution vision tasks, such as object detection and instance segmentation, to further validate its generalizability. Besides, adaptive mechanisms to dynamically adjust memory usage based on image complexity will be investigated.

Acknowledgments

We sincerely thank the reviewers for their valuable and constructive comments, which have greatly improved this article. This work was supported in part by the National Key Research and Development Program of China under Grant 2021YFA0715203; in part by the Science Fund for Creative Research Groups of the National Natural Science Foundation of China under Grant 62221002; in part by the National Natural Science Foundation of China under Grant 62201207 and 62371185; in part by the Natural Science Foundation of Hunan Province under Grant 2023JJ40163; in part by the Science and Technology Innovation Program of Hunan Province under Grant 2023RC3124 and 2024RC1030.