Medical image feature map denoising method and system based on diffusion model

Abstract

The present invention proposes a method and system for denoising medical image feature maps based on a diffusion model, relating to the field of medical image segmentation. The method comprises extracting an original feature map from a model; progressively adding Gaussian noise to the original feature map to obtain a series of noise feature maps; sequentially inputting the resulting noise feature map into multiple FMD modules connected in series for denoising, obtaining denoising results corresponding to the series of noise feature maps; training the FMD module using the noise feature map and the corresponding denoising results; and inputting the original feature map into the trained FMD module for denoising. Based on a designed denoising diffusion probability model, the present invention denoises the feature map information, thereby enabling the medical image segmentation model to more accurately identify the features of organs or lesions. Furthermore, the plug-and-play feature map denoising module can be conveniently combined with various existing medical image segmentation models without the need for model retraining, significantly saving development time and costs.

Description

Technical Field

[0001] The present invention belongs to the technical field of medical image segmentation, and in particular relates to a method and system for denoising a medical image feature map based on a diffusion model. Background Art

[0002] The statements in this section merely provide background information related to the present invention and do not necessarily constitute prior art.

[0003] Medical image segmentation is a crucial task in medical image processing. Its primary goal is to separate and segment medical images to provide useful information, helping doctors more accurately visualize tumors, lesions, and organs, thereby improving disease diagnosis and treatment planning. Accurate segmentation is therefore crucial for patient care, physician decision-making, and biomedical research. To achieve fast and accurate medical image segmentation, many researchers have proposed various segmentation models.

[0004] Among them, the UNet model has become a classic model in the field of medical image segmentation due to its excellent performance. Its success has led to the U-shaped network structure becoming the mainstream paradigm of the current medical image segmentation model. With the deepening of research, some researchers have proposed segmentation models similar to UNet, and divided them into two groups according to the feature extraction method. The first group only uses convolution for feature extraction, which includes models such as VNet, UNet++, and Res-UNet. The second group uses transformers for feature extraction, including models such as TransUNet, Swin-UNet, UNETR, and MedT. These models have been widely used in medical image segmentation and have demonstrated impressive results. However, medical images usually have more noise, which makes it challenging to distinguish the lesion area or organ in the image from the background, which may hinder the accuracy of segmentation. Summary of the Invention

[0005] In order to overcome the shortcomings of the above-mentioned existing technologies, the present invention provides a medical image feature map denoising method and system based on a diffusion model. Based on the designed denoising diffusion probability model, the feature map information is denoised, so that the medical image segmentation model can more accurately identify the characteristics of organs or lesions, effectively remove noise and improve segmentation accuracy; at the same time, this plug-and-play feature map denoising module can be easily combined with various existing medical image segmentation models without the need to retrain the model, thereby greatly saving development time and cost.

[0006] To achieve the above objectives, one or more embodiments of the present invention provide the following technical solutions:

[0007] A first aspect of the present invention provides a medical image feature map denoising method based on a diffusion model.

[0008] The medical image feature map denoising method based on the diffusion model includes the following steps:

[0009] Extract the feature map from the model and use the extracted feature map as the original feature map;

[0010] Add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps;

[0011] The noise feature maps finally obtained are sequentially input into multiple FMD modules connected in series for denoising, and denoising results corresponding to a series of noise feature maps are obtained. The noise feature maps and the corresponding denoising results are used to train the FMD modules to obtain a trained FMD module;

[0012] The original feature map is input into the trained FMD module for denoising.

[0013] Preferably, Gaussian noise is added to the original feature map in a progressive manner, specifically including:

[0014] Gaussian noise is gradually added to the original feature map f extracted from the model. The noise level of the original feature map increases from 0 to T, that is, from no noise to maximum noise, and a series of noise feature maps f0, f1, ..., f are obtained. T .

[0015] Preferably, the process of gradually introducing noise into the original feature map f is:

[0016]

[0017] I represents the identity matrix, β1, ..., β T represents the noise variance at each step in the process.

[0018] Preferably, the noise characteristic graph is expressed as:

[0019]

[0020] in,

[0021] Preferably, the denoising process is expressed as:

[0022] p θ (f' t-1 |f t )~N(f' t-1 ;μ θ (f t ,t),ε θ (f t ,t))

[0023] Here, θ represents a parameter involved in the denoising process; εθ is the noise predicted by the FMD module;

[0024] From f t Prediction f′ t-1 The formula is:

[0025]

[0026] σ t represents the variance that the model can learn, and z is the random component.

[0027] Preferably, the FMD module is trained using the mean square error loss of the noise feature map and the corresponding denoising result.

[0028] Preferably, the training loss function of the FMD module is:

[0029]

[0030] A second aspect of the present invention provides a medical image feature map denoising system based on a diffusion model.

[0031] Medical image feature map denoising system based on diffusion model, including:

[0032] The acquisition module is configured to: extract the feature map from the model and use the extracted feature map as the original feature map;

[0033] The noise adding module is configured to: add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps;

[0034] The training module is configured to: sequentially input the noise feature map obtained finally into a plurality of FMD modules connected in series for denoising, thereby obtaining denoising results corresponding to the series of noise feature maps; and train the FMD module using the noise feature map and the corresponding denoising results to obtain a trained FMD module;

[0035] The denoising module is configured to input the original feature map into the trained FMD module for denoising.

[0036] A third aspect of the present invention provides a computer-readable storage medium having a program stored thereon, which, when executed by a processor, implements the steps in the medical image feature map denoising method based on a diffusion model as described in the first aspect of the present invention.

[0037] The fourth aspect of the present invention provides an electronic device, comprising a memory, a processor, and a program stored in the memory and executable on the processor. When the processor executes the program, the steps of the medical image feature map denoising method based on the diffusion model as described in the first aspect of the present invention are implemented.

[0038] One or more of the above technical solutions have the following beneficial effects:

[0039] The present invention provides a medical image feature map denoising method and system based on a diffusion model. Based on the designed denoising diffusion probability model, the feature map information is denoised, so that the medical image segmentation model can more accurately identify the characteristics of organs or lesions, effectively remove noise and improve segmentation accuracy. At the same time, this plug-and-play feature map denoising module can be easily combined with various existing medical image segmentation models without the need to retrain the model, thereby greatly saving development time and cost.

[0040] Advantages of additional aspects of the present invention will be given in part in the following description and in part will be obvious from the following description, or will be learned through practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] The accompanying drawings, which constitute a part of the present invention, are used to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute improper limitations on the present invention.

[0042] Figure 1 This is a flow chart of the method of the first embodiment.

[0043] Figure 2 It is the location map of the FMD module combined with the existing segmentation model.

[0044] Figure 3 This is a system structure diagram of the second embodiment. DETAILED DESCRIPTION

[0045] It should be noted that the following detailed descriptions are exemplary and intended to provide further explanation of the present invention. Unless otherwise specified, all technical and scientific terms used herein have the same meaning as commonly understood by those skilled in the art to which the present invention belongs.

[0046] It should be noted that the terms used herein are for describing particular embodiments only and are not intended to be limiting of exemplary embodiments according to the present invention.

[0047] In the absence of conflict, the embodiments of the present invention and the features thereof may be combined with each other.

[0048] The overall idea proposed by the present invention is:

[0049] The denoising diffusion probabilistic model (DDPM) has garnered significant attention as a powerful generative model. Although originally designed for image generation, DDPM has now found applications in various other fields, including image deblurring and denoising. In the current field of medical image segmentation, researchers have increasingly combined DDPM with existing segmentation models. Wolleb et al. employed a diffusion model to address the problem of 2D medical image segmentation, improving the robustness of the segmentation results by fusing the outputs of each diffusion step using a summation method during testing. The Diff-UNet model employs DDPM for the task of segmenting 3D medical images. In the MedSegDiff model, the original input image is considered to contain rich object information and is difficult to segment, while the segmentation map at any given time in the diffusion model contains enhanced object regions but is inaccurate. Based on this observation, the authors proposed a fusion strategy. Building on the MedSegDiff model, researchers have persistently explored a superior version, the MedSegDiff-V2 model. This new model integrates the Transformer with the DDPM, further advancing image segmentation accuracy.

[0050] The aforementioned work primarily leverages the generative capabilities of DDPM to enhance the target region to be segmented by generating a set of implicit ground truth distributions and merging them with the feature maps extracted by the segmentation model. Unlike existing work, we focus on DDPM's denoising capabilities for feature maps and design a general feature map denoising (FMD) module based on DDPM.

[0051] Example 1

[0052] This embodiment discloses a medical image feature map denoising method based on a diffusion model.

[0053] like Figure 1 As shown in FIG, the medical image feature map denoising method based on the diffusion model includes the following steps:

[0054] Extract the feature map from the model and use the extracted feature map as the original feature map;

[0055] Add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps;

[0056] The noise feature maps finally obtained are sequentially input into multiple FMD modules connected in series for denoising, and denoising results corresponding to a series of noise feature maps are obtained. The noise feature maps and the corresponding denoising results are used to train the FMD modules to obtain a trained FMD module;

[0057] The original feature map is input into the trained FMD module for denoising.

[0058] Specifically, the present invention adopts the following technical solutions:

[0059] This example proposes a plug-and-play medical image feature map denoising module, called the FMD module. The explanation of the principles and processes of the FMD module is divided into two parts. In the first part, we introduce the training process of the FMD module and its associated loss function. In the second part, we discuss in detail how to use the FMD module to remove noise from feature maps.

[0060] Part 1: Gaussian noise is gradually added to the feature map f extracted from the model to obtain a series of noise features f0, f1, ..., f t-1 , f t ,……,f T , where t∈{1,...,T}. The noise level of the feature map increases from 0 (no noise) to T (maximum noise). After that, the noise feature map f at time T T It is sent to the FMD module. The FMD module predicts the corresponding denoising result f′ based on the noise feature map. T-1 Then the denoising result is input into the FMD module to obtain f′ T-2 This process is repeated step by step to obtain f′ t-1 , where t∈{1,...,T}, and finally we get f′0.

[0061] Since f′ t-1 The corresponding f t-1 It is known that the FMD module can be trained with mean squared error loss to improve its ability to denoise feature maps.

[0062] In short, the process of learning denoising capabilities of the FMD module can be divided into two main steps. The first step is to add Gaussian noise to the feature map in a progressive manner. The second step is to use FMD to denoise the noisy feature map.

[0063] To elaborate, the process of gradually introducing noise to the feature map f can be described by the following formula:

[0064]

[0065] In this formula, I represents the identity matrix, β1, ..., β T represents the noise variance at each step in the process. This formula allows us to gradually add Gaussian noise to the image at each iteration.

[0066] After t steps, the feature map obtained is as follows:

[0067]

[0068] A reparameterization technique can be used to express f t As a function of f0, it looks like this:

[0069]

[0070] The denoising process can be expressed by the following formula:

[0071] p θ (f' t-1 |f t )~N(f' t-1 ;μ θ (f t ,t),ε θ (f t ,t))

[0072] Here, θ represents a parameter involved in the denoising process, μ θ (f t ,t) represents the mean. Through the above formula, we can get f t Prediction f′ t-1 :

[0073]

[0074] In the above formula, σ t Represents the variance that the model can learn, z is the random component, and this process leads to a random sampling process. Note that ε θ is the noise predicted by the FMD module, the input of which is

[0075] According to the above formula, during the denoising process, f t The noise scheme ε is subtracted from θ ( t ,) must be learned by the model. Therefore, the training loss function of the FMD module denoising is as follows:

[0076]

[0077] Part 2: After completing Part 1, continue with the process outlined in Part 2. Unlike Part 1, Part 2 does not add noise, but only denoises the original feature map f. This involves feeding the feature map f extracted by the model into the FMD module and performing T denoising iterations to obtain the denoised feature map f. T In short, the first part is executed first, using Loss FMD The function implements the end-to-end training of the FMD module, giving it denoising capabilities. Subsequently, the second part is to use the FMD module to denoise the feature maps f extracted by various models to improve the accuracy of their feature extraction.

[0078] like Figure 2 As shown, this embodiment 1 provides the integration of the FMD module and the commonly used segmentation models. These models are as follows from top to bottom and from left to right:

[0079] UNet: A classic medical image segmentation model with a network structure consisting of three parts: an encoder, a decoder, and skip connections. The UNet structure has become the reference point for almost all current medical image segmentation model designs.

[0080] UNeXt: This model is very similar in structure to UNet, but has fewer parameters. It primarily consists of convolutional layers and MLP layers. The MLP layers function similarly to the Transformer, capturing global information, but with significantly lower computational requirements.

[0081] TransUNet: To address the limitations of convolutional operations, TransUNet incorporates a transformer into the encoder component based on the UNet model. This approach effectively captures global features and improves the accuracy of feature extraction in the model.

[0082] IB-TransUNet: To further enhance TransUNet, an information bottleneck block is incorporated into the transformer layer. This block compresses irrelevant features and retains only key features, making the model learning more accurate.

[0083] It can be seen that whether it is a CNN-based UNet model or other models that add MLP layers (UNeXt) or Transformer layers (TransUNet, IB-TransUNet) to the encoder, as long as the model contains convolutional layers, the FMD module can be added. In addition, the addition of the FMD module is very flexible and can be placed after the bottom or middle layer feature maps. This means that the FMD module does not need to be trained separately and can be trained end-to-end together with the segmentation model. By denoising the feature maps extracted by the model, the accuracy of the model's learned features can be improved, thereby improving the quality of the model's segmentation results.

[0084] like Figure 1 As shown, in this embodiment 1, the training process and functional block diagram of the FMD module include:

[0085] In the first part, we gradually add noise to the original feature map f, generating a series of feature maps. The noise level of these feature maps gradually increases from 0 (no noise) to T (maximum noise level). Next, we use the FMD module to gradually denoise these noisy feature maps to train the FMD module's denoising capabilities.

[0086] In the second part, the FMD module is used to denoise the feature maps extracted by each model. Specifically, the FMD module improves the performance and robustness of the neural network when processing noisy data by learning how to transform noisy feature maps into clearer and more accurate feature maps.

[0087] (I) To better illustrate the plug-and-play and effectiveness of our FMD module, we plugged the FMD module into four other state-of-the-art segmentation models. Among these models, experiments on UNet and UNeXt were conducted using the ISIC skin cancer dataset (Table 1). Experiments on TransUNet and IBTransUNet were conducted on the Synapse abdominal multi-organ segmentation dataset (Table 2).

[0088] Here is a brief introduction to the datasets used and the common evaluation indicators in medical image segmentation.

[0089] ISIC Skin Cancer Dataset: This dataset includes 2,594 skin images captured by cameras and their corresponding segmentation maps of skin lesions. The dataset is randomly split into three 80-20 parts.

[0090] Synapse Abdominal Multi-Organ Segmentation Dataset: This dataset consists of 30 CT scans of eight abdominal organs, including the left and right kidneys, aorta, spleen, gallbladder, liver, spleen, pancreas, and stomach. A total of 3779 axial contrast-enhanced abdominal clinical CT images were obtained. We randomly selected 18 samples (2212 axial slices) as the training set, and the remaining 12 samples (1567 axial slices) as the test set.

[0091] Evaluation Metrics: The DSC metric determines the similarity between two sets, with a value ranging from 0 to 1. A value of 0 indicates the worst segmentation, while a value of 1 indicates the best. The HD metric calculates the distance between two sets; smaller values indicate greater similarity. The IoU (Interference over Union) measures the positional accuracy of the segmentation results.

[0092] Table 1

[0093] Model DSC↑ IoU↑ UNet 88.87 80.28 UNeXt 90.58 82.97 UNet+FMD 90.59 83.12 UNeXt+FMD 91.02 83.76

[0094] Table 2

[0095] Model DSC↑ HD↓ aorta gallbladder Left kidney Right kidney liver pancreatic spleen Stomach TransUNet 78.58 32.08 88.16 64.37 81.10 76.58 94.53 58.88 87.57 77.43 IB-TransUNet 81.07 22.45 88.30 65.70 81.25 80.33 95.34 65.84 90.02 81.82 TransUNet+FMD 81.30 20.35 86.36 68.24 83.56 80.50 95.26 63.29 91.12 82.07 IB-TransUNet+FMD 81.77 22.01 88.67 66.94 82.59 81.37 95.23 66.65 90.84 81.86

[0096] Experimental results show that after adding the FMD module, the accuracy of the segmentation results of all four baseline models has been significantly improved. On the ISIC dataset, the two evaluation indicators of the UNeXt and UNet models increased by 0.44%-1.72% and 0.79%-2.84%, respectively. This shows that after integrating the FMD module, both models can effectively acquire and understand the characteristics of the lesions, and the accuracy is higher. On the Synapse dataset, the overall segmentation results of the IBTransUNet and TransUNet models improved by 0.70%-2.72%, and the quality of boundary segmentation improved by 0.44mm-11.73mm. The results show that using the FMD module for denoising can not only improve the accuracy of feature extraction, but also make the boundary segmentation smoother and the positioning more accurate. It is worth noting that the improvement of the FMD module is more obvious in the segmentation of small organs such as the gallbladder, kidney (left), kidney (right) and pancreas. The two models improve by 1.24%-3.87%, 1.34%-2.46%, 1.04%-3.92%, and 0.81%-4.41% respectively. These results demonstrate that our proposed FMD module can improve the segmentation results of various models on multi-class datasets and has good robustness.

[0097] (II) To study the effect of the number of denoising and denoising times (T) mentioned in the Summary of the Invention on the performance of the FMD module, we conducted comparative experiments using UNeXt and TransUNet as baseline models on the ISIC and Synapse datasets, respectively (Table 3):

[0098] Table 3

[0099]

[0100] As can be seen, with the addition of the FMD module, both models achieve substantially better DSC performance than the original model within the range of T = 2 to 9. However, as T increases, the FMD module's denoising capability becomes too strong, filtering out some features as noise, causing segmentation accuracy (DSC) to initially increase and then decrease. Furthermore, UNeXt achieves optimal performance at T = 3, while TransUNet achieves optimal performance at T = 5. This is because the two models use different datasets; UNeXt uses the ISIC dataset, which tends to have less noise than CT images and requires less noise removal, resulting in a smaller T value. Therefore, when processing medical images with significant noise, such as CT images, it is recommended to set T higher to enhance the module's denoising capabilities and thus help extract more accurate features. However, for medical images captured by cameras, such as dermatoscopes, which have less noise, using overly strong denoising capabilities may result in filtering out key features. To avoid this, it is recommended to set T lower.

[0101] Example 2

[0102] This embodiment discloses a medical image feature map denoising system based on a diffusion model.

[0103] like Figure 3 As shown in FIG, a medical image feature map denoising system based on a diffusion model includes:

[0104] The acquisition module is configured to: extract the feature map from the model and use the extracted feature map as the original feature map;

[0105] The noise adding module is configured to: add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps;

[0106] The training module is configured to: sequentially input the noise feature map obtained finally into a plurality of FMD modules connected in series for denoising, thereby obtaining denoising results corresponding to the series of noise feature maps; and train the FMD module using the noise feature map and the corresponding denoising results to obtain a trained FMD module;

[0107] The denoising module is configured to input the original feature map into the trained FMD module for denoising.

[0108] Example 3

[0109] The purpose of this embodiment is to provide a computer-readable storage medium.

[0110] A computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps in the medical image feature map denoising method based on a diffusion model as described in Example 1 of the present disclosure.

[0111] Example 4

[0112] The purpose of this embodiment is to provide an electronic device.

[0113] An electronic device includes a memory, a processor, and a program stored in the memory and executable on the processor. When the processor executes the program, the steps of the medical image feature map denoising method based on the diffusion model as described in Example 1 of the present disclosure are implemented.

[0114] The steps involved in the apparatuses of Examples 2, 3, and 4 above correspond to those of Method Example 1. For detailed implementations, please refer to the relevant description of Example 1. The term "computer-readable storage medium" should be understood to mean a single medium or multiple media containing one or more instruction sets; it should also be understood to include any medium capable of storing, encoding, or carrying an instruction set for execution by a processor and causing the processor to perform any method of the present invention.

[0115] Those skilled in the art will appreciate that the modules or steps of the present invention described above can be implemented using a general-purpose computer device. Alternatively, they can be implemented using program code executable by a computing device, which can then be stored in a storage device and executed by the computing device. Alternatively, they can be fabricated into separate integrated circuit modules, or multiple modules or steps can be fabricated into a single integrated circuit module for implementation. The present invention is not limited to any specific combination of hardware and software.

[0116] Although the above describes the specific embodiments of the present invention in conjunction with the accompanying drawings, it is not intended to limit the scope of protection of the present invention. Those skilled in the art should understand that various modifications or variations that can be made by those skilled in the art on the basis of the technical solution of the present invention without any creative work are still within the scope of protection of the present invention.

Claims

1. A medical image feature map denoising method based on a diffusion model, characterized in that: The following steps are involved: Extract the feature map from the model and use the extracted feature map as the original feature map; Add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps. Specifically, Gaussian noise is gradually added to the original feature map extracted from the model. In the figure, the noise level of the original feature map increases from 0 to T, that is, from no noise to maximum noise, and a series of noise feature maps are obtained. , ,……, ; For the original feature map The process of gradually introducing noise is: I represents the identity matrix, ,……, represents the noise variance at each step in the process; The noise characteristic graph is expressed as: in, ; The noise feature maps finally obtained are sequentially input into multiple FMD modules connected in series for denoising, and denoising results corresponding to a series of noise feature maps are obtained. The noise feature maps and the corresponding denoising results are used to train the FMD modules to obtain a trained FMD module; The denoising process is specifically as follows: in, Represents a parameter involved in the denoising process; is the noise predicted by the FMD module; represents the mean; from predict The formula is: Represents the variance that the model can learn, and z is the random component; The original feature map is input into the trained FMD module for denoising.

2. The medical image feature map denoising method based on a diffusion model according to claim 1, wherein: The FMD module is trained using the mean squared error loss of the noise feature map and the corresponding denoising results.

3. The medical image feature map denoising method based on a diffusion model according to claim 2, wherein: The training loss function of the FMD module is: 。 4. A medical image feature map denoising system based on a diffusion model, characterized by: include: The acquisition module is configured to: extract the feature map from the model and use the extracted feature map as the original feature map; The noise adding module is configured to add Gaussian noise to the original feature map in a progressive manner to obtain a series of noise feature maps, specifically: gradually adding Gaussian noise to the original feature map extracted from the model In the figure, the noise level of the original feature map increases from 0 to T, that is, from no noise to maximum noise, and a series of noise feature maps are obtained. , ,……, ; For the original feature map The process of gradually introducing noise is: I represents the identity matrix, ,……, represents the noise variance at each step in the process; The noise characteristic graph is expressed as: in, ; The training module is configured to: sequentially input the final noise feature map into multiple FMD modules connected in series for denoising, obtain denoising results corresponding to the series of noise feature maps, and train the FMD module using the noise feature map and the corresponding denoising results to obtain a trained FMD module; the denoising process is specifically as follows: in, Represents a parameter involved in the denoising process; is the noise predicted by the FMD module; represents the mean; from predict The formula is: Represents the variance that the model can learn, and z is the random component; The denoising module is configured to input the original feature map into the trained FMD module for denoising.

5. A computer-readable storage medium having a program stored thereon, characterized in that: When the program is executed by a processor, the steps of the medical image feature map denoising method based on a diffusion model as described in any one of claims 1 to 3 are implemented.

6. An electronic device comprising a memory, a processor, and a program stored in the memory and executable on the processor, wherein: When the processor executes the program, the steps in the medical image feature map denoising method based on the diffusion model are implemented as described in any one of claims 1 to 3.

Images