An image segmentation method based on Transformer segmentation network and regularized training

Abstract

The present invention provides an image segmentation method based on a Transformer segmentation network and regularized training. The method comprises: performing progressive feature extraction on the original image through a backbone network to generate a feature map; optimizing and reconstructing the feature map through a multi-level feature extraction layer of a pyramid feature extraction structure, using the high-dimensional features output by the pyramid feature extraction structure as input through a Transformer decoder, optimizing and compressing the feature map to obtain a feature vector; performing matrix operations on the feature vector and the feature map output by the pyramid feature extraction structure to obtain a preliminary mask prediction image of the same size as the original image, performing quality optimization training on the preliminary mask prediction image using a regularized training method to obtain a mask image; and performing image segmentation processing on the original image using the mask image. In the present invention, image features are sampled and optimized end-to-end, and the quality and robustness of the image segmentation mask are significantly improved by relying on the regularized training method of the Transformer network.

Description

Technical Field

[0001] The present invention relates to the technical field of image segmentation, and in particular to an image segmentation method based on a Transformer segmentation network and regularization training. Background Art

[0002] Image segmentation technology aims to divide an image into distinct regions based on object categories or instances. It is a pixel-level classification task. Image segmentation has attracted widespread attention due to its applications in fields such as image editing. Traditional image segmentation techniques rely on shallow image features such as color and texture, using pixel clustering to divide images.

[0003] Traditional pixel clustering methods have been eliminated by various deep learning algorithms because they are not applicable to classifying specific instances and categories in images. Some solutions use fully convolutional neural networks (FCN) to apply image classification networks to pixel-level image semantic segmentation tasks. Mathematically, the goal of the semantic segmentation network is to classify each pixel in the image {p ij}∈R H×W×3 , mapped to a specific category through the fully connected layer Here, P ij Represents the specific pixel points in the image, and C represents the number of categories.

[0004] As the more powerful network structure Transformer is widely used in the field of deep learning, some new deep learning segmentation methods have been proposed. Some scholars have proposed a method of compressing image features into feature vectors and unifying the above two methods through matrix operations. This method is different from the pixel-level classification method. Instead, it calculates the similarity between the feature vector and the high-dimensional feature. The final prediction result is expressed as

[0005] The disadvantages of the above-mentioned image segmentation methods in the prior art include: ignoring the differences in feature vectors in different tasks, and having the problem that the Transformer segmentation network structure training is difficult to converge. Summary of the Invention

[0006] The embodiment of the present invention provides an image segmentation method based on a Transformer segmentation network and regularized training to achieve efficient image segmentation processing.

[0007] In order to achieve the above-mentioned purpose, the present invention adopts the following technical solutions.

[0008] An image segmentation method based on a Transformer segmentation network and regularized training, comprising:

[0009] Perform progressive feature extraction on the original image through the backbone network to generate a feature map;

[0010] The feature map is optimized and reconstructed through the multi-level feature extraction layer of the pyramid feature extraction structure, and the high-dimensional features output by the pyramid feature extraction structure are used as input by the Transformer decoder to obtain a set of feature vectors by optimizing and compressing the feature map;

[0011] Performing a matrix operation on the set of feature vectors and a feature map output by the pyramid feature extraction structure to obtain a preliminary mask prediction image having the same size as the original image, and performing quality optimization training on the preliminary mask prediction image using a regularized training method to obtain a mask image;

[0012] The mask image is used to perform image segmentation processing on the original image.

[0013] Preferably, the step of performing progressive feature extraction on the original image through the backbone network to generate a feature map includes:

[0014] The backbone network consists of multiple convolutional layers and downsampling layers. The classification task is pre-trained on the original image through the backbone network, and the original image is progressively extracted to generate a corresponding number of low-dimensional feature maps that meet the given measurement rate.

[0015] Preferably, the step of optimizing and reconstructing the feature map through the multi-level feature extraction layers of the pyramid feature extraction structure includes:

[0016] The low-dimensional feature map output by the backbone network low-dimensional feature extraction module is input into the pyramid high-dimensional feature extraction module. The multi-level feature extraction layer in the pyramid feature extraction structure uses a multi-level CNN or Transformer structure to optimize and reconstruct the low-dimensional feature map, and gradually upsamples the low-resolution features. The pyramid high-dimensional feature extraction module outputs high-resolution high-dimensional features.

[0017] Preferably, the Transformer decoder uses the high-dimensional features output by the pyramid feature extraction structure as input, optimizes and compresses the feature map, and obtains a set of feature vectors, including:

[0018] A set of initialized feature vectors and high-dimensional features output by the pyramid feature extraction structure are input into a Transformer decoder. The initialized feature vectors are updated through a cross-attention mechanism in the Transformer decoder. The high-dimensional features are optimized and compressed through multiple iterative processing processes to obtain a set of feature vectors containing high-dimensional information.

[0019] Preferably, performing a matrix operation on the set of feature vectors and the feature map output by the pyramid feature extraction structure to obtain a preliminary mask prediction image having the same size as the original image includes:

[0020] A matrix dot multiplication operation is performed on the feature vector and the feature map output by the pyramid feature extraction structure to generate a unique mask prediction result and classification probability for each feature vector. The prediction results of all feature vectors are upsampled to finally obtain a preliminary mask prediction image with the same size as the original image.

[0021] Preferably, the method of performing quality optimization training on the preliminary mask prediction image using a regularized training method to obtain a mask image includes:

[0022] The quality of the preliminary mask prediction images is optimized using a regularized training method. By optimizing the matching mechanism during model training, a fixed ratio of preliminary mask prediction images is randomly discarded to ensure similar prediction results in different feature vectors. In the semantic segmentation application scenario, end-to-end training and prediction are performed by adding a new training method to obtain an optimized mask image.

[0023] Preferably, the performing image segmentation processing on the original image using the mask image includes:

[0024] An argmax operation is performed on the optimized mask image to generate a probability map, and each pixel in the probability map is assigned to a specific category to complete the segmentation process of the original image.

[0025] It can be seen from the technical solutions provided by the above-mentioned embodiments of the present invention that the method of the present invention uses the Transformer segmentation network structure to predict image segmentation results, ensuring the versatility and high performance of the model structure in different task scenarios; by adding regularization means during the model training process, the model convergence is optimized, and the prediction performance of the model in different segmentation scenarios is directly improved.

[0026] Additional aspects and advantages of the present invention will be set forth in part in the following description, will become apparent from the following description, or may be learned by practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0027] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0028] Figure 1 A schematic diagram illustrating the implementation principle of an image segmentation method based on a Transformer segmentation network and regularized training provided in an embodiment of the present invention.

[0029] Figure 2 The present invention provides a flowchart of an image segmentation method based on a Transformer segmentation network and regularized training.

[0030] Figure 3 This is a comparison chart of the prediction results when the new regularization technology is used and when not used according to an embodiment of the present invention (the left figure shows the mean of the prediction results of several images without using the regularization method, and the right figure shows the opposite).

[0031] Figure 4 This is a comparison chart of single prediction results using and not using the new regularization technology described in an embodiment of the present invention (the upper chart shows the prediction results without the regularization method, and the lower chart shows the opposite). DETAILED DESCRIPTION

[0032] The embodiments of the present invention are described in detail below, examples of which are shown in the accompanying drawings, wherein the same or similar reference numerals throughout represent the same or similar elements or elements having the same or similar functions. The embodiments described below with reference to the accompanying drawings are exemplary and are only used to explain the present invention, and are not to be construed as limiting the present invention.

[0033] It will be understood by those skilled in the art that, unless expressly stated otherwise, the singular forms "a", "an", "said" and "the" used herein may also include the plural forms. It should be further understood that the term "comprising" used in the description of the present invention refers to the presence of the features, integers, steps, operations, elements and / or components, but does not exclude the presence or addition of one or more other features, integers, steps, operations, elements, components and / or groups thereof. It should be understood that when we refer to an element as being "connected" or "coupled" to another element, it may be directly connected or coupled to the other element, or there may be intermediate elements. In addition, "connected" or "coupled" as used herein may include wireless connections or couplings. The term "and / or" used herein includes any unit and all combinations of one or more associated listed items.

[0034] It will be understood by those skilled in the art that, unless otherwise defined, all terms (including technical and scientific terms) used herein have the same meaning as commonly understood by those skilled in the art in the art to which the present invention pertains. It should also be understood that terms such as those defined in common dictionaries should be understood to have meanings consistent with their meanings in the context of the prior art and, unless defined as such herein, will not be interpreted in an idealized or overly formal sense.

[0035] To facilitate understanding of the embodiments of the present invention, several specific embodiments will be further explained below with reference to the accompanying drawings, and each embodiment does not constitute a limitation on the embodiments of the present invention.

[0036] The implementation principle diagram of an image segmentation method based on a Transformer segmentation network and regularized training provided by an embodiment of the present invention is as follows: Figure 1 The specific processing flow is as shown in Figure 2 As shown, the following processing steps are included:

[0037] Step S1: progressively extract features from the original image through a backbone network consisting of multiple convolutional layers and downsampling layers to generate a corresponding number of feature maps that meet a given measurement rate.

[0038] Step S2: The feature map is optimized and reconstructed through the multi-level feature extraction layer of the pyramid feature extraction structure, and the high-dimensional features output by the pyramid feature extraction structure are used as input through the Transformer decoder to optimize and compress the feature map to obtain a set of feature vectors.

[0039] Step S3: Perform a matrix dot multiplication between the feature vectors and the feature map output by the pyramid feature extraction structure to generate a unique mask prediction result and classification probability for each feature vector. Upsample the prediction results of all feature vectors to obtain a preliminary mask prediction image that is the same size as the original image.

[0040] Regularized training is used to optimize the quality of preliminary mask prediction images. By optimizing the matching mechanism during model training and randomly discarding a fixed percentage of preliminary mask prediction images, we ensure similar prediction results across different feature vectors. In semantic segmentation applications, we incorporate novel training methods for end-to-end training and prediction, directly improving model performance and obtaining optimized mask images.

[0041] In models without regularized training, the model will adopt a best-choice optimization strategy. Specifically, the optimization algorithm will update the model based solely on the best results from all preliminary predictions, then weight these results to produce an optimized mask image. This results in severe biases and limitations in model predictions. The regularization technique employed in this invention constrains the optimization algorithm's selection strategy, expanding the number of samples that influence model updates. This results in more reasonable weights for preliminary predictions and improves the model's generalization performance and final prediction results.

[0042] Step S4: performing an argmax operation on the optimized mask image to generate a probability map, and assigning each pixel in the probability map to a specific category to complete the segmentation process of the original image.

[0043] from Figure 2 As can be seen from the figure, the feature extraction network of the present invention includes: a backbone network low-dimensional feature extraction module, a pyramid-style high-dimensional feature extraction module, and a Transformer decoder module. Below we will introduce their detailed structures in order.

[0044] The backbone network low-dimensional feature extraction module is used to extract low-dimensional feature information from raw RGB images. Common backbone network models include the ResNet series and the Swin Transformer series, which use the CNN network structure or the Transformer structure respectively to perform pre-training on image classification tasks. This is a common network structure in deep learning.

[0045] The pyramid-type high-dimensional feature extraction module is used to take the low-dimensional feature information output by the backbone network low-dimensional feature extraction module as input, reconstruct the low-dimensional feature information using a multi-level CNN or Transformer structure, and gradually upsample the low-resolution features to finally output high-resolution high-dimensional features.

[0046] The Transformer decoder takes a set of initialized feature vectors and high-dimensional features as input and updates the initialized feature vectors using the Transformer's unique cross-attention mechanism. After multiple iterations, the high-dimensional features are optimized and compressed, ultimately outputting a set of feature vectors containing high-dimensional information.

[0047] The Hungarian algorithm is used to match the multiple predicted results and the multiple true results generated by the method of the present invention. The matching results are used to directly calculate the classification loss and mask loss. The classification loss uses the cross entropy loss, and the mask loss uses the Dice loss and cross entropy loss.

[0048] Algorithm steps:

[0049] 1. Input: Input image X

[0050] 2. The image segmentation model of the present invention predicts the output Y = {y0, ..., y 99}

[0051] 3. Randomly discard Y according to the ratio r:

[0052] 4. Match the predicted results with the actual results, calculate the loss function and update the model parameters

[0053] Comparative test

[0054] (1) Training and testing process

[0055] Experiments were conducted in the PyTorch framework, using various datasets including ADE20K, Cityscapes, COCOStuff, and VSPW. To promote network convergence, we cropped the training sets and augmented the original images with methods such as flipping and color transformation. A learning rate of 0.0001 was used for training. A batch size of 16 was used for all training data. During testing, we used validation sets from the corresponding datasets to verify the results. No additional processing was required, and the model directly generated predictions end-to-end.

[0056] (2) Comparison of experimental results

[0057] First, we compare our method with the state-of-the-art image segmentation methods, namely MaskFormer and Mask2Former. In Table 1, we give the objective performance comparison of the mean IoU value, where the best results are highlighted in bold.

[0058] Table 1 Comparison of average IoU performance of different datasets

[0059]

[0060] As shown in Table 1, the Transformer model with DQ regularization achieves the highest mean Intersection over Union (MIOU) values ​​across all four datasets and two different Transformer models. We can see that the new regularization method achieves significantly better results. On the ADE20K and Cityscapes datasets, the performance improvement is greater than 0.5% mIoU. On the COCOStuff and VSPW datasets, the improvement exceeds 1% mIoU. Furthermore, the regularization method can be easily applied to several Transformer segmentation models without requiring any adjustments to the original network structure.

[0061] Table 2 shows the basis for selecting the hyperparameters of the regularization method. The Mask2Former model is used as the base model, and the selection is based on the average IoU value of the experimental results.

[0062] Table 2 Comparison of average IoU under different measurement random dropout ratios

[0063]

[0064] like Figure 3 As shown in , we show the average distribution of the prediction results of the existing model's 100 feature vectors on 5000 pictures, and the left and right represent the prediction results before and after the regularization operation. Figure 3 It can be observed that the regularization technology significantly improves the proportion of effectively activated feature vectors to all feature vectors, thereby enhancing the segmentation performance and robustness of the model.

[0065] like Figure 4 As shown in , we show the best 4 prediction results and their classification scores in the single image prediction scenario. The upper and lower parts represent the prediction results before and after the regularization operation. Figure 4 It can be seen that the use of regularization increases the number of valid prediction results, thereby reducing the possibility of the model making mistakes.

[0066] In summary, the method of the embodiment of the present invention uses the Transformer network structure to predict image segmentation results, ensuring the versatility and high performance of the model structure in different task scenarios; by adding regularization means during the model training process, the model convergence is optimized, which directly improves the prediction performance of the model in different segmentation scenarios.

[0067] In this paper, image features are sampled and optimized end-to-end, relying on the regularization training technology of the Transformer network architecture, so that the quality and robustness of the image segmentation mask are significantly improved.

[0068] Those skilled in the art will appreciate that the accompanying drawings are merely schematic diagrams of an embodiment, and the modules or processes in the accompanying drawings are not necessarily required to implement the present invention.

[0069] From the above description of the embodiments, it can be seen that those skilled in the art can clearly understand that the present invention can be implemented by means of software plus the necessary general-purpose hardware platform. Based on this understanding, the technical solution of the present invention, or the portion that contributes to the prior art, can be embodied in the form of a software product. This computer software product can be stored in a storage medium such as ROM / RAM, a magnetic disk, or an optical disk, and includes a number of instructions for enabling a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in various embodiments of the present invention or certain parts of the embodiments.

[0070] Each embodiment in this specification is described in a progressive manner. The same or similar parts between the embodiments can be referred to each other. Each embodiment focuses on the differences from other embodiments. In particular, for the device or system embodiments, since they are basically similar to the method embodiments, the description is relatively simple. For the relevant parts, refer to the partial description of the method embodiments. The device and system embodiments described above are merely schematic, wherein the units described as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units, that is, they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules can be selected according to actual needs to achieve the purpose of the scheme of this embodiment. A person of ordinary skill in the art can understand and implement it without making any creative efforts.

[0071] The above description is merely a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by a person skilled in the art within the technical scope disclosed in the present invention should be included in the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be based on the scope of protection of the claims.

Claims

1. An image segmentation method based on Transformer segmentation network and regularized training, characterized in that: include: Perform progressive feature extraction on the original image through the backbone network to generate a feature map; The feature map is optimized and reconstructed through the multi-level feature extraction layer of the pyramid feature extraction structure, and the high-dimensional features output by the pyramid feature extraction structure are used as input by the Transformer decoder to obtain a set of feature vectors by optimizing and compressing the feature map; Performing a matrix operation on the set of feature vectors and a feature map output by the pyramid feature extraction structure to obtain a preliminary mask prediction image having the same size as the original image, and performing quality optimization training on the preliminary mask prediction image using a regularized training method to obtain a mask image; Performing image segmentation processing on the original image using the mask image; The optimizing and reconstructing of the feature map through the multi-level feature extraction layers of the pyramid feature extraction structure includes: The low-dimensional feature map output by the backbone network low-dimensional feature extraction module is input into the pyramid high-dimensional feature extraction module. The multi-level feature extraction layer in the pyramid feature extraction structure uses a multi-level CNN or Transformer structure to optimize and reconstruct the low-dimensional feature map, and gradually upsamples the low-resolution features. The pyramid high-dimensional feature extraction module outputs high-resolution high-dimensional features; The Transformer decoder uses the high-dimensional features output by the pyramid feature extraction structure as input, optimizes and compresses the feature map, and obtains a set of feature vectors, including: A set of initialized feature vectors and high-dimensional features output by the pyramid feature extraction structure are input into a Transformer decoder. The initialized feature vectors are updated through a cross-attention mechanism in the Transformer decoder. The high-dimensional features are optimized and compressed through multiple iterative processing processes to obtain a set of feature vectors containing high-dimensional information.

2. The method according to claim 1, characterized in that The method of performing progressive feature extraction on the original image through the backbone network to generate a feature map includes: The backbone network consists of multiple convolutional layers and downsampling layers. The classification task is pre-trained on the original image through the backbone network, and the original image is progressively extracted to generate a corresponding number of low-dimensional feature maps that meet the given measurement rate.

3. The method according to claim 1, characterized in that The step of performing a matrix operation on the set of feature vectors and the feature map output by the pyramid feature extraction structure to obtain a preliminary mask prediction image having the same size as the original image includes: A matrix dot multiplication operation is performed on the feature vector and the feature map output by the pyramid feature extraction structure to generate a unique mask prediction result and classification probability for each feature vector. The prediction results of all feature vectors are upsampled to finally obtain a preliminary mask prediction image with the same size as the original image.

4. The method according to claim 3, characterized in that The method of using the regularized training method to perform quality optimization training on the preliminary mask prediction image to obtain the mask image includes: The quality of the preliminary mask prediction images is optimized using a regularized training method. By optimizing the matching mechanism during model training, a fixed ratio of preliminary mask prediction images is randomly discarded to ensure similar prediction results in different feature vectors. In the semantic segmentation application scenario, end-to-end training and prediction are performed by adding a new training method to obtain an optimized mask image.

5. The method according to claim 4, characterized in that The performing image segmentation processing on the original image by using the mask image includes: An argmax operation is performed on the optimized mask image to generate a probability map, and each pixel in the probability map is assigned to a specific category to complete the segmentation process of the original image.

Images