Leveraging Adaptive Implicit Representation Mapping for Ultra High-Resolution Image Segmentation

The progress in semantic segmentation, which aims to assign distinct categories or classes to individual pixels within an image, has been significantly propelled by the achievements of deep neural networks [long2015fully, krizhevsky2012imagenet, simonyan2014very, lin2017feature, chen2017deeplab, chen2016attention, farabet2012learning, he2004multiscale, mostajabi2015feedforward, shotton2009textonboost, zhang2018context]. However, these methods are primarily focused on tackling the issue of low-resolution image segmentation but fall short of producing high-quality segmentation results. Thus, the task of refining segmentation is introduced to improve the quality of segmentation outcomes for high-resolution (1K$\sim$2K) (HR) or even ultra high-resolution (4K$\sim$6K) (UHR) images. Unlike low-resolution images, HR and UHR require a deeper understanding of semantic details. Conventional techniques rely on graphical models like CRF [chen2014semantic, chen2017deeplab, lin2016efficient, krahenbuhl2011efficient, liu2015semantic] or region growing [dias2019semantic] techniques to achieve HR segmentation. But they mainly focus on low-level color boundaries [zheng2015conditional], limiting their segmentation performance. In contrast, Multi-scale Cascade Networks [chen2019collaborative, cheng2020cascadepsp, he2019bi, qi2018sequential, sun2013deep, zhao2018icnet] employ a recursive approach to progressively integrate multi-scale information. Nevertheless, such a recursive approach inadvertently introduces more computational burdens. Therefore, a Continuous Refinement Network [shen2022high] is proposed to improve the efficiency by continually aligning the feature map with the refinement target. However, it tends to overlook the significance of the limited receptive field of the Convolution Operation in the CNN backbone and also lacks the generalization capability to effectively handle specific categories. To address these concerns, we employ a Transformer-based adaptive implicit representation function to mitigate these limitations and improve segmentation accuracy.

Vision Transformers (ViT) [dosovitskiy2020image], as demonstrated by their remarkable success, is adapted for various computer vision tasks such as semantic segmentation [liu2021swin, zheng2021rethinking] and object segmentation [zhu2020deformable, dai2021dynamic]. However, applying plain Vision Transformer (ViT) models directly to segmentation tasks is challenging due to their absence of segmentation-specific heads. To solve this problem, SETR [zheng2021rethinking] reframes semantic segmentation as a sequence-to-sequence prediction task. [liu2021swin] proposes a versatile hierarchical Vision Transformer backbone, known as Swin Transformer (Hierarchical Vision Transformer using Shifted Windows), using shifted windows to efficiently capture long-range dependencies while maintaining scalability. Segmenter [strudel2021segmenter] employs pre-trained image classification models fine-tuned on moderate-sized segmentation datasets for semantic segmentation tasks. Instead of ViT, [ru2022learning] utilizes Transformers to generate more integral initial pseudo labels for the end-to-end weakly-supervised semantic segmentation task and proposes an Affinity from Attention (AFA) module to learn semantic affinity from the multi-head self-attention (MHSA) in Transformers. Similar to [ru2022learning], we propose to integrate the Transformer architecture into our UHR image refinement segmentation task and leverage the learned semantic affinity from multi-head self-attention to enhance the precision of coarse masks.

Hypernetworks or meta-models, as referenced in [article, littwin2019deep], involve designing models specifically tailored for generating parameters for other models. This parameterization approach not only enhances the model’s expressiveness, as highlighted in [galanti2020comparing, galanti2020modularity], but it also supports data compression through weight sharing across a meta-model [article]. Unlike conventional deep neural networks, where the weights remain fixed during inference, hypernetworks possess the ability to generate adaptive weights or parameters for target networks [li2022misf, guo2021jpgnet, Li_2023_ICCV, zhang2023superinpaint]. We leverage this technique and propose an Adaptive Implicit Representation Mapping Function. Essentially, an additional Convolution Network is responsible for learning the parameters of the implicit function in the segmentation network, enhancing its adaptability to a variety of high-resolution image segmentation tasks.

Image encoder. To investigate the impact of receptive fields on IRM, we construct a straightforward image encoder, i.e. $\phi$, comprising only several convolutional layers, followed by ReLU and batch normalization operations. We then utilize the encoder $\phi$ to extract features $\mathbf{F}\in\mathbb{R}^{h\times w\times D}$ from the input image $\mathbf{I}$ and coarse segmentation mask $\mathbf{M}$, \linenomathAMS

$\displaystyle\mathbf{F}=\phi(\mathbf{I},\mathbf{M}).$

(1)

SIRMF. Following [chen2021learning], the SIRMF is implemented by an MLP network, denotes as $\textit{f}_{\theta}$, and the refined segmentation result $\tilde{\mathbf{M}}_{q}$ at location ${q}$, can be represented as \linenomathAMS

$\displaystyle\tilde{\mathbf{M}}_{q}=\sum_{p\in\mathcal{N}_{q}}\frac{S_{p}}{\sum{S_{p}}}\cdot\textit{f}_{\theta}(z_{p},p-q),$

(2)

where f is the mapping function, $\theta$ is the trainable parameter of f, $\mathcal{N}_{q}$ are neighbor locations around $q$, and $z_{p}$ is the pixel-wise feature vector related to location $p$. The aggregation weights, i.e. area value $\mathbf{S}_{p}$ are calculated based on the relative coordinate offset between $q$ and $p$.

We train the image encoder $\phi$ with various kernel sizes and the SIRMF $\textit{f}_{\theta}$ on the high-resolution image segmentation dataset. Subsequently, we evaluate the performance based on the Intersection over Union (IoU) and illustrate the results in Fig. 2 (a) (see the blue line). We can observe a substantial performance improvement in the segmentation results when enlarging the kernel size. This highlights that larger receptive fields significantly enhance the performance of IRM in the context of image segmentation tasks.

To further analyze the limitations of the current SIRMF, we randomly select three categories of images from the segmentation dataset. Subsequently, we conduct the following experiments:

Exp.1 We utilize the pre-trained image encoder $\phi$ from Sec. 3.1 to extract image features. We train a SIRMF on the entire selected image set and then evaluate the performance of the SIRMF on each category separately (see Fig. 2 (b) blue bar).

Exp.2 We use the same encoder as in Exp.1 to extract image features. But we train three SIRMF, each dedicated to one of the selected categories, and then evaluate the performance of each SIRMF on its corresponding category (see Fig. 2 (b) yellow bar).

From Fig. 2 (b) we can observe that training three SIRMF separately for each category yields better performance compared to training one SIRMF across all three categories. This observation indicates that the existing SIRMF lacks the generalization capability to effectively handle specific categories and tends to overlook global semantic information.

As discussed in Sec. 3, larger receptive fields significantly enhance the performance of IRM. Therefore, instead of employing a CNN-based image encoder like CRM [shen2022high] and LIIF [chen2021learning], we propose an affinity-empowered image encoder that takes advantage of the transformer architecture and semantic affinity. Specifically, given the RGB image $\mathbf{I}\in\mathbb{R}^{H\times W\times 3}$ and the corresponding coarse segmentation mask $\mathbf{M}\in\mathbb{R}^{H\times W\times 1}$, we employ the transformer-based encoder $\omega$ to extract features $\mathbf{F}$, which can be represented as follows:

$$\mathbf{F},\mathbf{ATT}=\omega(\mathbf{I},\mathbf{M}),$$

(3)

where $\mathbf{ATT}\in\mathbb{R}^{hw\times hw\times c}$ denotes the attention map produced by $\omega$ and $hw$ represents the flattened spatial dimensions and $c$ is the number of channels. Note that the coarse segmentation mask is obtained by applying an existing segmentation method such as [long2015fully, chen2018encoder, lin2017refinenet].

To further empower the capability of the transformer-based encoder, we propose learning the affinity from the attention map $\mathbf{ATT}$ to intentionally impose constraints on the encoder $\omega$ during training. Specifically, we obtain the predicted affinity $\mathbf{Aff}_{pre}\in\mathbb{R}^{hw\times hw}$ by feeding the sum of $\mathbf{ATT}$ and its transpose into an MLP network, i.e., $\varphi$, which can be represented as:

$$\mathbf{Aff}_{pre}=\varphi(\mathbf{ATT}+\mathbf{ATT}^{\top}).$$

(4)

We expect $\mathbf{Aff}_{pre}$ to be close to the ground truth affinity $\mathbf{Aff}_{gt}\in\mathbb{R}^{hw\times hw}$ obtained from the ground truth segmentation mask $\hat{\mathbf{M}}$. Specifically, we resize $\hat{\mathbf{M}}$ to $hw\times hw$ to align with the shape of $\mathbf{Aff}_{gt}$ and the affinity between two coordinates $(x_{i},y_{i})$ and $(x_{j},y_{j})$ is assigned as 1 if their classes are identical, and 0 otherwise.

In Sec. 3.2, we identified the limitation in the existing SIRMF (2) — it lacks the capacity to effectively generalize across diverse inputs (see Fig. 2 (b)). To tackle this issue, we introduce SIRMF which can be represented as \linenomathAMS

$\displaystyle\tilde{\mathbf{M}}_{q}=\sum_{p\in\mathcal{N}_{q}}\frac{S_{p}}{\sum{S_{p}}}\cdot\textit{f}_{\psi(\mathbf{F})}(z_{p},p-q),$

(5)

where $\psi(\cdot)$ denotes a hypernetwork that is used to predict the adaptive mapping parameters. Instead of directly optimizing the shared parameter $\theta$ in (2), we employ $\psi$ to predict it adaptively by inputting the feature $\mathbf{F}$. Intuitively, the predicted mapping parameter should capture global semantic information from the input feature $\mathbf{F}$. Consequently, AIRMF can leverage both local information from pixel-wise features and global information from the entire image feature. This design allows AIRMF to yield precise segmentation results (see the red points in Fig. 2 (a) and the red bar in Fig. 2 (b)).

Assume we aim to generate the weight ${\textit{{W}}^{\textit{l}}}\in\mathbb{R}^{n_{\rm in}\times n_{\rm out}}$ and bias ${\textit{{b}}^{\textit{l}}}\in\mathbb{R}^{n_{\rm out}}$ of the l-th linear layer of $f(\cdot)$. To optimize computational efficiency, we do not directly predict the parameter. Instead, we adopt a strategy similar to [skorokhodov2021adversarial], predicting two intermediary rectangular matrices ${\textit{{A}}^{\textit{l}}}\in\mathbb{R}^{n_{\rm out}^{l}\times r}$ and ${\textit{{B}}^{\textit{l}}}\in\mathbb{R}^{{r\times n_{\rm in}^{l}}}$ through $\psi(\cdot)$, where r = 20 for all the layers of $\textit{f}{(\cdot)}$. The linear modulating projection weight ${\textit{{W}}}^{l}$ is then modulated by multiplying these matrices as follow,

$${\textit{{W}}}^{l}={\textit{{A}}^{\textit{l}}}\times{\textit{{B}}^{\textit{l}}}.$$

(6)

Bias parameter ${\textit{{b}}^{\textit{l}}}$ is directly produced due to its small volume. We streamline the AIRMF in Fig. 3.

Network architectures. We employ the Mix Transformer (MiT) introduced in Segformer [xie2021segformer]) as the backbone of image encoder $\omega$. MiT is a more suitable backbone for image segmentation tasks compared to the vanilla ViT [dosovitskiy2020image]. In brief, MiT incorporates overlapped patch embedding to maintain local consistency, spatial-reductive self-attention to accelerate computation, and employs FFN with convolutions to safely replace position embedding. The network of $\phi$ to predict the semantic affinity is implemented by an MLP layer. The deep network $\psi$ employed to generate the flattened parameters of the mapping function $\textit{f}(\cdot)$ is constructed by several Conv+ReLU+BatchNorm layers, followed by a linear layer.

Loss. During training, we use the combined loss functions to optimize overall performance. To improve AIRMF, we utilize the cross-entropy loss $\mathcal{L}_{ce}$, L1 loss $\mathcal{L}_{1}$, L2 loss $\mathcal{L}_{2}$ and gradient loss $\mathcal{L}_{grad}$

$$\mathcal{L}_{AIRMF}=\lambda_{1}\mathcal{L}_{1}+\lambda_{2}\mathcal{L}_{2}+\lambda_{3}\mathcal{L}_{ce}+\lambda_{4}\mathcal{L}_{grad},$$

(7)

where $\lambda_{1}=0.2$, $\lambda_{2}=0.2$, $\lambda_{3}=0.3$ and $\lambda_{4}=0.4$ balance the contributions in $\mathcal{L}_{AIRMF}$.

To calculate the affinity loss $\mathcal{L}_{aff}$, we concentrate on adjacent coordinates during training to leverage their sufficient context while minimizing computational costs. We use a search mask on $\hat{\mathbf{M}}$ with a predefined radius $R$ as a threshold, efficiently limiting the distance between selected coordinate pairs. This approach ensures that only pixel pairs within the same local search mask are considered, excluding the affinity between distant pixel pairs. The affinity loss term is formulated as follows:

	$\displaystyle\mathcal{L}_{aff}$	$\displaystyle=-\frac{1}{\lvert\mathcal{S}^{-}\rvert}\sum_{(i,j)\in\mathcal{S}^{-}}\log{(1-\mathbf{Aff}_{pre}^{i,j})}$		(8)
		$\displaystyle-\frac{1}{\lvert\mathcal{S}^{+}\rvert}\sum_{(i,j)\in\mathcal{S}^{+}}\log{(\mathbf{Aff}_{pre}^{i,j})},$

where $\mathcal{S}^{+}$ and $\mathcal{S}^{-}$ refer to two subsets of positive and negative affinity pixel pairs:

$$\mathcal{S}^{+}=\{{(i,j)\mid(i,j)\in\mathcal{P},\mathbf{Aff}_{pre}^{i,j}=1}\},$$

(9)

$$\mathcal{S}^{-}=\{{(i,j)\mid(i,j)\in\mathcal{P},\mathbf{Aff}_{pre}^{i,j}=0}\}.$$

(10)

Our final loss can be written as:

$\displaystyle\mathcal{L}=\frac{1}{2}\mathcal{L}_{AIRMF}+\frac{1}{2}\mathcal{L}_{aff}.$

(11)

Training details. During training, we utilize the Adam optimizer [kingma2017adam] with an initial learning rate $2.25\times 10^{-5}$ to train our network. The learning rate is decayed to one-tenth at iterations 22,500 and 37,500 in a total of 45,000 iterations. We randomly crop input with a resolution of 224x224 from the original images. During the training, the coarse masks are generated by introducing random perturbations to the ground truth masks with a random IoU threshold ranging from 0.8 to 1.0. The radius $R$ of the local window size when computing the affinity loss is set to 14.

We follow the multi-resolution inference strategy proposed in [shen2022high], which can be seen as progressively coarse-to-fine operations. Specifically, the inference begins at a very low resolution around the training image resolution and gradually increases the input’s resolution along the continuous ratio axis $Rs\in(0.125,0.25,0.5,1)$. At each refinement stage, the refined mask is concatenated as a coarse mask for the next refinement step.

Datasets. To acquire object information and evaluate our models, we adhere to a dataset setup similar to that of CascadePSP [cheng2020cascadepsp]. Our training dataset comprises a combination of several individual datasets, totaling 36,572 RGB images accompanied by their corresponding ground truth masks. These images encompass a wide spectrum of over 1,000 categories, drawing from MSRA-10K [6871397], DUT-OMRON [6619251], ECSSD [7182346], and FSS-1000 [Wei2019FSS1000A1]. For testing, CascadePSP introduced a high-resolution image segmentation dataset, named BIG, ranging from 2K to 6K. Additionally, recognizing that the commonly used original PASCAL VOC 2012 dataset lacks pixel-perfect segmentations and designates regions near object boundaries as ’void’, [cheng2020cascadepsp] relabeled 500 objects to create a new dataset suitable for the task of UHR image segmentation.

Metrics. We evaluate our model using Intersection over Union (IoU) and mean Boundary Accuracy (mBA) [cheng2020cascadepsp], which robustly computes the segmentation accuracy across different image sizes within a specified radius from the ground truth boundary. Our model is evaluated without resorting to finetuning in various settings, demonstrating the enhancements it achieves.

Table 1 provides a comprehensive comparison of class-agnostic IoU and mBA among our AIRM, SegFix [yuan2020segfix], CascadePSP [cheng2020cascadepsp] and CRM [shen2022high]. All models are trained on low-resolution images and evaluated on their original high-resolution counterparts. Across both datasets, our method consistently outperforms other methods in terms of IoU and mBA. It’s worth noting that SegFix’s refinement performances are comparatively lower, as it is not specifically designed for the UHR image segmentation task. In the BIG dataset, our method exhibits an average improvement of around 1.7% for IoU and 1.2% for mBA based on four fundamental segmentation approaches. Similarly, in the relabeled PASCAL VOC 2012 dataset, our method achieves an average increase of approximately 1.1% for IoU and 0.6% for mBA compared to CRM and CascadePSP upon four approaches. These quantitative findings underscore the effectiveness of our model in capturing global semantic features and exploring pixel-wise affinity interrelationships in UHR image segmentation.

In Fig. 4, we present a qualitative comparison of the performance of CascadePSP, CRM, and our AIRM. Remarkably, even without any prior exposure to high-resolution training images, our method consistently produces refined results of exceptional quality across various scales. When employing four distinct input coarse masks, it becomes evident that our approach excels at reconstructing missing parts in comparison to CascadePSP and CRM. For instance, as depicted in Fig. 4, our method accurately distinguishes the pilot and the passenger seated in the rear of the aeroplane (the second row) from the segmented object. Additionally, our method precisely delineates the shape of the right rear-view mirror (the third row), achieving a notably high level of precision when compared to the ground truth, which other methods fail to accomplish.

Furthermore, Fig. 5 underscores the substantial visual enhancements achieved by our refinement technique on the relabeled PASCAL VOC 2012 dataset, which consists of low-resolution images. A direct comparison highlights the generation of finer details in our outputs, particularly along the boundary region, such as the gap in the back of the chair and intricate feather textures. More results in supplementary material further demonstrate the efficacy of our refinement algorithm.

Affinity learning. Fig. 6 illustrates that the attention map Fig. 6 (b) can only represent the coarse-level semantic affinity relations, prompting us to investigate the reliable semantic affinities Fig. 6 (c). In Table 2, we showcase the implementation of affinity learning across diverse backbones. For simplicity, We follow the method in [ahn2018learning] to acquire the affinity prediction map in CNN-based backbone. In this approach, the semantic affinity between a pair of feature vectors is defined in terms of their L1 distance. Notably, leveraging affinity from attention within the Transformer encoder yields a performance boost, averaging 1.8%/3.2% higher than its counterpart in the CNN-based encoder. This outcome strongly validates both our initial motivation and the effectiveness of the AEE module.

Adaptive implicit representation mapping function. We employ the TSNE [van2008visualizing] to reduce the dimension of the predicted adaptive parameters and illustrate that in Fig. 7. We observe that for input images with a similar semantic appearance, the corresponding mapping parameters are closely related. This demonstrates that the predicted adaptive parameters possess the capability to capture global semantic information.

Table 3 showcases the effect of the presence of AIRMF. For a fair comparison, we adopt ResNet50 [he2016deep] as an alternative image encoder for extracting image features when AEE is not employed. Additionally, we utilize a naive 5-layer MLP network as the substitute implicit function for AIRMF to decode the extracted latent code, following the setup outlined in [chen2021learning]. As shown in Table 3, we have the following observations: ❶ Without our proposed empowered encoder and adaptive implicit function, the segmentation results decrease by 2.24%/2.43% and 2.33%/2.75%, respectively, compared with a CNN-based encoder and conventional implicit function (see the first-third rows). This highlights our empowered encoder can efficiently capture strong and dependable global semantic information than the CNN-based encoder. Furthermore, our adaptive implicit function effectively utilizes these global features to compute the adaptive parameters, cooperatively enhancing segmentation performance. ❷ Particularly noteworthy is the synergistic performance achieved by the cohesive implementation of both modules in our proposed method, achieving the highest performance at 92.99%/74.77% (see the fourth row). This underscores the essential roles of AEE and AIRMF within AIRM, showcasing their synergistic effects.

Search mask radius in affinity loss. In Table 4, we also study the segmentation performance on different search mask radius $R$ when generating affinity label for computing affinity loss. Through comparison, we observe that a small value of $R$ may not yield sufficient affinity pairs, whereas a large $R$ might compromise the reliability of distant affinity pairs. Thus, we choose $R$ = 14 as our priority.