Learning to Downsample for Segmentation of Ultra-High Resolution Images

1 INTRODUCTION . Many computer vision applications such as auto-piloting , geospatial analysis and medical image processing rely on good semantic segmentation over ultra-high resolution images . Exemplary applications include urban scene analysis with camera array images ( > 25000×14000 pixels ) ( Wang et al. , 2020 ) , geospatial analysis with satellite images ( > 5000×5000 pixels ) ( Maggiori et al. , 2017 ) and histopathological analysis with whole slide images ( > 10,000×10,000 pixels ) ( Srinidhi et al. , 2019 ) . Challenges arise when applying deep learning ( DL ) methods on those ultra-high resolution images , like those in Srinidhi et al . ( 2019 ) ; Chen et al . ( 2019 ) ; Marin et al . ( 2019 ) , due to GPU memory constraints . To speed up performance and adapt to memory requirements or data transmission latency , standard pipeline ( Figure 2 ( a ) ) uniformly downsample both input images and labels , then train the segmentation network at lower resolution , whose predictions are upsampled to the original resolution at testing . However the uniform downsampling wrongly assumes all pixels are equally important thus compromises the performance of applying DL approaches that have proved powerful on standard-sized images without downsampling preprocessing , like those in Chen et al . ( 2016 ) ; Sun et al . ( 2019 ) . To tackle such problem , Marin et al . ( 2019 ) recently postulated that for better segmentation quality more pixels should be picked near semantic boundaries and formulated it as objective to independently train a content-adaptive downsampling network prior to the training of segmentation network , hence we referred to as `` edge-based '' and summarise in Figure 2 ( b ) , which achieved state-of-the-art in low-cost segmentation accuracy . However the `` edge-based '' downsampling is not end-to-end optimised whose sampling density distribution is fixed to empirical single “ best fit ” hyper-parameter . In this work , we start with a motivational experiment to investigate the barriers of current SOTA . We demonstrate that the segmentation performance of the `` edge-based '' method ( Marin et al. , 2019 ) is sub-optimal , as the optimal sampling density distribution varies spatially according to different local patterns . Motivated by this finding , we then introduce deformation module , a lightweight end-to-end trainable `` downsampler '' with segmentation module to adapt sampling density according to the local segmentation difficulty ( Figure 2 ( c ) ) . The downsampled image guided by the deformation module results in a deformed low-resolution image and hence the name . The sampling operation is not directly differentiable , hence we formulate the deformation module to predict sampling density at each location to shift each sampling coordinate in a convolutional form . This idea is inspired by Recasens et al . ( 2018 ) who proposed a jointly learnable downsampler for image classification . However , without further assumptions , jointly learn downsampling to optimise pixel-level segmentation loss is ill-posed because the model can learn to densely sample both input and label at trivial locations like background . We propose to simulate a label-edge based sampling target to regularise the training , inspired by Marin et al . ( 2019 ) . We demonstrate the general utility of our approach using two public and one local datasets from different domains . DeepGlobe ( Demir et al. , 2018 ) and a local Prostate Cancer Histology dataset ( `` PCa-Histo '' ) are datasets consisting of ultra high-resolution aerial and whole slide pathological images . We also use the Cityscapes dataset ( Cordts et al. , 2016 ) since it provides a familiar platform for demonstration and evaluation . In all datasets , we show that a lightweight implementation of the deformation module can more flexibly adjust the sampling to content hence consistently boost the segmentation performance with little extra computational cost . 2 RELATED WORKS . Fix-sized sampling window restricted information `` visible '' to the network . Various multi-scale architectures have shown the importance of sampling multiple-scale information , either at input patches ( He et al. , 2017 ; Jin et al. , 2020 ) or feature level ( Chen et al. , 2014 ; 2016 ; Hariharan et al. , 2015 ) , but less efficient because multiple re-scaling or inference steps . Dilated convolution ( Yu & Koltun , 2015 ) shift sampling locations of convolutional kernel at multiple distances to collect multiscale information therefore avoided inefficient re-scaling . Such shifting is however limited to fixed geometric structures of the kernel . Deformable convolutional networks ( DCNs ) ( Dai et al. , 2017 ) learn sampling offsets to augment each sampling location at feature level . Although DCNs share similarities with our approach , they are complementary to ours because we focus on learning optimal image downsampling as pre-processing so that keeps computation to a minimum at segmentation and a flexibly approach can be plug-in to existing segmentation networks . Learnt sampling methods have been developed for image classification , arguing better image-level prediction can be achieved by an improved sampling while keeps computation cost low . Spatial Transformer Networks ( STN ) ( Jaderberg et al. , 2015 ) introduce a layer that estimates a parametrized affine , projective and splines transformation from an input image to recover data distortions and thereby improve image classification accuracy . Recasens et al . ( 2018 ) proposed to jointly learn a saliency-based network and `` zoom-in '' to salient regions when downsampling an input image for classification . Talebi & Milanfar ( 2021 ) jointly optimise pixel value interpolated ( i.e . super-resolve ) at each fixed downsampling location for classification . However , an end-to-end trained downsampling network has not been proposed so far for per-pixel segmentation . Joint optimising the downsampling network for segmentation is more challenging than the per-image classification , as we experimentally verified , due to potential oversampling at trivial locations and hence we propose to simulate a sampling target to regularise the training . On learning the sampling for image segmentation , post-processing approaches ( Kirillov et al. , 2020 ; Huynh et al. , 2021 ) refining samplings at multi-stage segmentation outputs are complementary to ours , but we focusing on optimising sampling at inputs so that keeps computation to a minimum . Non-uniform grid representation modifying the entire input images to alleviate the memory cost , such as meshes ( Gkioxari et al. , 2019 ) , signed distance functions ( Mescheder et al. , 2019 ) , and octrees ( Tatarchenko et al. , 2017 ) . However none of those approaches is optimised specifically for segmentation . Recently Marin et al . ( 2019 ) proposed to separately train an `` edge-based '' downsampling network , to encourage denser sampling around object edges , hence improving segmentation performance at low computational cost and achieving state-of-the-art . However the `` edge-based '' approach is sub-optimal with sampling density distribution fixed to manual designed objective , the distance to edge , rather than segmentation performance . We verify this point empirically in Section 3.2 and hence propose our jointly trained downsampling for segmentation . 3 METHODS . In this section , we first formulate the basic components associated with this study in Section 3.1 . We then perform a motivation experiment in Section 3.2 to illustrate the impact of manual tuned objective , particularly the sampling density around edge , on segmentation performance and its spatial variation across the image . Motivated by the finding , we propose the deformation module , the jointly trained downsampler for segmentation , in Section 3.3 . To ease the joint optimisation challenge of oversampling at the trivial locations , a regularisation term is proposed in Section 3.4 . 3.1 PROBLEM FORMULATION . Sampling method . Let I ∈ RH×W×C be a high-resolution image of an arbitrary size H , W , C . A sampler G takes I as input and computes a downsampled image Î = G ( I ) , where Î ∈ Rh×w×C . Consider a relative coordinate system1 such that I [ u , v ] is the pixel value of I where u , v∈ [ 0,1 ] . And an absolute coordinate system such that Î [ i , j ] is the pixel value of Î at coordinates ( i , j ) for i∈ { 1,2 , ... h } , j∈ { 1,2 , ... w } . Essentially , the sampler G computes a mapping between ( i , j ) and ( u , v ) . Practically , sampler G contains two functions { g0 , g1 } such that : Î [ i , j ] : =I [ g0 ( i , j ) , g1 ( i , j ) ] 2 ( 1 ) Segmentation and non-uniform upsampling . In this work we discuss model agnostic downsampling and upsampling methods , therefore any existing segmentation model can be applied . Here we denote 1A relative instead of absolute coordinate system is selected for sampling be calculated in a continues space . 2The `` uniform '' approach will have sampler Gu= { g0u ( i , j ) = ( i−1 ) / ( h−1 ) , g1u ( i , j ) = ( j−1 ) / ( w−1 ) } . the segmentation network Sϕ , parameterised by ϕ , that takes as an input a downsampled image Î and makes a prediction Sϕ ( Î ) in the low-resolution space . At testing , the upsampling process consists of a reverse sampler G−1 ( ) that reverse mapping each pixel at coordinates ( i , j ) from the sampled image Î back to coordinates ( u , v ) in the high-resolution domain . Then an interpolation function Π ( ) is applied to calculate the missing pixels to get the final prediction P=Π ( G−1 ( Sϕ ( Î ) ) ) . The nearest neighbour interpolation is used as Π ( ) in this work . Disentangle the intrinsic interpolation error . Typical evaluation of downsampled segmentation with Intersection over Union ( IoU ) is performed after non-uniform upsampling , between the final prediction P and the label Y . Hence IoU ( P , Y ) incorporated both segmentation error and interpolation error . To improve the interpretability of the joint trained system , we propose to disentangle the intrinsic interpolation error from the segmentation error by assuming a perfect segmentation at the downsampled space and calculates the error introduced after interpolation . In specific we calculate IoU ( Y ′ , Y ) , where Y ′ = Π ( G−1 ( G ( Y ) ) ) , indicting the same non-uniform downsampling and upsampling processes applied to the label Y . 3.2 MOTIVATIONAL STUDY : INVESTIGATING THE BARRIER OF THE SOTA . The `` edge-based '' approach ( Marin et al. , 2019 ) , separately train a non-uniform down-sampler Ge minimizing two competing energy terms : `` sampling distance to semantic boundaries '' and `` sampling distance to uniform sampling locations '' , as the first and second term in equation 2 , where b [ i , j ] is the spatial coordinates of the closest pixel on the semantic boundary . A temperature term λ is used to balance the two energies , whose value is empirically recommend to 1 by Marin et al . ( 2019 ) and decided the sampling density around edge . E ( Ge ( i , j ) ) = ∑ i , j ∥Ge ( i , j ) −b [ i , j ] ∥2+λ ∑ i , j ∥Ge ( i , j ) −Ge ( i ’ , j ’ ) ∥2 ( 2 ) The first part of our work performs an empirical analysis to investigate a key question : `` How does the empirically tuned temperature term λ affect the segmentation performance ? '' . We hence perform binary segmentation on the Cityscapes dataset ( Cordts et al. , 2016 ) . We generate a set of simulated `` edge-based '' samplers each at different sampling density around edge ( i.e . λ ) 3 and either directly calculate IoU ( Y ′ , Y ) ( Figure 3 ( a ) ) or train segmentation with each fixed sampler to calculate IoU ( Figure 3 ( b ) ) . It is clear that there is no “ one-size-fits-all ” sampling density configuration that leads to the best performance for all individual classes , neither for intrinsic interpolation error ( IoU ( Y ′ , Y ) ) nor segmentation error ( IoU ) , which is also verified visually in Figure 3 ( c ) . These observations altogether imply that the SOTA `` edge-based '' sampling scheme with a pre-set sampling density around the edge is sub-optimal , highlighting the potential benefits of a more intelligent strategy that can adapting sampling strategy to informative regions according to local patterns .

Summary: This paper aims to improve the performance of segmenting ultra high-resolution images by replacing the commonly used downsampling functions with the learned deformation sampling module. The biggest difference of the proposed method is, instead of sampling all areas uniformly, the sampling density is estimated based on the contents. The experiment results on the Cityscape, DeepGlobe and PCa-Histo datasets show some improvements on mIoU and cost-performance trade-offs compared with uniform downsampling.