An image segmentation training method based on soft label self-distillation

Abstract

The application provides an image segmentation training method, comprising the following steps: A1, sampling a batch of multiple images from a preset image data set; A2, constructing an image pair by using each reference image obtained in the current batch and each image in the candidate set of the reference image, and obtaining multiple image pairs; A3, determining a total loss based on an image-level sub-loss and an image block-level sub-loss according to the feature maps obtained by inputting the image pairs and the multiple images in the current batch into online branches and target branches respectively; A4, fixing the parameters of the target branches, and updating the parameters of the online branches according to the total loss; A5, updating the parameters of the target branches by using the historical parameters of the multiple online branches through an exponential moving average strategy; A6, repeating steps A1-A5 to perform self-supervised training until the online branches are trained to convergence, and then proceeding to step A7; A7, obtaining an image segmentation model; and A8, training the image segmentation model by using a supervised image segmentation training set, and obtaining a trained image segmentation model.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence in computer technology, specifically to the field of image segmentation training based on soft label self-distillation.

[0002] Self-supervised learning and image segmentation Background Technology

[0003] With the rapid expansion of deep learning models and data scale, traditional supervised learning frameworks that rely on a large amount of manual annotation can no longer meet the needs of modern deep learning, especially the pixel-level annotations involved in image segmentation tasks. Therefore, self-supervised learning has become a development trend in deep learning technology in recent years, which means learning the embedding of images or pixels without relying on manual annotation, but only using the invariants hidden in the image as supervision information.

[0004] Mainstream self-supervised learning algorithms can be broadly divided into two schools: contrastive learning and non-contrast learning. Both constrain the embedding of the same image to remain unchanged under different data augmentations. Contrastive learning uses both positive and negative sample pairs to train the model, while non-contrast learning uses only positive sample pairs. The main difference lies in how to prevent trivial solutions (i.e., constant solutions). Contrastive learning explicitly increases the distance between the embeddings of different samples, while non-contrast learning uses asymmetric techniques to prevent mode collapse when narrowing the distance between the embeddings of positive sample pairs.

[0005] While existing methods have achieved significant performance improvements on image-level tasks (such as whole-image classification) and even surpassed supervised learning in some downstream tasks, most methods show limited performance gains on pixel-level tasks (such as dense prediction tasks like image semantic segmentation). The inventors discovered that one reason for this is the over-dispersion phenomenon in self-supervised learning. Existing methods constrain the embeddings of different transformations of the same image patch to be consistent, while the embeddings of different sample pairs of the same class are explicitly or implicitly amplified. These sample pairs are called latent positive sample pairs. The embedding distance of latent positive sample pairs should be reduced so that they can be classified into the same class, but existing methods increase this distance. This over-dispersion phenomenon leads to the over-segmentation of different parts of the same category or the same object into multiple categories in downstream tasks, resulting in performance degradation (such as decreased accuracy in downstream semantic segmentation tasks). Taking "human" as an example, it is generally expected that all parts belonging to a human are classified into the same class, but over-segmentation leads to different parts of a "human" being classified into categories such as "head," "hand," and "foot."

[0006] Therefore, self-supervised training needs to address the problem of over-segmentation in existing technologies. Solving this problem faces three main challenges: First, in the absence of manual annotation, potential positive sample pairs are difficult to locate; second, most image patch pairs are negative examples, meaning the two image patches do not belong to the same class, causing the model's learning to be dominated by negative examples and making it difficult to learn positive example information; third, the number of image patch pairs far exceeds the number of image pairs (an image patch pair refers to two image patches, which can belong to the same image or different images), so directly constraining the similarity of all patch pairs would result in unmanageable computational costs. Summary of the Invention

[0007] Therefore, the purpose of this invention is to overcome the shortcomings of the prior art and provide an image segmentation training method.

[0008] The objective of this invention is achieved through the following technical solution:

[0009] According to a first aspect of the present invention, an image segmentation training method is provided, comprising the steps of:

[0010] A1. Sample multiple images in a batch from a preset image dataset, including: sampling multiple reference images and constructing a candidate set for each reference image by sampling positive and negative sample images through clustering. Positive sample images are selected from multiple images with high similarity to the reference images. A2. Construct an image pair using each reference image obtained in the current batch and each image in its candidate set, resulting in multiple image pairs. A3. Based on the image pairs and the multiple images in the current batch, input the feature maps obtained from the online branch and the target branch respectively, and determine the total loss based on image-level sub-loss and image block-level sub-loss, wherein the image block-level sub-loss is configured to correspond to the feature maps of the reference images and positive sample images. Image segmentation is related to: A4. Fix the target branch parameters and update the parameters of the online branch according to the total loss; A5. Use the parameters of multiple historical online branches to update the parameters of the target branch through an exponential moving average strategy; A6. Repeat steps A1-A5 for self-supervised training until the online branch is trained to convergence, then proceed to step A7; A7. Obtain an image segmentation model, which includes a feature extractor for extracting image features from the input image and a linear module for determining the image segmentation result based on the image features, wherein the feature extractor uses an online branch trained to convergence; A8. Train the image segmentation model using a supervised image segmentation training set to obtain a trained image segmentation model.

[0011] Optionally, the total loss is determined in the following manner:

[0012]

[0013] in, Represents image-level sub-loss, Let α represent the image block-level sub-loss. The weight, β represents The weight.

[0014] Optionally, the candidate set for each benchmark image includes positive sample images and multiple negative sample images, wherein the positive sample images are randomly sampled from multiple images with high similarity to the benchmark image in each batch, wherein the image-level sub-loss is configured to correlate with the difference between the ranking of a first similarity and the ranking of a second similarity for all image pairs in the current batch, wherein the first similarity is the similarity between the feature map extracted from the benchmark image by the online branch and the feature map extracted from another image in the image pair, and the second similarity is the similarity between the feature map extracted from the benchmark image by the target branch and the feature map extracted from another image in the image pair.

[0015] Optionally, the image of the input online branch is first processed by a first enhancement method, while the image of the input target branch is first processed by a second enhancement method. The image-level sub-loss is configured to be determined as follows:

[0016]

[0017] Where x represents the reference image, C x Let x represent the candidate set of the benchmark image. c Let v1(·) represent an image in the candidate set, v2(·) represent the first enhancement method, and v2(·) represent the second enhancement method. Represents the reference image and image x c The first similarity between the outputs obtained from the online branch after processing by the first enhancement method is calculated. Represents the reference image and image x c The second similarity between the outputs obtained from the input target branch after processing by the second enhancement method, where τ1 represents the adjustable hyperparameter corresponding to the image-level sub-loss.

[0018] Optionally, the image block-level sub-loss is configured to correlate with the difference between the ranking of the third similarity and the ranking of the fourth similarity of each pixel in the feature maps of the reference image and its positive sample image, wherein the third similarity is the similarity of corresponding pixels in the feature maps extracted by the online branch from the reference image and its positive sample image, and the fourth similarity is the similarity of corresponding pixels in the feature maps extracted by the target branch from the reference image and its positive sample image.

[0019] Optionally, the image block-level sub-loss is determined in the following manner:

[0020]

[0021] Where x represents the reference image, x + S represents the positive sample image of the reference image. p (v1(x), v1(x) + )) ij Represents the baseline image x and its positive sample image x. + After being processed by the first enhancement method, the third similarity S is input to the pixel in the i-th row and j-th column of the feature map extracted online. z (v2(x),v2(x + )) ij Represents the baseline image x and its positive sample image x. + The fourth similarity is the pixel in the i-th row and j-th column between the feature maps extracted from the target branch after being processed by the second enhancement method, and τ2 represents the adjustable hyperparameter corresponding to the image block-level sub-loss.

[0022] Optional, for the above Determined in the following manner:

[0023]

[0024] Where S represents the set of similarities of related elements determined based on the feature map output by the online branch, s i ,s i ∈S, s i Let s represent the similarity of the i-th element in set S. j Let represent the similarity of the j-th element in set S, and let T represent the set of similarities of related elements determined based on the feature map output by the target branch. The similarities of the elements in T are used as soft labels. i ,t i ∈T,t i Let t represent the similarity of the i-th element in set T. j Let w(t) represent the similarity of the j-th element in set T, τ represent an adjustable hyperparameter, N represent the total number of elements in the set, and w(t) represent the similarity of the j-th element in set T. i ) represents the weighted function of the similarity of the i-th element, σ(·) represents the monotonically increasing function, l t (t i -t j ) represents the soft label difference between elements i and j, l s (s i -s j ;τ) represents the similarity difference between elements i and j; I[t i ≤t j ] represents an indicator function, if t i ≤t j Output 1 if true, otherwise output 0; I[s i ≤s j] represents an indicator function, if s i ≤s j Output 1 if true, otherwise output 0.

[0025] According to a second aspect of the present invention, an image segmentation method is provided, the method comprising: acquiring an image to be segmented; obtaining a trained image segmentation model obtained according to the method of the first aspect; extracting image features of the image to be segmented using a feature extractor in the trained image segmentation model; and determining an image segmentation result based on the image features of the image to be segmented using the linear module.

[0026] According to a third aspect of the present invention, an electronic device is provided, comprising: one or more processors; and a memory for storing executable instructions; wherein the one or more processors are configured to implement the steps of the methods of the first aspect and / or the second aspect by executing the executable instructions. Attached Figure Description

[0027] The embodiments of the present invention will be further described below with reference to the accompanying drawings, wherein:

[0028] Figure 1 This is a flowchart illustrating the image segmentation training method according to an embodiment of the present invention;

[0029] Figure 2 This is another flowchart illustrating the image segmentation training method according to an embodiment of the present invention. Detailed Implementation

[0030] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be understood that the specific embodiments described herein are merely illustrative and are not intended to limit the invention.

[0031] As mentioned in the background section, self-supervised training needs to address the oversegmentation problem in existing techniques. Existing self-supervised learning suffers from oversegmentation, meaning it only constrains the embeddings of different transformations of the same image patch to be consistent, while the embeddings of different sample pairs of the same class are explicitly or implicitly amplified; these sample pairs are called latent positive sample pairs. Oversegmentation leads to the same category or different parts of the same object being over-segmented into multiple categories in downstream tasks, resulting in performance degradation.

[0032] To address this, the present invention employs batch-based, multi-stage self-supervised training, with each self-supervised training session following steps S1-S5. This approach offers several advantages: First, it allows for the sampling of positive and negative sample images as candidate sets for the baseline image through clustering, enabling the localization of potential positive sample pairs without manual annotation. Second, it enhances the online branch's ability to accurately extract image features (i.e., feature maps) by using the total loss determined by image-level and image patch-level sub-losses. Third, given the vast number of image patches, using all image patches for computation would not only hinder model learning but also negative examples (i.e., two image patches not being compatible) would further complicate the learning process. Since the image patch is dominated by images belonging to the same class or with low similarity, it is difficult to learn positive examples (i.e., two image patches belong to the same class or have high similarity). Therefore, this application configures the image patch-level sub-loss to be related to the feature maps corresponding to the benchmark image and the positive sample image. This is equivalent to introducing embedding constraints between image patches of potential positive sample pairs in self-supervised learning, thereby reducing over-segmentation, realizing self-supervised training of dense image features, and improving the performance of downstream image semantic segmentation tasks. The positive sample image is selected from multiple images with high similarity to the benchmark image, which enables the online branch to learn more positive example information. It can also greatly reduce the number of image patches that do not belong to the same class and reduce the amount of computation required for the training process.

[0033] According to one embodiment of the present invention, see Figure 1 This invention provides an image segmentation training method, including steps A1-A8. To better illustrate the solution of this invention, each step is described below in conjunction with specific embodiments.

[0034] Step A1: Sample multiple images in a batch from a preset image dataset, including: sampling multiple reference images and constructing a candidate set of positive and negative sample images for each reference image by clustering. Positive sample images are selected from multiple images with high similarity to the reference images.

[0035] According to one embodiment of the present invention, the preset image dataset can be an existing image dataset, such as the MSCOCO dataset or the ImageNet dataset. Alternatively, the preset image dataset can be a dataset composed of multiple image data collected by the implementer, and the present invention does not impose any limitations on this.

[0036] Additionally, to achieve better results in downstream segmentation tasks, the preset image dataset here can be selected from images of the same type as those in the supervised image segmentation training set in step A8 (but the images collected in the preset image dataset do not require manual labeling). For example, if the supervised image segmentation training set corresponds to a task based on segmenting the various components of the fundus from fundus images, then multiple fundus images (e.g., 100, 5000, or 20000) can be collected to construct the preset image dataset; if the supervised image segmentation training set corresponds to a task based on segmenting organs from CT images, then multiple CT images can be collected to construct the preset image dataset. Of course, collecting images of the same type is not mandatory; implementers can also collect multiple images of different types to construct the preset image dataset.

[0037] According to one embodiment of the present invention, the clustering method can employ existing unsupervised clustering algorithms, such as K-Means clustering, hierarchical clustering, or density-based scan clustering (DBSCAN), or a clustering method defined by the implementer. Based on the clustering method, images in a preset image dataset are clustered to obtain the clustering results.

[0038] According to one embodiment of the present invention, after clustering, multiple reference images can be randomly sampled, or a reference image can be determined according to a preset reference image sampling method. For example, the preset reference image sampling method can be set to select (e.g., randomly) one or more images from each of the clustered categories as reference images.

[0039] According to one embodiment of the present invention, after clustering, positive sample images can be selected from the K-nearest neighbor images after clustering for each reference image based on the clustering results. For example, one or more positive sample images can be randomly selected from the K-nearest neighbor images. The K-nearest neighbor images are multiple images with high similarity to the reference image. The K-nearest neighbor images of a reference image refer to the top K images selected from those sorted by spatial distance from the reference image in the clustering results from smallest to largest (the smaller the spatial distance, the higher the similarity). The size of K can be set according to the needs of the implementer, such as 2, 4, 6, or 8, etc., and the present invention does not limit this. Thus, through the preceding steps, potential positive sample pairs (i.e., the reference image and its positive sample images) can be roughly located through image-level clustering and similarity matching, thereby realizing the search for potential similar image block pairs. Subsequently, image block-level loss calculations can be performed based on the reference image and its positive sample images, which can improve online branching performance and significantly reduce the amount of computation.

[0040] According to one embodiment of the present invention, after clustering, one or more negative sample images can be selected for each reference image based on the clustering results. The negative sample images can be multiple images selected from those whose spatial distance from the reference image in the clustering results is in descending order, based on the clustering results. Alternatively, the negative sample images can also be images randomly selected from other images in a preset image dataset, excluding the K-nearest neighbor images of the reference image.

[0041] For simplicity, the following explanation will use the example of sampling one positive sample image and multiple negative sample images for each baseline image x, i.e., the candidate set is represented as follows: Where, x + This represents the positive sample image of the baseline image x. Let n represent the j-th negative sample image of the reference image x. - This indicates the number of negative sample images in the candidate set.

[0042] Step A2: Construct an image pair using each baseline image obtained in the current batch and each image in its candidate set, resulting in multiple image pairs.

[0043] According to one embodiment of the present invention, since the images sampled in each batch may be different, the image pairs are also reconstructed for each batch. Each reference image and its corresponding positive sample image constitute an image pair, and each reference image and its corresponding negative sample image constitute an image pair. Thus, multiple image pairs can be obtained for each batch.

[0044] Step A3: Based on the feature maps obtained from the online branch and the target branch, respectively inputting the image pair and multiple images in the current batch, determine the total loss based on image-level sub-loss and image block-level sub-loss, wherein the image block-level sub-loss is configured to be related to the feature maps corresponding to the reference image and the positive sample image.

[0045] According to one embodiment of the present invention, the online branch and the target branch are feature extractors built based on neural networks. Both the online branch and the target branch are used to extract feature maps from the input image. Preferably, the online branch and the target branch are feature extractors built based on deep neural networks. The online branch and the target branch can use feature extractors from existing models, such as both using feature extractors from the ViT model or the MobileViT model. Alternatively, the online branch and the target branch can adopt a custom structure defined by the implementer, which is not limited by the present invention.

[0046] While it's possible to directly input uncluttered images into both branches, further improvements can be made to enhance robustness. According to one embodiment of the invention, the image input to the online branch is first processed using a first enhancement method, while the image input to the target branch is first processed using a second enhancement method. Illustratively, for a given dataset containing n images... x i H represents the i-th image. I W represents the height of the image. I Indicates the width of the image, This indicates that the image has 3 channels; two viewpoints are generated for each image using two different data augmentation methods (or augmentation methods): v1(x) i ),v2(x i Images processed by different enhancement methods are input into two branches with the same structure but independent parameters, namely the online branch f. p and target branch f z The two branches map the input image to a feature map of scale H×W×D, where H is the height of the feature map, W is the width of the feature map, and D is the feature dimension or the number of channels of the feature map.

[0047] According to one embodiment of the present invention, based on the feature maps and image pairs extracted from the two branches, sub-losses are calculated at both the image level and the image block level, and the total loss is obtained by weighted summation of the two sub-losses. Preferably, the total loss is determined in the following manner:

[0048]

[0049] in, Represents image-level sub-loss, Let α represent the image block-level sub-loss. The weight, β represents The weights. The technical solution of this embodiment can achieve at least the following beneficial technical effects: This embodiment is equivalent to proposing a two-level learning strategy, that is, calculating the sub-loss from the image level and the image patch level respectively, constraining the learning of the feature maps of sample pairs (image pairs or image patch pairs) by the online branch, thereby reducing the amount of computation.

[0050] According to one embodiment of the present invention, the image-level sub-loss is configured to correlate with the difference between the ranking of a first similarity and the ranking of a second similarity among all image pairs in the current batch. The first similarity is the similarity between a feature map extracted from a reference image by an online branch and a feature map extracted from another image in the image pair. The second similarity is the similarity between a feature map extracted from a reference image by a target branch and a feature map extracted from another image in the image pair. Preferably, the image-level sub-loss is configured to be determined in the following manner:

[0051]

[0052] Where x represents the reference image, C x Let x represent the candidate set of the benchmark image. c Let v1(·) represent an image in the candidate set, v2(·) represent the first enhancement method, and v2(·) represent the second enhancement method. Represents the reference image and image x c The first similarity between the outputs obtained from the online branch after processing by the first enhancement method is calculated. Represents the reference image and image x c The second similarity between the outputs obtained from the input target branch after processing by the second enhancement method, where τ1 represents the adjustable hyperparameter corresponding to the image-level sub-loss.

[0053] There are several ways to calculate the similarity between the outputs. For example, using the ViT model, the similarity between images can be calculated in the following ways:

[0054] We can extract the attention map attn(x) output by the ViT model, and first calculate the image-level embedding:

[0055] Image-level embedding of online branches:

[0056] Image-level embedding of the target branch:

[0057] Where H represents the height of the feature map, W represents the width of the feature map, and f p (x) i Let attn(x) represent the representation of image x at position i in the feature map extracted through online branching (i.e., the data at position i in the feature map). i This represents the weight at position i in the attention map, that is, the weight given in the attention map for position i in the feature map.

[0058] Subsequently, for each image pair u,v within the batch, the similarity of the image pair is calculated based on the image-level embedding in the following manner:

[0059] For online branches: the first similarity of image pairs can be calculated as follows:

[0060]

[0061] For the target branch: the second similarity of the image pair can be calculated as follows:

[0062]

[0063] It should be understood that the above similarity calculation formula is for illustrative purposes only, and implementers can adjust it to obtain different implementation methods. For example, [the formula can be modified to...]. and Change the 2 in the formula to 4, or change the 0.5 in the two formulas to 0.501, etc.

[0064] According to one embodiment of the present invention, the image block-level sub-loss is configured to correlate with the difference between the ranking of the third similarity and the ranking of the fourth similarity corresponding to each pixel between the feature maps of the reference image and its positive sample image, wherein the third similarity is the similarity of corresponding pixels in the feature maps extracted by the online branch from the reference image and its positive sample image, and the fourth similarity is the similarity of corresponding pixels in the feature maps extracted by the target branch from the reference image and its positive sample image. Preferably, the image block-level sub-loss is determined in the following manner:

[0065]

[0066] Where x represents the reference image, x + S represents the positive sample image of the reference image. p (v1(x),v1(x + )) ij Represents the baseline image x and its positive sample image x. + After being processed by the first enhancement method, the third similarity S is input to the pixel in the i-th row and j-th column of the feature map extracted online. z (v2(x),v2(x + )) ij Represents the baseline image x and its positive sample image x. + The fourth similarity is the pixel in the i-th row and j-th column between the feature maps extracted from the target branch after being processed by the second enhancement method, and τ2 represents the adjustable hyperparameter corresponding to the image block-level sub-loss.

[0067] Illustratively, the similarity between outputs can be calculated in various ways. For example, for any image pair u, v, the similarity map S corresponding to their feature maps can be calculated. p (u,c),A z (u,c)∈R HW×HW The value in the i-th row and j-th column is determined as follows:

[0068]

[0069]

[0070] Among them, f p (x) i f represents the representation of position i in the feature map of image x extracted through online branching.z (x) i Let represent the representation of position i in the feature map of image x extracted through the target branch, and T denote the transpose. That is, for any image pair u, v, f p (u) i f represents the representation of position i in the feature map of image u extracted through online branching. p (v) j f represents the representation of position i in the feature map of image v extracted through online branching. z (u) i f represents the representation of position i in the feature map of image u extracted through the target branch. z (v) j This represents the representation of position i in the feature map of image v extracted through the target branch.

[0071] It should be understood that the above similarity calculation formula is for illustrative purposes only, and implementers can adjust it to obtain different implementation methods. For example, [the formula can be modified to...]. Know Change the 2 in the formula to 4, or change the 0.5 in the two formulas to 0.51, etc.

[0072] Optional, the above Determined in the following manner:

[0073]

[0074] Where S represents the set of similarities of related elements determined based on the feature map output by the online branch, s i s i ∈S, s i Let s represent the similarity of the i-th element in set S. j Let represent the similarity of the j-th element in set S, and let T represent the set of similarities of related elements determined based on the feature map output by the target branch. The similarities of the elements in T are used as soft labels. i , t i ∈T,t i Let t represent the similarity of the i-th element in set T. j Let w(t) represent the similarity of the j-th element in set T, τ represent an adjustable hyperparameter, N represent the total number of elements in the set, and w(t) represent the similarity of the j-th element in set T. i ) represents the weighted function of the similarity of the i-th element, σ(·) represents the monotonically increasing function, l t (t i -t j ) represents the soft label difference between elements i and j, l s (s i -s j ;τ) represents the similarity difference between elements i and j; I[ti ≤t j ] is an indicator function, if t i ≤t j Output 1 if true, otherwise output 0; I[s i ≤s j ] is an indicator function, if s i ≤s j Output 1 if true, otherwise output 0. This is illustrative; w(t) i )=[t i -β] + / (1-β), [a] + =max(0, a), where β is an adjustable hyperparameter. σ(·) is a monotonically increasing function, for example, σ(z) = z / (1+z). In addition, other monotonically increasing functions can be used, such as σ(z) = z / (2+z), to form other implementation methods.

[0075] In addition, the above The formula is merely an illustrative alternative, and those skilled in the art can adjust it to obtain other alternative implementations, such as:

[0076]

[0077] Where ρ represents the preset hyperparameter for adjustment.

[0078] The above In the middle, l t (t i -t j )=[t i -t j ] + For I[t i >t j A relaxation condition used to avoid similar soft labels (i.e., t) i ≈t j The losses between ) and ). t (t i -t j It can also be set as θ[t] i -t j ] + θ is a hyperparameter used for scaling, which leads to other implementation methods.

[0079] Furthermore, given the hyperparameter τ, l s The illustrative calculation method is as follows:

[0080]

[0081] Among them, l s In (x;τ), x is equivalent to l s(a i -a j ;τ) in a i -s j Additionally, l s (x; τ) can also be defined by the implementer, resulting in other implementation methods, such as modification as follows:

[0082]

[0083] It should be understood that, regarding the calculation process of image-level sub-loss, middle, τ = τ1; and one element corresponds to one image pair. For the calculation of image block-level sub-loss, middle, τ = τ²; and one element corresponds to one pixel pair in a feature map. The feature maps extracted by the online branch and the target branch are smaller in length and width than the length and width of the input image. Therefore, one pixel in the feature map corresponds to one image patch in the input image, and the corresponding sub-loss is the image patch-level sub-loss.

[0084] The following is a schematic illustration of how the inventors improved and obtained the above loss function during the research process.

[0085] Let the similarity expansion calculated based on the outputs predicted from the online branch and the target branch be denoted as the set: Where N is the number of element pairs. Treating T as a soft label, the loss function should measure the consistency of ranking between S and T. Considering that the soft labels are predominantly negative samples (i.e., samples with low scores), directly applying classification or regression loss would lead to the training process being dominated by negative samples. Furthermore, the inventors noted that T contains noise, and the proportion of noise is negatively correlated with the score; therefore, higher weights should be given to positive samples with higher scores. For this purpose, consider the following AUPRC loss:

[0086]

[0087] The indicator function I[A] = 1 if and only if A is true, otherwise I[A] = 0, and σ(x) = x / (1+x) is a monotonically increasing function. However, since the above loss only applies to binary classification of discrete labels, the inventors proposed the soft-label AUPRC loss, i.e., SLAP loss, which is calculated as follows:

[0088]

[0089] Among them, w(t) i )=[t i -β] + / (1-β) is used to measure the probability that both elements in element i are positive samples, [x]+ =max(0,x), where β is an adjustable hyperparameter. Compared to AUPRC loss, the above loss will sum the terms of negative samples. Replace with relatively negative samples (i.e., those with scores lower than t) i Summation of ) Similarly, the summation term over positive samples in the denominator was replaced. Building upon this, to optimize the loss using gradient descent, the non-differentiable indicator function was further replaced with a differentiable substitution loss:

[0090]

[0091] Among them, l t (t i -t j )=[t i -t j ] + For I[t i >t j A relaxation condition used to avoid similar soft labels (i.e., t) i ≈t j The loss between () can be used to calculate the gradient and backpropagate to update the parameters of the online branch.

[0092] Step A4: Fix the target branch parameters and update the parameters of the online branch according to the total loss.

[0093] According to one embodiment of the present invention, the target branch parameters are fixed, that is, the parameters of the target branch are not updated in step A4. Updating the parameters of the online branch based on the total loss includes: calculating the gradient based on the total loss and updating the parameters of the online branch through gradient backpropagation. For example, updating the parameters of the feature extractor of the ViT model, which serves as the online branch.

[0094] Step A5: Utilize the parameters of multiple historical online branches to update the parameters of the target branch using an exponential moving average strategy.

[0095] According to one embodiment of the present invention, an exponential moving average strategy is used. Updating the target branch using this strategy is equivalent to obtaining the target branch's parameters by weighted summation of the parameters of multiple online branch models. In this weighted summation, the target branch updated in batches further removed from the current batch has a smaller weight. After the first batch of updates is completed, since only the parameters of one online branch from the current batch are available, the parameters of the online branch updated in the current batch can be directly used to update the target branch's parameters; alternatively, the initial parameters of the online branches and the parameters of the online branches updated in the first batch can be used to update the target branch using the exponential moving average strategy. This improves the stability of the target branch.

[0096] Step A6: Repeat steps A1-A5 for self-supervised training until the online branch training converges, then proceed to step A7.

[0097] According to one embodiment of the present invention, determining whether an online branch has converged can be based on a method defined by the implementer. For example, the training is considered to have converged after the parameters of the online branch have been updated a preset number of times, or the training is considered to have converged after the total loss is less than or equal to a preset threshold for a consecutive preset number of batches.

[0098] Step A7: Obtain an image segmentation model, which includes a feature extractor for extracting image features from an input image and a linear module for determining the image segmentation result based on the image features, wherein the feature extractor employs an online branch trained to convergence.

[0099] According to one embodiment of the present invention, a linear module is used to perform feature processing, width and height adjustment, and channel adjustment on image features to obtain an image segmentation result. Optionally, the linear module can be implemented using convolutional and / or deconvolutional layers based on a neural network. The number of channels in the image segmentation result needs to correspond to the number of pixel categories in the pixel-level segmentation task achieved by the image segmentation training set.

[0100] Step A8: Train the image segmentation model using a supervised image segmentation training set to obtain the trained image segmentation model.

[0101] According to one embodiment of the present invention, a supervised image segmentation training set includes multiple training images and pixel-level segmentation labels for each image, whereby the pixel-level segmentation labels indicate the ground truth category of each pixel in the training image. Illustratively, the supervised image segmentation training set may be a subset extracted from datasets such as COCOStuff-27, Pascal VOC, or ADE20k. Alternatively, the supervised image segmentation training set may include training images that are CT images, where the pixel-level segmentation labels indicate the organ category (i.e., the ground truth category) corresponding to each pixel in the CT image. The corresponding task is to train a linear module on the downstream supervised image segmentation training set in step A8, with the parameters of the feature extractor of the image segmentation model fixed, to achieve semantic image segmentation; that is, only the parameters of the linear module are updated during training. Without considering the differences in the feature extractors used in the image segmentation model, the training process in step A8 can be regarded as a general supervised training process for image segmentation models.

[0102] According to one embodiment of the present invention, step A8 includes performing one or more training operations, with the trained image segmentation model obtained after the last training operation:

[0103] The training images in the supervised image segmentation training set are input into the image segmentation model. The image segmentation training set includes a variety of training images and pixel-level segmentation labels, which indicate the ground truth category of each pixel in the training image.

[0104] The feature extractor of the image segmentation model extracts image features from the training image, and the linear module of the image segmentation model determines the image segmentation result of the training image based on the image features.

[0105] Based on the image segmentation results of the training images and the corresponding pixel-level segmentation labels, determine the segmentation loss (which can be determined using the existing cross-entropy loss function);

[0106] The gradient is calculated based on the segmentation loss, and the parameters of the linear module are updated through backpropagation of the gradient.

[0107] In addition, the method of the present invention can also be expressed in different forms, for example, see [link to relevant documentation]. Figure 2 According to an embodiment of the present invention, a self-supervised image segmentation training method is provided, comprising:

[0108] S1. Search for potential similar image patch pairs to locate potentially similar image patch pairs through image-level similarity matching;

[0109] S2. Calculate the loss function at the image level and image patch level respectively to measure the difference between the similarity map of the online branch prediction and the target branch. Fix the target branch and update the online branch parameters through gradient backpropagation.

[0110] S3. Update the target branch using an exponential moving average strategy with multiple online branch model sets;

[0111] S4. Repeat steps S1-S4 until the model converges and proceed to S5.

[0112] S5. Fix the above online branch as the feature extractor, and train a linear model for image semantic segmentation on the downstream supervised dataset to obtain the trained image segmentation model.

[0113] To avoid going into too much detail, Figure 2 The corresponding embodiments are described in the prior applications to which the priority is claimed in this application, and are included in this embodiment by reference as if explicitly stated.

[0114] According to one embodiment of the present invention, an image segmentation method is also provided, the method comprising:

[0115] B1. Obtain the image to be segmented;

[0116] B2. Obtain the trained image segmentation model obtained by the image segmentation training method according to the foregoing embodiments;

[0117] B3. The image features of the image to be segmented are extracted using the feature extractor in the trained image segmentation model, and the image segmentation result is determined using the linear module based on the image features of the image to be segmented. It should be noted that in this invention, image features are also known as feature maps; this invention uses different names to distinguish between the feature maps extracted by the image segmentation model and those extracted in the online branch.

[0118] To verify the effectiveness of the invention, the inventors also conducted comparative experiments.

[0119] The inventors fixed the parameters of the trained online branch and added a linear layer (i.e., a linear module). The linear layer was trained on three semantic segmentation datasets: the scene segmentation dataset COCOStuff-27, the object segmentation dataset Pascal VOC, and the fine-grained segmentation dataset ADE20k. The experiments used 384-dimensional features (i.e., the feature map is 384-dimensional; it should be understood that this is only for experimental use, and the feature map can be of other dimensions, such as 256, 480, or 512-dimensional vectors, depending on the implementer's needs). The linear layer was mapped to the number of classes contained in each dataset, and the linear layer was updated using stochastic gradient descent.

[0120] The comparison method uses several state-of-the-art self-supervised training models. Testing metrics include mean intersection-over-union (mIoU) and accuracy (Acc). Experimental results are shown in Table 1, where the best-performing method is highlighted in bold, and the second-best result is underlined. The results demonstrate that the proposed method significantly outperforms the state-of-the-art methods on multiple datasets and metrics, validating the effectiveness of the proposed method.

[0121] Table 1 compares the experimental results on three benchmark datasets for image semantic segmentation.

[0122]

[0123] The following are the literatures for various comparison methods:

[0124] [1]. Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He.

[0125] Improved baselines with momentum contrastive learning.arXiv preprintarXiv:2003.04297,2020.

[0126] [2]. Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems, 33:9912–9924, 2020.

[0127] [3].Xin Wen,Bingchen Zhao,Anlin Zheng,Xiangyu Zhang,and XIAOJUANQI.Self-supervised visual representation learning with semantic grouping.InAdvances in Neural Information Processing Systems.

[0128] [4]. Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In International Conference on Computer Vision, pages 9650–9660, 2021.

[0129] [5].Mark Hamilton,Zhoutong Zhang,Bharath Hariharan,Noah Snavely,andWilliam T Freeman.Unsupervised semantic segmentation by distilling featurecorrespondences.In International Conference on Learning Representations.

[0130] [6].Adrian Ziegler and Yuki M Asano.Self-supervised learning of object parts for semantic segmentation.In IEEE / CVF Conference on ComputerVision and Pattern Recognition,pages 14502–14511,2022.

[0131] In summary, the method of this invention can solve the problem of excessive dispersion in self-supervised image semantic segmentation tasks, that is, the features learned by existing methods for different sample images tend to be unrelated, even if the corresponding sample images belong to the same category. This is due to the self-distillation learning framework proposed in this invention, which uses the similarity obtained from the prediction results of the target branch as a soft label (or pseudo label or soft pseudo label), and proposes a soft label AUPRC loss to improve the robustness of learning, thereby improving the performance of downstream semantic segmentation tasks.

[0132] The self-distillation framework proposed in this invention overcomes challenge one. Specifically, this scheme adopts a dual-branch structure, including an online branch and a target branch, and uses the block pair similarity output by the target branch as a soft pseudo-label for distilling the online branch. To overcome challenge two, this scheme proposes to use ranking loss to avoid decision bias caused by imbalanced data distribution. On the one hand, noting that pseudo-labels contain noise, and the noise ratio decreases as the sample score increases, the Area Under the Precision-Recall Curve (AUPRC) loss is adopted to give greater weight to sample pairs with high scores. Secondly, the original AUPRC loss is only applicable to binary classification labels (i.e., discrete values ​​of 0 or 1), while the pseudo-labels used in this invention belong to soft labels of [0,1]. That is, this invention proposes a soft-label AUPRC (SLAP) loss, where the soft label can be a continuous value between 0 and 1. To overcome challenge three, this scheme proposes a dual-level learning strategy, constraining element pair similarity from the image level and the image block level respectively, thereby reducing the computational cost. By combining the above technologies, this solution enables self-supervised training of dense image features, thereby improving the performance of downstream image semantic segmentation tasks.

[0133] It should be noted that although the steps are described in a specific order above, it does not mean that the steps must be executed in the above specific order. In fact, some of these steps can be executed concurrently, or even in a different order, as long as the required function can be achieved.

[0134] This invention can be a system, method, and / or computer program product. A computer program product may include a computer-readable storage medium having computer-readable program instructions loaded thereon for causing a processor to implement various aspects of the invention.

[0135] Computer-readable storage media can be tangible devices that hold and store instructions for use by an instruction execution device. Computer-readable storage media can be, for example, including but not limited to, electrical storage devices, magnetic storage devices, optical storage devices, electromagnetic storage devices, semiconductor storage devices, or any suitable combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), static random access memory (SRAM), portable compact disc read-only memory (CD-ROM), digital multifunction disc (DVD), memory sticks, floppy disks, mechanical encoding devices, such as punch cards or recessed protrusions storing instructions thereon, and any suitable combination thereof.

[0136] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or technical improvements to the embodiments in the market, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. An image segmentation training method, characterized in that, Including the following steps: A1. Sample multiple images in a batch from a preset image dataset, including: sampling multiple reference images and constructing a candidate set of positive and negative sample images for each reference image by clustering. Positive sample images are selected from multiple images with high similarity to the reference images. A2. Construct an image pair using each baseline image obtained in the current batch and each image in its candidate set to obtain multiple image pairs; A3. Based on the feature maps obtained from the online branch and the target branch, respectively, for the image pairs and multiple images in the current batch, determine the total loss based on image-level sub-loss and image block-level sub-loss. Each reference image's candidate set includes positive sample images and multiple negative sample images. The positive sample images are randomly sampled from multiple images with high similarity to the reference image in each batch. The image-level sub-loss is configured to correlate with the difference between the ranking of the first similarity and the ranking of the second similarity among all image pairs in the current batch. The first similarity is the similarity between the feature maps extracted from the reference image by the online branch and the feature maps extracted from another image in the image pair. The second similarity is the similarity between the feature map extracted from the reference image by the target branch and the feature map extracted from another image in the image pair. The image block-level sub-loss is configured to be related to the feature maps corresponding to the reference image and the positive sample image. The image block-level sub-loss is configured to be related to the difference between the ranking of the third similarity and the ranking of the fourth similarity corresponding to each pixel between the feature maps of the reference image and the positive sample image. The third similarity is the similarity of corresponding pixels in the feature maps extracted by the online branch from the reference image and the positive sample image. The fourth similarity is the similarity of corresponding pixels in the feature maps extracted by the target branch from the reference image and the positive sample image. A4. Fix the target branch parameters and update the parameters of the online branch according to the total loss; A5. Utilize the parameters of multiple historical online branches to update the parameters of the target branch using an exponential moving average strategy; A6. Repeat steps A1-A5 to perform self-supervised training until the online branch training converges, then proceed to step A7. A7. Obtain an image segmentation model, which includes a feature extractor for extracting image features from an input image and a linear module for determining the image segmentation result based on the image features, wherein the feature extractor employs an online branch trained to convergence; A8. Train the image segmentation model using a supervised image segmentation training set to obtain the trained image segmentation model.

2. The method according to claim 1, characterized in that, The total loss is determined in the following manner: in, Represents image-level sub-loss, Represents the image block-level sub-loss. express The weight, express The weight.

3. The method according to claim 2, characterized in that, The image input to the online branch is first processed by a first enhancement method, while the image input to the target branch is first processed by a second enhancement method. The image-level sub-loss is configured to be determined as follows: in, Represents the reference image. This represents the candidate set of the benchmark image. This represents an image from the candidate set. Indicates the first enhancement method. This indicates the second enhancement method. Representing the reference image and the image The first similarity between the outputs obtained from the online branch after processing by the first enhancement method is calculated. Representing the reference image and the image The second similarity is calculated between the outputs obtained from the input target branch after processing with the second enhancement method. This represents the adjustable hyperparameters corresponding to the image-level sub-loss.

4. The method according to claim 2 or 3, characterized in that, The image block-level sub-loss is determined as follows: in, Represents the reference image. The positive sample image representing the reference image. Representing the reference image and its positive sample images The third similarity is obtained by inputting the pixel in the i-th row and j-th column of the feature map extracted online after processing by the first enhancement method. Representing the reference image and its positive sample images The fourth similarity is the pixel corresponding to the i-th row and j-th column of the feature maps extracted from the target branch after processing by the second enhancement method. This represents the adjustable hyperparameters corresponding to the image block-level sub-loss.

5. The method according to claim 4, characterized in that, Regarding the above It shall be determined in the following manner: in, This represents the set of similarities between related elements determined based on the feature maps output by the online branch. , Represents a set The Middle The similarity of elements, Represents a set The Middle The similarity of elements, This represents the set of similarities between relevant elements determined based on the feature map output by the target branch. The similarity of elements in the text is used as a soft tag. , Represents a set The Middle The similarity of elements, Represents a set The Middle The similarity of elements, Indicates adjustable hyperparameters. This represents the total number of elements in the set. Indicates the first A weighted function for the similarity of elements. Represents a function that is monotonically increasing. Represents element and Differences in soft labels between them Represents element and Differences in similarity between them; Indicates an indicator function, if Output 1 if true, otherwise output 0; Indicates an indicator function, if Output 1 if true, otherwise output 0.

6. An image segmentation method, characterized in that, The method includes: Obtain the image to be segmented; A trained image segmentation model is obtained according to the method described in any one of claims 1-5. The feature extractor in the trained image segmentation model is used to extract image features of the image to be segmented, and the linear module is used to determine the image segmentation result based on the image features of the image to be segmented.

7. A computer-readable storage medium, characterized in that, It stores a computer program that can be executed by a processor to implement the steps of the method according to any one of claims 1 to 6.

8. An electronic device, characterized in that, include: One or more processors; as well as Memory, wherein the memory is used to store executable instructions; The one or more processors are configured to implement the steps of the method according to any one of claims 1 to 6 by executing the executable instructions.

Images