Small sample image segmentation method based on lightweight multi-scale feature enhancement network

Abstract

A small sample image segmentation method based on a lightweight multi-scale feature enhancement network provides a new type of small sample image segmentation network, namely the lightweight multi-scale feature enhancement network Lite‑FENet, which extracts multi-scale feature representations with rich contextual information, generates descriptive features for predicted pixels, and can effectively alleviate the deviation between support samples and query samples. In addition, the network can replace the back-end feature interaction module of the small sample image segmentation method to further improve performance. The present invention provides a spatial interaction module SIM, which realizes the transmission of contextual information between different spatial scales over long distances with a lower number of parameters, wherein the top-down cross-scale transmission adaptively provides the required information for low-resolution features and outputs a refined feature map. The method can accurately and efficiently predict and segment specific target categories of the query image, suppress erroneous predictions, and is still robust when deviations occur between support samples and query samples.

Description

Technical Field

[0001] The present invention relates to the fields of deep learning and computer vision, and in particular to a small sample image segmentation method based on a lightweight multi-scale feature enhancement network. Background Art

[0002] Image segmentation is a fundamental task in computer vision. It classifies an image at the pixel level, assigning each pixel to a specific class and thereby dividing the entire image into multiple object regions. Image segmentation has become central to a wide range of real-world applications, such as autonomous driving and medical imaging. Driven by deep learning techniques, some fully supervised image segmentation methods have achieved excellent results. However, meticulous pixel-by-pixel annotation of large numbers of samples is time-consuming and expensive, and in the real world, many class examples are difficult to collect, such as rare tumor images. To address these two issues, few-shot image segmentation, a method that segments objects of unseen classes using only a small number of labeled examples, has recently been proposed and has attracted increasing attention. In few-shot image segmentation, the training and test sets for these tasks are disjoint, necessitating a model with strong generalization performance that can learn class-specific knowledge from a small number of support examples and maintain accuracy when transferring to unseen classes.

[0003] Current methods for small-shot image segmentation are primarily meta-learning methods based on metric learning. These methods split the training and test sets into two sets: a support set and a query set. The support set is the image to be segmented, and the query set is the small number of labeled support examples available (for example, when only one support example is available, this is called a one-shot scenario). Furthermore, most existing models employ a two-branch framework, using a weighted backbone network to extract query and support features separately. Feature processing techniques are then used to extract more information. For example, by designing metric functions to calculate the similarity between query and support features, a rough segmentation result of the query image can be obtained. The query features, support features, and other additional information are then input into a feature decoder for matching and mapping the query features to predict the final segmentation result. To improve model performance, multi-scale feature fusion methods are often used to capture spatial context. However, some existing multi-scale structural interactions are shallow and fail to effectively capture the spatial contextual relationship between query and support features, resulting in suboptimal results and a high computational cost. Summary of the Invention

[0004] To address the above issues, this paper proposes a small-sample image segmentation method based on a lightweight multi-scale feature enhancement network. This method employs a novel feature enhancement network structure for small-sample image segmentation, namely the Lightweight Feature Enhancement Network (Lite-FENet). This network can effectively utilize information from query and support samples to enrich discriminative features after multiple interactions, ensuring the accuracy of the final prediction results. Furthermore, it can replace the feature decoding module in current prototype-based small-sample image segmentation methods to optimize performance. Furthermore, a lightweight and efficient Spatial Interaction Module (SIM) is designed for this network to enable the transfer of contextual information between different spatial scales over long distances.

[0005] The small sample image segmentation method based on a lightweight multi-scale feature enhancement network includes the following steps:

[0006] Step 1: Obtain a dataset for small-sample image segmentation and divide it into a training set and a test set. For each training or test, randomly select sample pairs. Each sample pair consists of a support set and a query set. The query set contains one query sample, and the support set contains K annotated support samples. All samples and mask labels undergo unified data preprocessing.

[0007] Step 2: Use the pre-trained backbone network to extract the features of the support samples and query samples, and use the high-level features of the support samples and the high-level features of the query samples to calculate the cosine similarity to obtain the prior mask of the query sample.

[0008] Step 3: Using the support sample labels and the intermediate support features, the support prototype feature vector is extracted through the mask average pooling operation, and this one-dimensional feature vector is expanded to the same spatial size as the prior mask to obtain

[0009] Step 4: Construct a lightweight multi-scale feature enhancement network Lite-FENet to support the prototype feature P s , intermediate query feature F q and the prior mask It takes as input, outputs feature maps of different spatial resolutions, and performs splicing and information fusion;

[0010] Step 5: Use the training data set constructed in step 1 to train the Lite-FENet network constructed in steps 2 to 4. main and auxiliary loss L auxCalculate the total model loss L and perform backpropagation to optimize the model parameters; the backbone network parameters are not updated during training;

[0011] Step 6: Load the Lite-FENet network model trained in step 5 and evaluate the model on the test data set constructed in step 1.

[0012] Step 7: After completing the evaluation and verification, use the Lite-FENet network model to process the small sample image to be segmented and output the segmentation result.

[0013] The beneficial effects of the present invention are:

[0014] (1) The lightweight multi-scale feature enhancement network Lite-FENet proposed in this paper can extract multi-scale feature representations with rich contextual information, generating descriptive features for the final predicted pixel classification. Compared with other advanced small sample image methods, this paper focuses on the design of the network interaction layer, enabling the supporting features to fully interact with the query features at multiple scales, achieving higher performance with a smaller network size.

[0015] (2) The SIM proposed in this paper benefits from the clever application of depthwise convolution and 1×1 convolution, which realizes the contextual information transmission between different spatial scales over long distances with a lower number of parameters. Moreover, the top-down cross-scale transmission can adaptively provide the required information for low-resolution features and output refined feature maps.

[0016] (3) The small sample image segmentation method based on a lightweight multi-scale feature enhancement network proposed in the present invention can accurately and efficiently predict and segment specific target categories in the query image, suppress erroneous predictions, and remain robust when there are deviations between the support samples and the query samples (in terms of quantity, size, posture, occlusion).

[0017] (4) The lightweight multi-scale feature enhancement network proposed in this invention can be simply replaced by the back-end feature interaction module in some small sample image segmentation methods to optimize their performance. BRIEF DESCRIPTION OF THE DRAWINGS

[0018] Figure 1 Schematic diagram of the structure of the lightweight multi-scale feature enhancement network Lite-FENet in an embodiment of the present invention.

[0019] Figure 2 Schematic diagram of the structure of the spatial interaction module SIM in an embodiment of the present invention.

[0020] Figure 3 This is a visual analysis comparison chart of the segmentation results predicted by the lightweight multi-scale feature enhancement network Lite-FENet in an embodiment of the present invention and other advanced small sample image segmentation methods. DETAILED DESCRIPTION

[0021] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings.

[0022] The present invention discloses a small sample image segmentation method based on a lightweight multi-scale feature enhancement network, comprising the following steps:

[0023] Step 1: Obtain a dataset for small-sample image segmentation and divide it into a training set and a test set (validation set). For each training or test, randomly select pairs of samples. Each pair consists of a support set and a query set. The query set contains one query sample, and the support set contains K annotated support samples. All samples and mask labels undergo uniform data preprocessing. This preprocessing includes random scaling by a factor of [0.9, 1.1], random rotation by [-10°, 10°], and random horizontal flipping.

[0024] Step 2: Select a pre-trained backbone network as a feature extractor, such as ResNet, to extract features of support samples and query samples, and use the high-level features of the support samples (the output of the last layer of the backbone network) and the high-level features of the query samples (the output feature map of the fourth stage of the backbone network) to perform cosine similarity calculation to obtain the prior mask of the query sample.

[0025] Step 3: Using the support sample labels and the intermediate support features (the output feature maps of the second and third layers in the backbone network), the support prototype feature vector is extracted through the mask average pooling operation, and this one-dimensional feature vector is expanded to the same spatial size as the prior mask to obtain Both intermediate support features and intermediate query features are intermediate features (the outputs of the second and third stages of the backbone network), but the images used to extract features are different (query images and support images).

[0026] Step 4: Construct a lightweight multi-scale feature enhancement network Lite-FENet. The lightweight multi-scale feature enhancement network Lite-FENet consists of four stages: Figure 1 As shown, each contains 1, 2, 3, and 4 parallel network branches, each with feature representations of different resolutions.

[0027] The network supports the prototype feature P s , intermediate query feature F q and the prior mask As input, in the first stage, the above three types of information are spliced ​​in the channel dimension. After using 1×1 convolution to fuse the channel information, the channel dimension is reduced to 256, and this branch is used as the highest resolution branch:

[0028]

[0029] Among them, ReLU() represents the ReLU nonlinear activation function, Conv 1×1 () represents a 1×1 convolution operation, and Concat() represents concatenation in the channel dimension. After that, at the beginning of each stage, the lowest resolution feature map of the previous stage is downsampled twice to generate a new lower resolution feature map, which is used as a new branch. In the next stage, the feature map output by the previous stage is used as the new branch. (The multiple relationship between each resolution is and the newly generated lowest resolution feature map The input passes through the spatial interaction module in parallel. The superscript n indicates the nth stage, with n branches in that stage. The subscript {1, 2, ..., n} represents the branch index in stage n, with 1 representing the highest-resolution branch. The spatial interaction module transfers more information from adjacent scales to low-resolution features in a top-down manner, generating refined feature maps at that scale. In the final stage, Lite-FENet outputs feature maps at four different spatial resolutions.

[0030] like Figure 2 As shown, SIM consists of two skip connections. The first skip connection receives the refined feature map output by the adjacent higher-resolution branch in the current stage and the feature map of the current branch. It first resizes the higher-resolution feature map to the same spatial size as the smaller-resolution feature map, then undergoes feature concatenation and a 1×1 convolution before performing element-wise addition with the input lower-resolution feature map. The second skip connection consists of a 3×3 depthwise convolution with a stride of 1, two 1×1 convolutions, and a 3×3 depthwise convolution with a stride of 1. The first 1×1 convolution uses a reduction factor to control the channel reduction factor, and the second 1×1 convolution restores the channel to its original size. Each convolution is followed by a batch normalization layer. The ReLU activation function is used after the first and third batch normalization layers, and the output is finally element-wise added to the input feature map.

[0031] The number of channels and spatial resolution of SIM remain the same before and after input.

[0032] In step 4, the feature maps of four different spatial resolutions output by Lite-FENet are upsampled to the same size as the highest resolution and spliced ​​in the channel dimension:

[0033]

[0034] Then a 1×1 convolution is used to fuse feature information of different scales, reducing the number of channels to 256 for the final classification:

[0035]

[0036] Step 5: Use the training data set constructed in Step 1 to train the Lite-FENet network constructed in Steps 2 to 4. Perform backpropagation to optimize the model parameters based on the calculated total model loss L. The backbone network parameters are not updated during the training process.

[0037] Initialize 5 classifiers composed of convolution, one of which is used for final classification and calculate the main loss L main , the other 4 are used for the corresponding 4 branches, and each branch shares a classifier to calculate the auxiliary loss L aux .

[0038] The binary cross entropy function BCE is used to calculate the loss of the model, where the main loss L main The calculation formula is:

[0039]

[0040] Among them, i, j represent the coordinate location of the pixel in the feature map in space; H, W represent the height and width of the feature map respectively; M q (i,j) represents the true label of the pixel, Indicates the final classification result of the pixel predicted by the model.

[0041] The auxiliary loss is the supervised loss calculated for the refined feature maps generated by each branch in each stage. The auxiliary loss L aux The calculation formula is:

[0042]

[0043] in, It represents the classification loss calculated by the feature map generated by the l-th branch in the n-th stage. Since SIM is added to start feature enhancement in the second stage, it starts from n=2:

[0044]

[0045] The total loss L of the model is:

[0046]

[0047] Step 6: Load the Lite-FENet network model trained in step 5 and evaluate the verification model on the test data set constructed in step 1.

[0048] Step 7. After completing the evaluation and verification, use the Lite-FENet network model to process the small sample image to be segmented and output the segmentation result.

[0049] In order to verify the effectiveness of the method, this example uses the PASCAL-5 dataset i 、COCO-20 i We conducted multiple 4-fold cross-comparison experiments and ablation experiments to compare the performance of this method with some classic or advanced methods, and analyzed the impact of various technical details of this method on the overall effect of the method.

[0050] Regarding the experimental environment: All experiments in this example were conducted on a computer running the Ubuntu 20.04LTS operating system and an NVIDIA GeForce RTX3090 graphics card. The software environment included Python 3.7 and PyTorch 1.10.0.

[0051] Regarding model training: In this example, in PASCAL-5 i and COCO-20 i Lite-FENet was trained on the dataset using a GTX 3090 graphics card. The backbone network was ResNet50, and the parameters were kept independent of the update. In both datasets, the batch size was set to 8, the SGD optimizer was used, the initial learning rate was set to 0.005, and the "poly" strategy was used to decay the learning rate. The exponential factor power was equal to 0.9. PASCAL-5 was used. i The input size of the dataset is set to 473×473 and trained for 200 cycles; while using COCO-20 i The dataset model's input size was set to 641×641, and training was performed for 50 epochs. The reduction factor in SIM was 2. In the K-shot scenario, consistent with other small-shot image segmentation methods, the model trained in 1-shot mode was frozen, and the prototype feature vector was averaged over K samples. To ensure fairness in the comparative experiments, a series of data augmentation preprocessing was performed on the input images (for details, see the relevant papers and open-source code for classic small-shot image segmentation methods such as PFENet).

[0052] Regarding model testing: In this embodiment, when evaluating the model, PASCAL-5 i The dataset randomly selects 1000 sample pairs for COCO-20 iThe dataset contains 20,000 randomly selected sample pairs. To eliminate the influence of sample pair selection on the experimental results, all results are the average of the results generated using five different random seeds. To verify the accuracy and lightweightness of the proposed model, the model's test accuracy was evaluated using the Mean Intersection over Union (MIoU), a popular evaluation metric for semantic segmentation tasks, and the model's parameter count (#Params) was also calculated.

[0053] It is particularly important to note that due to PASCAL-5 i The dataset has 20 foreground categories, COCO-20 i The dataset has a total of 80 foreground categories. All categories of the two datasets are divided into 4 parts on average, and a 4-fold cross-validation is performed: three of the category samples are selected during training, and the remaining unseen class sample during training is used for testing. In this way, a test is performed on each category sample, and finally the MIoU of all categories in the dataset is calculated.

[0054] Experimental results:

[0055] (1) Quantitative analysis

[0056] Table 1 in PASCAL-5 i Experimental results on the dataset

[0057]

[0058] Among them, bold represents the best result of this indicator, and underline represents the second best result of this indicator.

[0059] As shown in Table 1, the PASCAL-5 i The Lite-FENet network model trained on the dataset was subjected to 4-fold cross-validation. Compared with previous small-sample image segmentation models, Lite-FENet has the lowest number of learnable parameters (#Params), only 6.7M. In addition, under the 1-shot setting, Lite-FENet achieved the maximum or second-maximum MIoU in the four different folds. Under the 5-shot setting, the MIoU for all classes was also the second highest, second only to DPNet. In addition, after extracting the support prototype using mask average pooling, this method does not further refine the prototype vector. With its lightweight model volume, it surpasses most methods, which is sufficient to prove the effectiveness of this method.

[0060] Table 2. COCO-20 i Comparative experimental results on the dataset

[0061]

[0062] Among them, bold represents the best result of this indicator, and underline represents the second best result of this indicator.

[0063] As shown in Table 2, the results of the 1-shot and 5-shot scenarios using COCO-20 are shown in Table 2. i The Lite-FENet network model trained on the dataset was subjected to 4-fold cross-validation. In terms of the MIoU metric across 80 categories, the Lite-FENet proposed by this method surpassed existing methods to varying degrees in both 1-shot and 5-shot performance. In particular, it surpassed the baseline model PFENet by 7.1 percentage points and 5.7 percentage points in 1-shot and 5-shot performance, respectively, demonstrating that the network can still improve performance on a wider range of datasets.

[0064] (2) Ablation experiment

[0065] Table 3 in PASCAL-5 i Ablation experiment results on the dataset to explore the impact of the number of stages on the model

[0066]

[0067] Among them, bold represents the optimal result of this indicator.

[0068] As shown in Table 3, in PASCAL-5 i The dataset explored the impact of varying numbers of stages on the performance of our Lite-FENet. The number of stages controls the number of branches used to process features in parallel. Increasing the number of stages generates more new branches, ranging from large to small scales. This enhances the model's ability to capture multi-scale information and results in more robust feature representations, leading to improved performance. However, with five stages, performance degraded due to the low resolution of the branch with the smallest resolution, which prevented it from capturing small objects.

[0069] Table 4 in PASCAL-5 i Exploring the effectiveness of the Spatial Interaction Module (SIM) on the dataset

[0070]

[0071] Among them, bold represents the optimal result of this indicator.

[0072] As shown in Table 4, in PASCAL-5 iThis dataset explores the impact of the Spatial Interaction Module (SIM), included in Lite-FENet, on final performance. After adding the SIM, due to its lightweight nature, repeated additions at different stages did not significantly increase the number of model parameters, yet resulted in a 7 percentage point improvement in average performance. This demonstrates that the SIM maintains long-range perception and adaptability while ensuring lightweightness and efficiency.

[0073] Table 5 in PASCAL-5 i Exploring the generalization performance of Lite-FENet on the dataset

[0074]

[0075] Among them, bold represents the optimal result of this indicator.

[0076] As shown in Table 5, in PASCAL-5 i This dataset validates the generalization performance of Lite-FENet. In addition to the baseline model PFENet, the DCP and SCL small-shot image segmentation models were selected. These advanced methods all borrow ideas from the prototype. By replacing the structure within the model with Lite-FENet as the feature interaction module, the performance of the original model was significantly improved, demonstrating the generalization performance of Lite-FENet and its applicability to other small-shot image segmentation tasks.

[0077] (3) Qualitative analysis

[0078] like Figure 3 As shown in the figure, when there is a deviation between the target class objects of the support samples and the query samples (number, size, posture, occlusion), the Lite-FENet network can suppress the error prediction to the greatest extent and can segment specific objects more accurately than other advanced small sample image segmentation methods.

[0079] The above description is only a preferred embodiment of the present invention, and the protection scope of the present invention is not limited to the above embodiment. Any equivalent modifications or changes made by ordinary technicians in this field based on the contents disclosed in the present invention should be included in the protection scope recorded in the claims.

Claims

1. A small sample image segmentation method based on a lightweight multi-scale feature enhancement network, characterized by: The method comprises the following steps: Step 1: Obtain a dataset for small sample image segmentation and divide it into a training set and a test set. For each training or test, randomly select a sample pair. Each sample pair includes a support set and a query set. The query set contains 1 query sample, and the support set contains K annotated support samples. All samples and mask labels undergo unified data preprocessing; Step 2: Use the pre-trained backbone network to extract the features of the support samples and query samples, and use the high-level features of the support samples and the high-level features of the query samples to calculate the cosine similarity to obtain the prior mask of the query sample. Step 3: Using the support sample labels and the intermediate support features, the support prototype feature vector is extracted through the mask average pooling operation, and this one-dimensional feature vector is expanded to the same spatial size as the prior mask to obtain Step 4: Construct a lightweight multi-scale feature enhancement network Lite-FENet to support the prototype feature P s , intermediate query feature F q and the prior mask It takes as input, outputs feature maps of different spatial resolutions, and performs splicing and information fusion; In step 4, the lightweight multi-scale feature enhancement network Lite-FENet transforms the prototype feature P s , intermediate query feature F q and the prior mask After splicing in the channel dimension and fusing the information of each channel using 1×1 convolution, the channel dimension is reduced and the branch is used as the highest resolution branch: Among them, ReLU() represents the ReLU nonlinear activation function, Conv 1×1 () represents a 1×1 convolution operation, and Concat() represents concatenation in the channel dimension; After that, at the beginning of each stage, the lowest resolution feature map of the previous stage is downsampled by two times to generate a new lower resolution feature map, and a new branch is started with this feature map; In the next stage, the feature map output from the previous stage and the newly generated lowest resolution feature map As input, it passes through the spatial interaction module SIM in parallel; where the superscript n indicates the nth stage, the subscript {1,2,...,n} indicates the branch index in the nth stage, and 1 indicates the highest resolution branch; the spatial interaction module transfers more information to the low-resolution features at adjacent scales in a top-to-bottom manner to generate a refined feature map at that scale; in the final stage, Lite-FENet outputs feature maps of four different spatial resolutions; The spatial interaction module includes two skip connections; in the first skip connection, the refined feature map output by the adjacent higher-resolution branch in the current stage and the feature map input to the current branch are received, the spatial size of the higher-resolution feature map is first adjusted to the same spatial size as the smaller-resolution feature map, and then after feature splicing and a 1×1 convolution, the element-level addition operation is performed with the input lower-resolution feature map; the second skip connection is composed of a 3×3 depth convolution with a stride of 1, two 1×1 convolutions, and a 3×3 depth convolution with a stride of 1, and each convolution is followed by a batch normalization layer. The ReLU activation function is used for activation after the batch normalization layer of the first and third convolutions, and the final output is added to the feature map input by the second skip connection at the element level; Step 5: Use the training data set constructed in step 1 to train the Lite-FENet network constructed in steps 2 to 4. main and auxiliary loss L aux Calculate the total model loss L and perform backpropagation to optimize the model parameters; the backbone network parameters are not updated during training; Step 6: Load the Lite-FENet network model trained in step 5 and evaluate the model on the test data set constructed in step 1. Step 7: After completing the evaluation and verification, use the Lite-FENet network model to process the small sample image to be segmented and output the segmentation result.

2. The small sample image segmentation method based on a lightweight multi-scale feature enhancement network according to claim 1, characterized in that: In the lightweight multi-scale feature enhancement network Lite-FENet, the multiple relationship between the resolutions of the feature map is 3. The small sample image segmentation method based on a lightweight multi-scale feature enhancement network according to claim 1, characterized in that: The number of channels and spatial resolution of the spatial interaction module SIM remain the same before input and after output.

4. The small sample image segmentation method based on a lightweight multi-scale feature enhancement network according to claim 1, characterized in that: In step 4, the feature maps of four different spatial resolutions output by Lite-FENet are upsampled to the same size as the highest resolution and spliced ​​in the channel dimension: Then, 1×1 convolution is used to fuse feature information of different scales, and the number of channels is reduced to 256 for the final classification, which is expressed as follows:

5. The small sample image segmentation method based on a lightweight multi-scale feature enhancement network according to claim 1, characterized in that: In step 6, five classifiers composed of convolutions are initialized, one of which is used for final classification and calculation of the main loss; the other four are used for the corresponding four branches, and each branch shares a classifier to calculate the auxiliary loss.

6. The small sample image segmentation method based on a lightweight multi-scale feature enhancement network according to claim 5, characterized in that: The binary cross entropy function BCE is used to calculate the loss of the model, where the main loss L main The calculation formula is: Among them, i, j represent the spatial coordinates of the pixels in the feature map; H, W represent the height and width of the feature map respectively; M q (i,j) represents the true label of the pixel, Indicates the classification result of the pixel point finally predicted by the model; The auxiliary loss is the supervised loss calculated for the refined feature maps generated by each branch in each stage. The auxiliary loss L aux The calculation formula is: in, Represents the classification loss calculated by the feature map generated by the l-th branch of the n-th stage, starting from n=2: The total loss L of the model is:

Images