Image segmentation loss function design method for guiding attention by label

Abstract

The invention relates to an attention-guided image segmentation loss function design method, belongs to the field of deep learning image processing, and aims to solve the problem that an attention mechanism of a deep learning image segmentation model is difficult to effectively reinforce different types of features. Under the condition that an image segmentation task is given, a calculation result of an attention module in a deep learning image segmentation model is directly incorporated into calculation of a model loss function, the attention module is directly associated with a training label, and the attention module is guided through the label, so that the attention module effectively focuses on a category feature, and the image segmentation efficiency is improved. The accuracy of the model is improved; and meanwhile, a gradient updating chain of the attention module is shortened, so that the convergence speed of model training is accelerated.

Description

Technical Field

[0001] The present invention relates to a method for designing a loss function for image segmentation based on deep learning, and in particular to a method for designing a loss function for image segmentation guided by attention, and belongs to the field of deep learning image processing. Background Art

[0002] Image segmentation is one of the basic tasks in the field of image processing. Its purpose is to divide an image into several meaningful regions for better subsequent image analysis and understanding. This technology has been widely used in computer vision, medical image analysis, autonomous driving and other fields.

[0003] Traditional image segmentation methods include threshold segmentation and region growing algorithms. These methods are simple to implement, but are sensitive to noise and have poor robustness. They are difficult to achieve ideal results in multi-category complex image segmentation tasks. Therefore, many scholars have proposed image segmentation methods based on deep learning. These methods use deep learning models to extract high-dimensional abstract features of images, making the accuracy of image segmentation far surpass traditional methods.

[0004] The attention mechanism is a deep learning feature extraction technology commonly used in the fields of natural language processing and computer vision. Its core idea is to give different attention weights to different input data in the process of deep learning, so as to strengthen relevant features and weaken redundant features. This mechanism has been introduced by many scholars into multi-category image segmentation tasks. For example, Oktay et al. (Ozan Oktay, Jo Schlemper, Loic Le Folgoc et al. Attention U-Net: Learning Where to Look for the Pancreas[J].arXiv preprint,arXiv:1804.03999,2018.) used multiple attention modules to strengthen features of different categories in pancreatic image segmentation. At present, most image segmentation models based on deep learning include attention modules for feature extraction. However, in practical applications, it is sometimes found that not only can each attention module not extract features as expected, but also multiple attention modules have similar weights, which makes it difficult to effectively strengthen features of different categories and fail to achieve the purpose of distinguishing different features. Therefore, Zhang et al. (Zhang Wenyao, Li Xue, Kang Zijian, et al. A brain tissue segmentation method combining UNet and volume rendering prior knowledge [P]. Beijing: CN202210701401.4, 2024-07-19.) proposed a method of strengthening the attention mechanism with prior knowledge, multiplying the prior knowledge weight with the calculation result of the attention module as the weight of the attention module to strengthen the relevant features. This method expands the application of the attention mechanism in multi-category image segmentation, which not only simplifies the learning process of the attention module, but also improves the convergence speed of the entire model training. However, this method has high requirements on image quality, because image noise can seriously affect the accuracy of the prior knowledge weight, which in turn leads to a decrease in the segmentation accuracy of the model. In addition, the requirements for prior knowledge are too strict, because in many cases there is no additional prior knowledge available.

[0005] Therefore, how to apply the attention mechanism to effectively enhance the features of different categories in image segmentation, especially multi-category image segmentation tasks, remains a problem that needs to be further explored.

[0006] At present, the existing deep learning image segmentation models generally embed the attention module as a basic component into the backbone network; then during the model training process, the overall calculation results of the model are combined with the training data labels to design the loss function and calculate the corresponding loss; then the parameters of the attention module are updated through the back-propagation mechanism. In this process, the attention module is no different from other network modules. Although the original intention of the attention module is to strengthen relevant features and weaken irrelevant features, the attention module is not treated differently during the model training process, or the attention module is not effectively guided. This is one of the main reasons why the existing image segmentation models using the attention mechanism have slow convergence speed and low prediction accuracy.

[0007] In view of the above situation, the present invention proposes a method for designing an image segmentation loss function guided by labels. This method directly incorporates the calculation results of the attention module into the calculation of the model loss function, directly associates the attention module with the training labels, and guides the attention module through the labels, so that it can focus on relevant features faster, thereby improving the accuracy and convergence speed of the model. Summary of the invention

[0008] The purpose of the present invention is to solve the problem that the attention mechanism of the deep learning image segmentation model is difficult to effectively strengthen the features of different categories, and provide a method for designing an image segmentation loss function guided by labels. The method directly incorporates the calculation results of the attention module into the calculation of the model loss function, directly associates the attention module with the training label, and guides the attention module through the label. This can not only improve the gradient update efficiency of the attention module to accelerate convergence, but also enable the model to effectively strengthen the features of different categories, thereby improving the accuracy of the model.

[0009] The objectives of the present invention are achieved through the following technical solutions.

[0010] A method for designing an image segmentation loss function guided by label attention, comprising the following steps:

[0011] Step 1: Obtain a multi-category image segmentation dataset, perform necessary preprocessing on the image data to make it meet the neural network processing requirements, and divide it into a training set S and a test set T.

[0012] Step 2: For the multi-category image segmentation task, build a deep learning model M with an attention mechanism, so that the attention sub-modules of each category together constitute a multi-category attention module A.

[0013] Step 3: Let A = [A1, A2, …, A c ] represents the calculation result of the multi-category attention module, Y = [Y1, Y2, ..., Y c] represents the labels corresponding to each category, and the attention loss function of model M is designed according to the following pattern:

[0014] Loss Attention =Loss(f(A)),Y)

[0015] Where Loss(·) represents any valid loss function, and f(·) represents the compression transformation function. The compression transformation function f(·) is implemented using a neural network module, which integrates the calculation results of the attention module into a tensor of the label size in order to calculate the loss relative to the label. The structure and form of the compression transformation function module are not limited as long as it meets its function.

[0016] Step 4: Embed the module that implements the compression transformation function f(·) into the model M, and embed the attention loss function Loss Attention Incorporated into the loss function of model M.

[0017] Step 5: Use the training set S to train the model M.

[0018] Step 6: Use the test set T to test the model M and evaluate the effect of its loss function.

[0019] Beneficial Effects

[0020] The image segmentation loss function design method guided by label attention described in the present invention has the following technical advantages and characteristics compared with the existing deep learning image segmentation technology and method using attention mechanism:

[0021] (1) The method of the present invention directly links the output of the attention module with the loss function. The loss function designed in this way provides a direct and effective guidance for the attention module, which can enable the attention module to effectively focus on category features, thereby improving the model accuracy.

[0022] (2) The method of the present invention directly uses the calculation results of the attention module and the label to calculate the loss, shortening the gradient update chain of the attention module, which can effectively improve the model update efficiency and thus accelerate the convergence speed of model training. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] Figure 1 Flowchart of a method for designing loss functions for image segmentation guided by label attention;

[0024] Figure 2 The network structure diagram of model M, where the gray marked box is the multi-category attention module, the green marked box is the newly embedded neural network module that implements the compression transformation and the corresponding attention loss calculation module;

[0025] Figure 3Comparison of image segmentation results of model M under different conditions, where the first column is the true value of the test sample, the second column is the segmentation result when the model uses attention loss, and the third column is the segmentation result when the model does not use attention loss;

[0026] Figure 4 Comparison of the convergence speed and segmentation accuracy of model M under different conditions, where the blue curve is the case with attention loss and the orange curve is the case without attention loss. Specific implementation methods

[0028] The specific implementation of the present invention is described below in conjunction with the accompanying drawings and examples.

[0029] Figure 1 A flowchart of a method for designing an image segmentation loss function guided by label attention according to the present invention is given, and its main steps are as follows:

[0030] Step 1: Obtain a multi-category image segmentation dataset, perform necessary preprocessing on the image data to make it meet the neural network processing requirements, and divide it into a training set S and a test set T.

[0031] Image segmentation datasets usually contain two parts: (1) the original image data and (2) the label data corresponding to the original image data. The label data marks the category to which each pixel in the original image data belongs. Preprocessing mainly involves standardizing and cropping the image data to meet the input and output requirements of the neural network.

[0032] The multi-category image segmentation dataset used in this embodiment is ISEG-2017 (L.Wang et al., "Benchmark on Automatic Six-Month-Old Infant Brain Segmentation Algorithms: TheiSeg-2017 Challenge," in IEEE Transactions on Medical Imaging, vol.38, no.9, pp.2219-2230, Sept.2019.). This dataset contains 10 manually annotated infant brain MR image samples. The MR image of each sample has the skull removed and has two modes, T1 and T2. The sample labels contain four categories, namely background, cerebrospinal fluid, gray matter, and white matter.

[0033] In this embodiment, samples 0 to 8 of the 10 samples in the ISEG-2017 dataset are selected as the training set S, the 9th sample is selected as the test set T, and the image data and label data of the two modalities are divided into small blocks of 64×64×64.

[0034] Step 2: For the multi-category image segmentation task, build a deep learning model M with an attention mechanism, so that the attention sub-modules of each category together constitute a multi-category attention module A.

[0035] The network structure of model M depends on the specific image segmentation task and has no special restrictions as long as it includes the attention mechanism. Attention module A is a part of model M and is composed of attention submodules of various categories of model M.

[0036] This embodiment needs to segment the infant brain MR image into background, cerebrospinal fluid, gray matter and white matter, a total of four categories. To this end, this embodiment constructs a deep learning model M based on AttentionUNet. In the model M, a respective attention submodule is set for each category, respectively called A1, A2, A3 and A4. These four submodules together constitute the multi-category attention module A of this embodiment.

[0037] The AttentionUNet involved here is a relatively classic image segmentation model that includes an attention mechanism and was proposed earlier (Ozan Oktay, Jo Schlemper, Loic Le Folgoc et al. Attention U-Net: Learning Where to Look for the Pancreas[J].arXiv preprint,arXiv:1804.03999,2018.).

[0038] This embodiment builds a model M for segmenting infant brain MR images according to the framework of AttentionUNet. Its main structure consists of four pairs of encoders and decoders, such as Figure 2 As shown in the figure, each encoder contains a convolution module and a downsampling module, and each decoder contains a convolution module and a deconvolution module. The convolution module is composed of a convolution layer, a normalization layer, an activation layer, a Dropout layer, a convolution layer, a normalization layer, and an activation layer in sequence, where the normalization layer uses the BatchNormalization method and the activation layer uses the ReLU activation function. The downsampling module uses the MaxPooling method. There are three jump connections between the encoder and the decoder. Each jump connection contains an attention module, whose input is the output of the encoder in the same layer and the output of the decoder in the next layer. To adapt to the training set S, the network input size of model M is set to 2×64×64×64, and the corresponding label size is 4×64×64×64.

[0039] Step 3: Let A = [A1, A2, ..., A c ] represents the calculation result of the multi-category attention module, Y = [Y1, Y2, ..., Y c] represents the labels corresponding to each category, and the attention loss function of model M is designed according to the following pattern:

[0040] Loss Attention =Loss(f(A)),Y), (1)

[0041] Where Loss(·) represents any valid loss function, and f(·) represents the compression transformation function. The compression transformation function f(·) is implemented using a neural network module, which integrates the calculation results of the attention module into a tensor of the label size in order to calculate the loss relative to the label. The structure and form of the compression transformation function module are not limited as long as it meets its function.

[0042] In this step, the attention loss function Loss(·) can be selected from existing common loss functions or set according to actual needs.

[0043] The attention loss function Loss used in this embodiment Attention As shown in formula (2):

[0044] Loss Attention (x′,y)= Loss CE (x′,y)+Loss Dice (x′,y), (2)

[0045] Among them, x′=f(A)(A=[A1,A2,…,A4]) represents the output of each attention submodule after compression transformation and integration of the calculation results of each category; y=[Y1,Y2,…,Y4] represents the true label value corresponding to each category; Loss CE is the cross entropy loss function, Loss Dice is the Dice loss function, and their definitions are as follows:

[0046]

[0047] Among them, N and C represent the number of samples and the number of categories respectively, n and c represent the sample number and the category number respectively, and x nc is the model output result of a specific sample, y nc is the label value of a specific sample.

[0048] The function of the compression transformation function f(·) is to integrate the calculation results of the attention module into a tensor of the label size in order to calculate the loss relative to the label. In the specific implementation, the compression transformation function is implemented using a neural network module. The module can be designed by itself, and its structure and form are not limited as long as it meets the expected function.

[0049] This embodiment uses a self-designed neural network module FA to implement the function of the compression transformation function. The neural network module FA first converts the four tensors of size 32×64×64×64 output by the attention module into 1×64×64×64 tensors respectively through a convolution layer with a convolution kernel size of 1×1×1; then the four 1×64×64×64 tensors are spliced ​​into a 4×64×64×64 tensor in the order of background, cerebrospinal fluid, gray matter, and white matter. In this way, the output of the multi-category attention module is aligned with the category label after compression transformation processing.

[0050] Step 4: Embed the module that implements the compression transformation function f(·) into the model M, and embed the attention loss function Loss Attention Incorporated into the loss function of model M.

[0051] In this embodiment, the original loss function of model M is the sum of cross entropy loss and Dice loss, that is:

[0052] Loss M (x,y)=Loss CE (x,y)+Loss Dice (x,y). (5)

[0053] In order to incorporate the attention loss function in step 3, this embodiment modifies the loss function of model M as follows:

[0054] Loss M (x,x′,y)=Loss CE (x,y)+Loss Dice (x,y)+Loss Attention (x′,y), (6)

[0055] Among them, x is the result of model prediction, x′ is the result of the output of the attention module after compression transformation and integration, and y is the true label value; Loss CE 、Loss Dice and Loss Attention The definitions of are shown in equations (3), (4) and (2) respectively.

[0056] After determining the final loss function of the model, the model needs to be modified as necessary to embed the neural network module FA that implements the compression transformation function f(·) and update the loss function calculation module of the model M. After the neural grid module FA that implements the compression transformation is embedded in the model M of this embodiment, the final network structure is as follows: Figure 2 As shown, the gray marked box is the multi-category attention module, the green marked box is the newly embedded neural network module that implements the compression transformation and the corresponding attention loss calculation module.

[0057] Step 5: Use the training set S to train the model M.

[0058] Step 6: Use the test set T to test the model M and evaluate the effect of its loss function.

[0059] After fully training the model M, this embodiment uses the test set data to perform an ablation experiment on it in order to evaluate the effect of the newly designed loss function.

[0060] Figure 3 The comparison of image segmentation results of model M with and without the attention loss function is given, where the first column is the true value of the test sample, the second column is the segmentation result when the model uses the attention loss function, and the third column is the segmentation result when the model does not use the attention loss function; the first to third rows are the cross-sectional views of the axial, coronal and sagittal planes of the test sample, respectively. Figure 3 It can be seen that when the attention loss function designed by the method of the present invention is used, the model can accurately segment the gray matter and white matter in the red circle. If it is not used, there will be an erroneous segmentation in which part of the white matter is confused with the gray matter.

[0061] Figure 4 It is a comparison result of the convergence speed and segmentation accuracy of model M when or not using the attention loss function, where the blue curve is the case of using attention loss, and the orange is the case of not using attention loss. Obviously, the use of the attention loss function can effectively improve the convergence speed and segmentation accuracy of the model. Table 1 gives the segmentation accuracy of each category of the test samples under two different conditions, as well as the average accuracy of all categories, where CSF represents cerebrospinal fluid, GM is gray matter, WM is white matter, and Average is the average of the first three. It should be noted that the background category is relatively simple, with an accuracy of 100%, so it is not included in Table 1. The improvement in the segmentation accuracy of the test samples further demonstrates the effectiveness of the image segmentation loss function designed by the method of the present invention.

[0062] Table 1 Segmentation accuracy of test samples

[0063]

[0064] Finally, it should be understood that this embodiment is only a specific example of the implementation of the present invention and should not be a limitation of the scope of protection of the present invention. Without departing from the spirit and scope of the present invention, any equivalent modification or change to the above content should be included in the scope of protection claimed by the present invention.

Claims

1. A method for designing image segmentation loss function guided by label attention, characterized in that: The following steps are involved: Step 1: Obtain a multi-category image segmentation dataset, perform necessary preprocessing on the image data to make it meet the neural network processing requirements, and divide it into a training set S and a test set T; Step 2: For the multi-category image segmentation task, a deep learning model M with an attention mechanism is constructed, and the attention submodules of each category together constitute a multi-category attention module A; Step 3: Let A = [A1, A2, ..., A c ] represents the calculation result of the multi-category attention module, Y = [Y1, Y2, ..., Y c ] represents the labels corresponding to each category, and the attention loss function of model M is designed according to the following pattern: Loss Attention =Loss(f(A)),Y), Where Loss(·) represents any valid loss function, and f(·) represents the compression transformation function; the function of the compression transformation function is to integrate the calculation results of the attention module into a tensor of the label size in order to calculate the loss relative to the label; Step 4: Embed the module that implements the compression transformation function f(·) into the model M, and embed the attention loss function Loss Attention Incorporate into the loss function of model M; Step 5: Use the training set S to train the model M; Step 6: Use the test set T to test the model M and evaluate the effect of its loss function.

2. The method for designing a loss function for image segmentation guided by label attention according to claim 1, characterized in that: The compression transformation function f(·) in step 3 is implemented by a neural network module. The structure and form of the neural network module are not limited as long as it can meet its functions.

3. The method for designing a loss function for image segmentation guided by label attention according to claim 1, characterized in that: The specific implementation of the attention loss function Loss(·) in step 3 is: Loss Attention (x′,y)=Loss CE (x′,y)+Loss Dice (x′,y), Where x′=f(A)(A=[A1,A2,...,A m ]) represents the output of each attention submodule after the calculation results of each category are compressed and transformed and integrated; y = [Y1, Y2, ..., Y m ] represents the true label value corresponding to each category; Loss CE is the cross entropy loss function, Loss Dice is the Dice loss function, and their definitions are as follows: Among them, N and C represent the number of samples and the number of categories respectively, n and c represent the sample number and the category number respectively, and x nc is the model output result of a specific sample, y nc is the label value of a specific sample.

Images