Research on land cover type classification method based on improved MaskFormer for remote sensing images

Dataset construction

Dataset production

In this study, the remote sensing images of the Qilian County area were acquired between June and October 2020. A total of 24 scenes are from Gaofen-1, Gaofen-6 and Ziyuan-3 satellites. These original satellite images were from the Natural Resources Remote Sensing Center of Qinghai Province. The original image is a multispectral image with four bands: red, green, blue, and near-infrared. We enhance the vegetation in the image in order to differentiate vegetation from other land covers. Specifically, we remove the near-infrared band, multiply the near-infrared value by 0.2 and add it to the green band. The images with less cloud interference are selected from the original images and then processed by orthorectification, image fusion and image cropping to obtain high-resolution images with a spatial resolution of 2 m in the study area (Li et al., 2022). Qilian County is located at a high altitude and highland hills and gentle slopes dominate the terrain. In this study, several locations were randomly selected in the area containing the above two types of landforms to annotate and create a grassland dataset. The images are annotated as grassland and other categories using ArcGIS tools, and saved as shape-format data. We crop both the label image and the satellite image with base size 256 × 256 pixels, resulting in a total of 820 images. The grassland dataset of Qilian County is randomly divided into training set (80%, 656 images), the validation set (10%, 82 images) and the test set (10%, 82 images). Some of the dataset images and corresponding labels are shown in Fig. 2.

Overall model framework

The MaskFormer network structure (Cheng, Schwing & Kirillov, 2021) treats the semantic segmentation task of pixel-by-pixel classification as a mask classification task, which predicts a set of binary masks, and each mask is associated with a global category label to complete the image segmentation task, and the overall framework is shown in Fig. 3.

The model framework consists of three parts: the pixel-level module, the transformer module, and segmentation module, which complete the feature extraction and feature map upsampling functions, the mask prediction function, and the final segmentation function, respectively. Specifically, the pixel-level module uses an encoder–decoder structure, with the input being an image of C × H × W and the output being ɛ_pixel ∈ ℝ^C_ɛ×N×W; the transformer module uses the standard decoder structure to take as input the feature map output from the encoder in the pixel-level module and the positional embedding of the N learnable positions, and the output is Q ∈ ℝ^C_Q×N; the segmentation module uses Q output from the transformer module to obtain the category probability ${\left\{p_i\in R^{K+1}\right\}}_{i=1}^N$ for each segment using a linear classifier, followed by a softmax activation. And transform Q through the multi-layer perceptron of the two hidden layers into ɛ_mask ∈ ℝ^C_ɛ×N. Then the N predicted binary masks are generated by dot product between the pixel embedding ɛ_pixel and the mask embedding ɛ_mask. Finally, in the segmentation module, we use simple matrix multiplication to get the final prediction matrix K × H × W after removing the empty target category. In this study, we use a more practical feature extraction network and upsampling method based on the MaskFormer network to reduce network complexity and improve the accuracy of image segmentation.

Encoders

We use the Shunted Transformer (Ren et al., 2022) as the feature extraction network, and the network structure is shown in Fig. 4. The whole framework consists of a Patch Embedding module and four cascade modules to produce four resolution outputs. Each module contains a linear embedding and a Shunted Transformer Block module. The Shunted Transformer Block module contains two normalisation layers, Shunted self-attention and Detail Specific FeedForwad. The LN layer normalises the data to make the training process more stable, while the Shunted self-attention layer captures information at different granularities for each attention head, reducing computational effort while fusing multi-scale attention information. Compared with the traditional feedforward layer, the Detail Specific FeedForwad layer adds a detail convolution branch to specify the details in the feedforward layer to supplement feature information. The details are as follows: Given an input of 3 × H × W (3 for RGB channels), the patch embedding module first generates a non-overlapping input sequence of size $\frac{H}{4}\times \frac{W}{4}$ using the convolutional layers of 7 × 7 and 3 × 3.

At this point, the resolution of the original input H × W is reduced to $\frac{H}{4}\times \frac{W}{4}$ and the dimensionality is changed from 3 to 48. The linear embedding layer then maps the dimensionality from 48 to C and feeds the features into the Shunted Transformer module. Each attention head in the module computes attention at different scales and captures information of different granularities for global modelling. In each subsequent iteration module, the resolution of the feature map is reduced to half of the output of the previous module, and the number of channels is doubled.

Decoders

In the decoder part, the original model adopts the feature pyramid network (FPN) (Lin et al., 2017) structure, using bilinear interpolation to upsample the feature map twice, and then fuses the feature map layer by layer from deep to shallow to the original image size. However, the relatively coarse use of nearest neighbor interpolation to upsample feature maps, whose unlearnability may lead to ineffective transfer of high-level feature information. The decoder of the Shunted-MaskFormer network adopts a data-dependent upsampling method. First, the feature maps of different resolutions are uniformly downsampled to the same resolution as the deepest feature map. Second, the feature selection module is used to adaptively recalibrate channel-wise feature response. Finally use the DUpsample (Tian et al., 2019) module to restore the feature map to the original resolution. The entire upsampling process is shown in Fig. 5.

Feature selection module

Compared with the simple use the convolution of 3 × 3 in FPN, before using the feature map extracted by the encoder, it is important to emphasize the feature maps that contain a lot of spatial detail information, while suppressing redundant feature maps. The Squeeze and Excitation (SE) module is proposed by Hu, Shen & Sun (2018). This structure takes into account the different importance of each channel to different categories and obtains the weight vectors of different channels mainly through two operations, squeeze and excitation. The details are as follows, the squeeze operation is used for the down-sampled feature map to extract the global information of different channels using global average pooling, and the excitation operation is used to calculate the dependencies between different channels and output the weight vector. Next, the original feature map is scaled using the weight vectors and added to the original feature map to form a residual structure. The residual structure is used to avoid some channels being over-scaled or suppressed, and the process is defined as: (1)${\displaystyle y=x\ast u+x}$ (2)${\displaystyle u=f_E\left(z\right),}$ (3)${\displaystyle z=f_S\left(x\right)}$

where: x is the feature map input to the feature extraction module; $f_S\left(\cdot \right)$ is the squeeze operation; z is the global information for each channel; $f_E\left(\cdot \right)$ is the excitation operation; u is the calculated weight vector; y is the output of the feature selection module.

DUpsample module

The final layer of the decoder is usually a bilinear upsampling process that restores the feature map to its original resolution. This upsampling method is data-independent and does not consider the correlation between each pixel. Such an upsampling process may lead to suboptimal results. Tian et al. (2019) proposed a data-dependent upsampling method (DUpsample) to replace the bilinear interpolation method. DUpsample exploits the spatial redundancy in segmentation labels to accurately restore the feature map to the original scale, and does not require multiple upsampling strides, thus reducing the framework’s computation time and memory footprint.

In the training process, we no longer use the interpolation method to upsample the feature map, but complete the upsampling process of the feature map F by finding the reconstruction matrix W. The segmentation label Y is not independent and identically distributed, it contains structural information and can be compressed without causing too much loss. In order to minimize the reconstruction error, we use the linear projection method to compress Y to $\tilde{Y}\in {\mathbb{R}}^{\tilde{H}\times \tilde{W}\times \tilde{C}}$ with the same size as F. First, Y isdivided into sub-windows of size r × r(r represents the ratio of the original scale H to the compressed scale $\tilde{H}$), after which each sub-window is deformed into a vector v and multiplied by the matrix P to obtain x. Finally, stack x vertically and horizontally to obtain $\tilde{Y}$. The formula is expressed as: (4)${\displaystyle x=Pv;\tilde{v}=Wx.}$

Here $P\in {\mathbb{R}}^{\tilde{C}\times N}$ is used to compress v to x, W is the reconstruction matrix used to reconstruct x back to v, and $\tilde{\nu }$ represents the reconstructed v.W can be found by minimizing the error between v and $\tilde{\nu }$. Use the reconstruction matrix W to complete the upsampling of F and calculate the error with Y as a loss function. (5)${\displaystyle L\left(F,Y\right)=Loss\left(softmax\left(DUpsample\left(F\right)\right),Y\right).}$

With linear reconstruction, DUpsample (F) applies a linear upsampling to each feature in the tensor F. This upsampling process is essentially the same as applying a 1 × 1 convolution along the spatial dimension, with the convolution kernel stored in W. Decompression is shown in Fig. 6.

Experiments and analysis

Dataset

Two datasets are chosen for the experiment, the Qilian grassland dataset and the Gaofen Image Dataset (GID) (Tong et al., 2020). The construction of the Qilian grassland dataset has been discussed previously. The GID dataset contains 150 images from the Gaofen-2 satellite, each with a size of 7,200 × 6,800 pixels, containing five categories: buildings, farmland, forest, grassland and water. The original data set provides RGB images and near-infrared images. In this study only selected the red–green–blue version of the GID dataset. We choose the typical morphology of each category in the GID dataset as shown in Fig. 7.

This large dataset has an unbalanced number of samples in each category. In order to balance the number of samples in each category and accommodate the limited computational resources. We randomly selected 15 images and cropped them to a size of 256 × 256. Subsequently, 3,000 images were selected from the cropped images to constitute a small GID dataset for the experiment. The small GID dataset is randomly divided into training set (80%, 2,400 images), the validation set (10%, 300 images) and the test set (10%, 300 images). The percentage of data in each sample category is shown in Fig. 8.

Qilian grassland dataset

We compared Shunted-Transformer with several other state-of-the-art computer vision networks on the Qilian grassland dataset for land cover types segmentation. Table 1 shows the results. DeepLab v3+ achieves the best segmentation results among several models based on convolutional neural networks, but there is still a relatively large gap in segmentation accuracy compared to the Vision Transformer network-based models. Shunted-MaskFormer achieves the highest segmentation accuracy. The MIoU and MPA of our method on this dataset are 80.75 and 88.89%. Furthermore, regarding the single-class segmentation results, the grass achieved the best result of 85.74% and the other category also achieved the best segmentation result of 75.76%.

Figure 9 illustrates the visualization of the results of our method compared with the comparison method on the Qilian grassland dataset. From region 1 and region 2 in the figure, it can be seen that in the comparison method of a convolutional neural network, due to the inevitable spatial smooth processing of convolution kernel in the process of convolution, some small areas of other categories and grassland cannot be clearly distinguished and the edge segmentation is not satisfactory enough. Through the attention mechanism, the two segmentation networks based on Vision Transformer are modeled globally, and the segmentation results are better than convolutional neural networks overall, but Upernet overcomes the disadvantage of too smooth boundary segmentation while also leading to more fragmented meadows misclassified into other categories such as region 1. And because only at a single scale calculating attention, it leads to discontinuous boundary segmentation of grassland formed by some small rivers such as region 2 and region 3. Our method used a multi-scale feature extraction network and a data-dependent upsampling process to obtain the finest boundary segmentation results while ensuring the integrity of objects with different land types.

Model size analysis

Through the accuracy analysis of the above two datasets, we proved the validity and generalization ability of the model. In this section, we discuss the scale of the model. We used the number of model parameters and the number of FlOPs as the indicators of model size and complexity respectively. The image size of the input network is uniformly 256 × 256, and the number of parameters and FlOPs measured for each model are shown in Table 3. SegNet has the lowest number of parameters and FlOPs. Upernet has the largest model size with 104M parameters. Our model not only achieves the highest accuracy but also has a relatively small model scale and small complexity among all the models. Our approach achieves a good balance between segmentation accuracy and model complexity.