Fetal brain tissue image semi-supervised segmentation method and system

Abstract

The invention relates to a semi-supervised segmentation method for a fetal brain tissue image, and the method comprises the steps: carrying out the noise addition of a label of a labeled image, inputting a diffusion U-Net encoder and a diffusion U-Net decoder, obtaining a segmentation result outputted by a supervised diffusion model, and calculating a loss function Ldeno; for a diffusion U-Net encoder, decoding by using a U-Net decoder to obtain a decoding segmentation result, and calculating a loss function Ldiff; inputting the label-free image into a diffusion U-Net encoder, respectively obtaining segmentation results of the label-free image by using a diffusion U-Net decoder and a U-Net decoder, and calculating a loss function Lu; the method comprises the following steps of: for a U-Net decoder, calculating a loss function Luconsis by using a guide consistency module; and a final loss function is calculated, a diffusion U-Net encoder is used for encoding, a U-Net decoder is used for decoding, and a final segmentation result of the fetal brain tissue image is obtained. According to the method, less label data can be used, more attention is paid to learning of difficult categories, the problem of category imbalance encountered by medical image segmentation is solved, and a high-accuracy image segmentation effect is achieved.

Description

Technical Field

[0001] The invention relates to a semi-supervised segmentation method and system for fetal brain tissue images. Background Art

[0002] At present, the diagnosis of fetal brain developmental abnormalities mainly relies on manual judgment, but manual judgment is highly dependent on the level of experts, cannot be quantitatively analyzed, has a large workload, and different experts make different judgments on the same data, and the same expert makes different judgments at different time periods. In recent years, with the development of computer vision technology, the latest research has begun to automatically segment different tissues based on artificial intelligence methods, which can not only speed up the expert's diagnosis process, but also can quantitatively analyze brain capacity through segmentation results, and can also remove subjective factors between different experts.

[0003] At the same time, as a medical image segmentation problem, the fetal brain tissue segmentation task has many medical imaging-related difficulties: First, medical images involve the privacy of case data, and the amount of data is very limited. Second, the annotation of medical images usually requires professionals with medical background knowledge, which leads to high cost and heavy workload of the annotation task.

[0004] Prior Art Huang et al. designed a lightweight fetal brain tissue segmentation model that integrates a CNN-based encoder-decoder architecture with a transformer module called context transformer block (CoT). The CoT block integrates context information mining and self-attention learning into a unified architecture. In addition, a hybrid dilated convolution module is introduced in the last layer of the decoder.

[0005] However, the prior art of Huang et al. has at least the following disadvantages:

[0006] First, most of them focus on the field of fully supervised learning and cannot solve the problem of data labeling difficulties.

[0007] Second, no weights are set for learning of all categories, and learning of difficult and easy categories is treated equally.

[0008] Third, the information transfer effect of labeled samples on unlabeled samples is not fully utilized. Summary of the invention

[0009] In view of this, it is necessary to provide a semi-supervised segmentation method and system for fetal brain tissue images.

[0010] The present invention provides a semi-supervised segmentation method for fetal brain tissue images. The method is applied to a diffusion model. The diffusion model includes: a diffusion U-Net encoder, a diffusion U-Net decoder and a U-Net decoder electrically connected to each other. The method includes the following steps: S1, obtaining a fetal brain tissue image; the fetal brain tissue image includes: a labeled image and an unlabeled image; S2, performing noise processing on the label of the labeled image to generate a noise image, and then inputting the noise image, the time of generating the noise image and the labeled image into the diffusion U-Net encoder and the diffusion U-Net decoder to obtain a segmentation result output by the supervised diffusion model, and finally calculating the loss function L with the label of the labeled image. deno ; S3, for the diffusion U-Net encoder, use the U-Net decoder to decode the intermediate information to obtain the decoded segmentation result, and then calculate the loss function L with the label of the labeled image diff ; S4 inputs the unlabeled image into the diffusion U-Net encoder to extract features, uses the diffusion U-Net decoder and the U-Net decoder to obtain the segmentation results of the unlabeled image, and then calculates the loss function L u ; S5, use the guided consistency module for the U-Net decoder to calculate the loss function L u_consis ; S6, according to the obtained loss function L deno , L diff , L u , L u_consis , calculate the final loss function, use the diffusion U-Net encoder for encoding, and then use the U-Net decoder for decoding to obtain the final segmentation result of the fetal brain tissue image

[0011] Preferably, the labeled image includes: a labeled image X L and its label Y L , X is a preprocessed three-dimensional fetal brain MRI image, the size of X is CxDxHxW, where: C represents the number of channels, D represents the depth, H represents the height, and W represents the width; the unlabeled image includes: unlabeled image X U .

[0012] Preferably, the step S2 comprises:

[0013] Use Diffused U-Net encoder and Diffused U-Net decoder for perception, given a labeled image X L ∈R D×W×H and label Y L ∈R D×W×H , first convert the labels into one-hot encoding format Then we continuously add t steps of noise ∈ to obtain the noisy label This is the forward process of diffusion:

[0014]

[0015] Then, the noise label With labeled image X L Connect as the input of the diffusion U-Net encoder;

[0016] Next, and time t as the input of the diffusion U-Net encoder to generate multi-scale features embedded in the time step Among them, i indicates which scale, F indicates the basic feature size, Further diffused U-Net decoder is used to predict clear labels

[0018] Preferably, the step S3 comprises:

[0019] The U-Net decoder deconvolves the features of the last layer output by the diffuse U-Net encoder in step S1 to expand the feature map size and reduce the number of channels, and then performs a convolution operation after splicing each layer with the intermediate output of the diffuse U-Net encoder, and finally outputs the segmentation result.

[0020] Preferably, the step S3 further comprises:

[0021] Define τ as the number of cumulative iterations, λ k,e is the Dice score of the kth class in the eth iteration, and Respectively represent the categories that have not been successfully learned and successfully learned after the τth iteration, d k,e Defined as the difficulty of the kth category after the eth iteration, then: categories with faster learning speed have smaller d k,e , the corresponding weight in the loss function is smaller to slow down the learning speed;

[0022]

[0023] definition is the inverse Dice weight, which is inversely proportional to the Dice score. The lower the Dice score, the greater the weight, forcing the diffusion model to pay more attention to the learning of difficult categories. The overall difficulty perception weight of the kth category is defined as

[0024]

[0025] The difficulty-aware loss function is defined as:

[0026]

[0027] Preferably, the step S4 comprises:

[0028] For the same unlabeled image X U , using the diffusion U-Net decoder and the U-Net decoder to obtain the unlabeled image X U Based on the two segmentation results above, the loss function L is calculated. u .

[0029] Preferably, the step S5 comprises:

[0030] Features F obtained from the Diffused U-Net encoder for labeled and unlabeled images l and F u ,After passing through the U-Net decoder, the features of the middle layer are obtained respectively Combination becomes Feat_l and Combined into Feat_u, Feat_l obtains Feat_Maps_l and attention query updates_Qs_l through the guided consistency module, Feat_u obtains Feat_Maps_consis through the guided consistency module, and then Feat_Maps_consis and updates_Qs_l are input into the guided consistency module to obtain Feat_Maps_u. Finally, the difference between Feat_Maps_l and Feat_Maps_u is calculated to obtain the loss function L u_consis .

[0031] Preferably, the step S5 further comprises:

[0032] When there is no input Giuded_Q, the guided consistency module is directly initialized to 0. After feature projection and feature normalization, Feat is sent to the classification decoder together with Giuded_Q. Then, for each scale feature, the fused multi-head attention feature and the updated Giuded_Q are output. Finally, the results of each scale are collected to obtain Feat_map and update_Q.

[0033] Preferably, the step S5 further comprises:

[0034] The classification decoder inputs the normalized Feat and Guided_Q, and then uses linear mapping to regard the linearly mapped Feat as V and K in the attention mechanism, and the mapped Guided_Q as Q. Q and K are multiplied to get a new Feat, and Feat and V are multiplied to get a new Guided_Q. Then, they are added to the original variables after norm, mlp, and dropout to obtain the fused multi-head attention features and updated Guided_Q.

[0035] Preferably, the step S6 comprises:

[0036] The final loss function L is given by L deno , L diff , L u , L u_consis composition:

[0037] L=L deno +L diff +μ(L u +L u_consis ).

[0038] The present invention solves the problem that there are a large number of unlabeled samples and a small number of labeled samples in medical image data, and further improves the accuracy of fetal brain tissue segmentation. The present invention introduces a diffusion model into the semi-supervised segmentation of fetal brain tissue images. Compared with ordinary fully supervised algorithms, the present invention can use less labeled data to achieve similar results. The network adopts a single encoder and dual decoder structure, and introduces difficulty-aware loss when calculating the loss. The difficulty-aware loss can make the diffusion model pay more attention to the learning of difficult categories to solve the category imbalance problem encountered in medical image segmentation. At the same time, the present invention introduces a guided consistency module, which can learn the information of labeled samples, and through the attention mechanism, use the learned information to guide the learning of unlabeled samples to improve the learning effect of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0039] Figure 1 A schematic diagram of the overall framework of the technical solution provided by the embodiment of the present invention;

[0040] Figure 2 A flowchart of the semi-supervised segmentation method of fetal brain tissue images of the present invention;

[0041] Figure 3 The loss function L provided in the embodiment of the present invention is u_consis A schematic diagram of the specific calculation method of

[0042] Figure 4 A schematic diagram of a specific flow of a boot consistency module provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0043] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0044] The overall framework diagram of the technical solution of the embodiment of the present invention is as follows: Figure 1 shown.

[0045] The diffusion model includes: a diffusion U-Net encoder, a diffusion U-Net decoder, and a U-Net decoder. The diffusion U-Net encoder, the diffusion U-Net decoder, and the U-Net decoder are electrically connected.

[0046] During the training phase, a diffusion model is used to capture invariant features, and a guided consistency module is used to guide the learning of more precise features.

[0047] Input images include: labeled images X L and its label Y L , unlabeled image X U , and further:

[0048] X is a preprocessed three-dimensional fetal brain MRI image, and the size of X is (CxDxHxW), where C represents the number of channels, D represents the depth, H represents the height, and W represents the width. In this embodiment, since the image is a three-dimensional black-and-white image, C is 1.

[0049] See also Figure 2 , which is a flowchart of a preferred embodiment of the semi-supervised segmentation method for fetal brain tissue images of the present invention.

[0050] Step S1, obtaining a fetal brain tissue image. The fetal brain tissue image includes: a labeled image and an unlabeled image.

[0051] Step S2, add noise to the label of the labeled image to generate a noise image, then input the noise image, the time of generating the noise image and the labeled image into the diffusion U-Net encoder and the diffusion U-Net decoder to obtain the segmentation result output by the supervised diffusion model, and finally calculate the loss function L with the label of the labeled image. deno Specifically:

[0052] In the embodiment of the present invention, for the labeled image X L Label Y L Noise processing is performed to generate a noise image, and then the noise image, the time t of generating the noise image and the labeled image X are combined. L Input the diffusion U-Net encoder and diffusion U-Net decoder to obtain the segmentation result output by the supervised diffusion model, and then use the above segmentation result and label Y L , calculate the loss function L deno. That is:

[0053] Use L deno The difference between the segmentation results produced by the diffuse U-Net decoder and the original labels is calculated to learn image information and enhance the accuracy of the diffusion model.

[0054] Further:

[0055] Use the diffusion model for perception, including the diffusion U-Net encoder part and the diffusion U-Net decoder part. Given a labeled image X L ∈R D×W×H and label Y L ∈R D×W×H , first convert the labels into one-hot encoding format Then we continuously add t steps of noise ∈ to obtain the noisy label This is the forward process of diffusion:

[0056]

[0057] Then, the noise label With labeled image X L The concatenation is used as the input of the diffusion U-Net encoder. It is worth noting that the high-level features of different data streams are different.

[0058] For the denoising process, first, and time t as the input of the diffusion U-Net encoder to generate multi-scale features embedded in the time step Where i indicates which scale, F indicates the basic feature size, Further diffused U-Net decoder is used to predict clear labels

[0059] Step S3: for the diffusion U-Net encoder, use the U-Net decoder to decode the intermediate information to obtain the decoded segmentation result, and then calculate the loss function L with the label of the labeled image. diff Specifically:

[0060] In the embodiment of the present invention, for the diffusion U-Net encoder, the U-Net decoder is used to decode the intermediate information to obtain a decoded segmentation result, and then the decoded segmentation result and the label Y are combined. L Calculate the loss function L diff . That is:

[0061] L diff is the difficulty-aware loss function, through L diff Calculate the difficult categories and give different learning weights to different categories so that the diffusion model can learn more information about difficult categories.

[0062] Further:

[0063] The U-Net decoder deconvolves the features of the last layer output by the diffuse U-Net encoder in step S1 to expand the feature map size and reduce the number of channels, and then performs a convolution operation after splicing each layer with the intermediate output of the diffuse U-Net encoder, and finally outputs the segmentation result.

[0064] In order to alleviate the class imbalance problem encountered in medical image segmentation, this embodiment adopts a difficulty-aware weight strategy, which forces the diffusion model to focus on the most difficult categories, that is, categories with slower learning speed and poorer performance. The difficulty is learned through the probability map p generated by the U-Net decoder.

[0065] Specifically, in this embodiment, τ is defined as the number of cumulative iterations, which is empirically set to 50; λ k,e is the Dice score of the kth class in the eth iteration; and Respectively represent the categories of failure and success after the τth iteration; d k,e is defined as the difficulty of the kth category after the eth iteration. Then: categories with faster learning speed have smaller d k,e , the corresponding weights in the loss function are smaller to slow down the learning speed.

[0066]

[0067] In this embodiment, the definition is the inverse Dice weight, which is inversely proportional to the Dice score. The lower the Dice score, the greater the weight, forcing the diffusion model to pay more attention to the learning of difficult categories. The overall difficulty perception weight of the kth category is defined as Empirically, α is set to 0.2 to mitigate the impact of outliers.

[0068]

[0069] The difficulty-aware loss function is defined as follows:

[0070]

[0071] Step S4, input the unlabeled image into the diffusion U-Net encoder to extract features, and then use the diffusion U-Net decoder and U-Net decoder to obtain the segmentation results of the unlabeled image respectively, and then calculate the loss function L u Specifically:

[0072] In this embodiment of the present invention, the unlabeled image X UInput to the diffusion U-Net encoder to extract features, and then use the diffusion U-Net decoder and U-Net decoder to obtain the unlabeled image X U Then, based on the above two segmentation results, calculate the loss function L u . That is:

[0073] Calculate the consistency L of the prediction results of the diffusion U-Net decoder and the U-Net decoder for the same unlabeled image u .

[0074] Step S5, using the guided consistency module for the U-Net decoder, calculates the loss function L u_consis Specifically:

[0075] In the embodiment of the present invention, a guided consistency module is used to calculate the consistency L of the unlabeled image segmentation result guided by the labeled image and the unguided unlabeled image segmentation result at multiple scales. u_consis .

[0076] Further:

[0077] The working principle of the guided consistency module is to achieve the effect of guided consistency by learning the attention query updates_Qs_l of the labeled image, and then comparing the consistency of the multi-scale features output when the attention query is not used with the multi-scale features output when the attention query is used.

[0078] like Figure 3 As shown, the features F obtained from the diffusion U-Net encoder for labeled and unlabeled images l and F u ,After passing through the U-Net decoder, the features of the middle layer are obtained respectively Combination becomes Feat_l and Combined into Feat_u, Feat_l obtains Feat_Maps_l and attention query updates_Qs_l through the guided consistency module, Feat_u obtains Feat_Maps_consis through the guided consistency module, and then Feat_Maps_consis and updates_Qs_l are input into the guided consistency module to obtain Feat_Maps_u. Finally, the difference between Feat_Maps_l and Feat_Maps_u is calculated to obtain the loss function L u_consis .

[0079] The specific process of booting the consistency module is as follows Figure 4As shown in the figure, when Giuded_Q is not input, it is directly initialized to 0. After feature projection and feature normalization, Feat is sent to the classification decoder together with Giuded_Q. Then, for each scale feature, the fused multi-head attention feature and the updated Giuded_Q are output. Finally, the results of each scale are collected to obtain Feat_map and update_Q.

[0080] The process of the classification decoder is as follows: input the normalized Feat and Guided_Q, then use linear mapping respectively, regard the linearly mapped Feat as V and K in the attention mechanism, regard the mapped Guided_Q as Q, multiply Q and K to get a new Feat, multiply Feat and V to get a new Guided_Q, then pass norm, mlp, dropout respectively and add to the original variable to get the fused multi-head attention feature and updated Guided_Q.

[0081] Step S6, calculating the final loss function, encoding using the diffusion U-Net encoder, and then decoding using the U-Net decoder to obtain the final segmentation result of the fetal brain tissue image.

[0082] Specifically:

[0083] The embodiment of the present invention uses the following formula to calculate the final loss function L, and the final loss function L is composed of L deno , L diff , L u composition:

[0084] L=L deno +L diff +μ(L u +L u_consis )

[0085] According to the calculated final loss function, the diffuse U-Net encoder is used for encoding, and then the U-Net decoder is used for decoding to obtain the final segmentation result.

[0086] The present invention can achieve high-accuracy image segmentation effect when there are only a small number of labeled samples.

[0087] Although the present invention has been described with reference to the current preferred embodiments, those skilled in the art should understand that the above-mentioned preferred embodiments are only used to illustrate the present invention and are not used to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the scope of protection of the present invention.

Claims

1. A semi-supervised segmentation method for fetal brain tissue images, characterized in that: The method is applied to a diffusion model, the diffusion model comprising: a diffusion U-Net encoder, a diffusion U-Net decoder and a U-Net decoder electrically connected to each other, and the method comprises the following steps: S1, obtaining a fetal brain tissue image; the fetal brain tissue image includes: a labeled image and an unlabeled image; S2, add noise to the label of the labeled image to generate a noisy image, then input the noisy image, the time of generating the noisy image and the labeled image into the diffusion U-Net encoder and the diffusion U-Net decoder to obtain the segmentation result output by the supervised diffusion model, and finally calculate the loss function L with the label of the labeled image deno ; S3, for the diffusion U-Net encoder, use the U-Net decoder to decode the intermediate information to obtain the decoded segmentation result, and then calculate the loss function L with the label of the labeled image diff ; S4 inputs the unlabeled image into the diffusion U-Net encoder to extract features, uses the diffusion U-Net decoder and the U-Net decoder to obtain the segmentation results of the unlabeled image, and then calculates the loss function L u ; S5, using the guided consistency module for the U-Net decoder, calculates the loss function L u_consis ; S6, according to the obtained loss function L deno , L diff , L u , L u_consis , calculate the final loss function, use the diffusion U-Net encoder for encoding, and then use the U-Net decoder for decoding to obtain the final segmentation result of the fetal brain tissue image.

2. The method according to claim 1, characterized in that The labeled images include: labeled images X L and its label Y L , X is a preprocessed three-dimensional fetal brain MRI image, the size of X is CxDxHxW, where: C represents the number of channels, D represents the depth, H represents the height, and W represents the width; the unlabeled image includes: unlabeled image X U .

3. The method according to claim 2, characterized in that The step S2 comprises: Use Diffused U-Net encoder and Diffused U-Net decoder for perception, given a labeled image X L ∈R D×W×H and label Y L ∈R D×W×H , first convert the labels into one-hot encoding format Then we continuously add t steps of noise ∈ to obtain the noisy label This is the forward process of diffusion: Then, the noise label With labeled image X L Connect as the input of the diffusion U-Net encoder; Next, and time t as the input of the diffusion U-Net encoder to generate multi-scale features embedded in the time step Among them, i indicates which scale, F indicates the basic feature size, Further diffused U-Net decoder is used to predict clear labels 4. The method according to claim 3, characterized in that The step S3 comprises: The U-Net decoder deconvolves the features of the last layer output by the diffuse U-Net encoder in step S1 to expand the feature map size and reduce the number of channels, and then performs a convolution operation after splicing each layer with the intermediate output of the diffuse U-Net encoder, and finally outputs the segmentation result.

5. The method according to claim 4, characterized in that The step S3 further comprises: Define τ as the number of cumulative iterations, λ k,e is the Dice score of the kth class in the eth iteration, and Respectively represent the categories that have not been successfully learned and successfully learned after the τth iteration, d k,e Defined as the difficulty of the kth category after the eth iteration, then: categories with faster learning speed have smaller d k,e , the corresponding weight in the loss function is smaller to slow down the learning speed; definition is the inverse Dice weight, which is inversely proportional to the Dice score. The lower the Dice score, the greater the weight, forcing the diffusion model to pay more attention to the learning of difficult categories. The overall difficulty perception weight of the kth category is defined as The difficulty-aware loss function is defined as:

6. The method according to claim 5, characterized in that The step S4 comprises: For the same unlabeled image X U , using the diffusion U-Net decoder and the U-Net decoder to obtain the unlabeled image X U Based on the two segmentation results above, the loss function L is calculated. u .

7. The method according to claim 6, characterized in that The step S5 comprises: Features F obtained from the Diffused U-Net encoder for labeled and unlabeled images l and F u ,After passing through the U-Net decoder, the features of the middle layer are obtained respectively Combination becomes Feat_l and Combined into Feat_u, Feat_l obtains Feat_Maps_l and attention query updates_Qs_l through the guided consistency module, Feat_u obtains Feat_Maps_consis through the guided consistency module, and then Feat_Maps_consis and updates_Qs_l are input into the guided consistency module to obtain Feat_Maps_u. Finally, the difference between Feat_Maps_l and Feat_Maps_u is calculated to obtain the loss function L u_consis .

8. The method according to claim 7, characterized in that The step S5 further comprises: When there is no input Giuded_Q, the guided consistency module is directly initialized to 0. After feature projection and feature normalization, Feat is sent to the classification decoder together with Giuded_Q. Then, for each scale feature, the fused multi-head attention feature and the updated Giuded_Q are output. Finally, the results of each scale are collected to obtain Feat_map and update_Q.

9. The method according to claim 8, characterized in that The step S5 further comprises: The classification decoder inputs the normalized Feat and Guided_Q, and then uses linear mapping to regard the linearly mapped Feat as V and K in the attention mechanism, and the mapped Guided_Q as Q. Q and K are multiplied to get a new Feat, and Feat and V are multiplied to get a new Guided_Q. Then, they are added to the original variables after norm, mlp, and dropout to obtain the fused multi-head attention features and updated Guided_Q.

10. The method according to claim 9, characterized in that The step S6 comprises: The final loss function L is given by L deno , L diff , L u , L u_consis composition: L=L deno +L diff +μ(L u +L u_consis )

Images