A cervical vertebra MRI image self-supervised segmentation method using a diffusion model to generate data

Abstract

The application discloses a cervical vertebra MRI image self-supervision segmentation method using a diffusion model to generate data. The method comprises image preprocessing such as cervical vertebra MRI slice screening and super-resolution reconstruction based on a diffusion model; and in order to solve the problems of insufficient training data amount and data lack of difference of a deep segmentation model, an improved diffusion model is used to generate high-quality MRI images; finally, a self-supervision method is used to combine real images and generated images, solve the problem of insufficient label amount, construct an Att-UNet network with an Encoder-Decoder structure to extract relevant features of input data, and use multi-scale consistency output measurement loss, so that more accurate vertebra and intervertebral disc segmentation results are obtained. The application combines MRI image generation based on a diffusion model and self-supervision segmentation, effectively expands the sample quantity, improves the model robustness and segmentation precision, and provides more accurate reference information for clinical cervical vertebra disease diagnosis.

Description

Technical Field

[0001] This invention relates to the field of medical image segmentation technology, and in particular to a self-supervised segmentation method for cervical spine MRI images using data generated by a diffusion model. Background Technology

[0002] Deep learning has had a tremendous impact on the medical field, with medical image analysis being one of the most important areas of computer vision. Faced with challenges such as the imbalance between medical supply and demand, the massive increase in clinical imaging data, and the uneven skill levels and resource distribution among radiologists, deep learning is being increasingly widely applied in medical image reconstruction, lesion detection, medical image segmentation, and computer-aided diagnosis.

[0003] In medical images with relatively complex information, techniques for generative modeling using neural networks are gradually developing. Diffusion models, as generative models proposed in the last two years, have received widespread attention due to their strong pattern coverage and high-quality generated samples. Diffusion models have wide applications in medicine, including segmentation, anomaly detection, image conversion, 2D / 3D data generation, reconstruction, denoising, and other medical-related challenges. Diffusion models typically include three general diffusion modeling frameworks: the Diffusion Probability Model (DDPM), the Noise Conditional Score Network (NBS), and stochastic differential equations. The improved denoising probabilistic diffusion model used in this patent for MRI image generation is a diffusion probability model, while the score-MRI diffusion model used for high-resolution reconstruction is a stochastic differential equation model.

[0004] Convolutional neural networks (CNNs) are one of the mainstream methods for medical image segmentation. Among CNNs, UNet and its variants perform well in various medical segmentation tasks. However, CNNs are limited by their local receptive fields, which is not conducive to capturing long-term feature dependencies. To overcome this limitation, this method applies a simple and efficient consistency regularization method to semi-supervised medical image segmentation to achieve better segmentation accuracy and robust generalization ability.

[0005] Unsupervised regularized semi-supervised learning is a method for semi-supervised medical image segmentation. Its basic principle is to train the model by simultaneously inputting labeled and unlabeled data. The encoder is used to extract features from the medical images, while the decoder is used to transform the extracted features. For labeled data, a supervised loss is used to train the network, while for unlabeled data, a pyramid structure is used to generate multi-scale predictions to enforce consistency constraints.

[0006] Current research on cervical vertebral mass segmentation and intervertebral disc segmentation indicates that cervical spine images are relatively scarce and have more complex background information compared to images of the lumbar and thoracic regions, making segmentation more difficult and accuracy harder to control. Classic image segmentation networks such as UNet and its variants struggle to accurately capture information from cervical spine MRI images. Furthermore, key challenges in cervical spine medical image research include limited sample size and diversity, as well as the variability in manual annotation. Summary of the Invention

[0007] To address the above problems, this invention aims to provide a method for generating cervical spine MRI image data and self-supervised segmentation of cervical vertebrae using a diffusion model. This method effectively alleviates the limitations of the number of medical images and the difficulties in annotation, while also significantly improving vertebral body segmentation performance using the generated data. The invention specifically includes the following steps:

[0008] A self-supervised segmentation method for cervical spine MRI images using data generated by a diffusion model includes the following steps:

[0009] S1. Image preprocessing: Annotated and clearly visible slices are extracted from DICOM format cervical MRI data, converted to PNG format, and super-resolution reconstruction is performed using the score-MRI diffusion model;

[0010] S2, Image Generation: The real images preprocessed in S1 are divided into training sets, and their relevant features are extracted as conditional information for image generation. The improved denoising probability diffusion model DDPM is used to generate cervical spine MRI images, and the generated images are selected based on structural similarity.

[0011] S3. Self-supervised segmentation: In the pre-task, the generated image and the unlabeled real image are input into the Encoder feature extraction network of the UNet model for self-supervised training to obtain a pre-trained model for feature extraction. The parameters of the pre-trained model are assigned to the Encoder module of the Att-UNet-URPC model in the downstream task, and the parameters are fine-tuned using labeled real images. Finally, the cervical vertebrae and intervertebral disc regions are segmented through the Decoder module of the Att-UNet-URPC model.

[0012] Preferably, in step S1, extracting the slice refers to extracting clearly visible slices of the cervical spine from a DICOM format medical image file, and some slices contain annotation information. The extracted slices are then converted to PNG format.

[0013] In S1, the main modules for generating diffusion models based on score-MRI include:

[0014] Forward process: By using stochastic differential equations (SDEs), linear Gaussian noise is continuously added to the original image for a total of T1 times, transforming the complex data distribution into a known Gaussian distribution. Simultaneously, a network is built to train a time-based score model.

[0015] The reverse process involves slowly removing noise over a total of T1 times, transforming the Gaussian distribution back to the data distribution, using the score model trained in the forward process to predict the score at the previous time step, and finally obtaining a distribution similar to the original image.

[0016] Preferably, in S2, the main modules of the improved denoising probability diffusion model DDPM include:

[0017] Diffusion process: A nonlinear cosine noise scheme is used instead of a linear noise scheme. Gaussian noise is continuously added to the original image for a total of T2 times, transforming the complex data distribution into a known Gaussian distribution. A 5-layer UNet network is built. An attention mechanism is added to the second layer of the Encoder module of the UNet network, while only the fourth layer of the Decoder module retains the attention mechanism. In summary, an improved Att-UNet network is built to learn the noise distribution.

[0018] Inverse diffusion process: Randomly generate Gaussian noise, use the Att-UNet network trained during the diffusion process to predict the noise from the previous step, and continuously denoise it for a total of T2 times to finally obtain the generated image.

[0019] Apply the exponential moving average (EMA) operation to the obtained model parameters.

[0020] Preferably, in S2, generated images are selected based on structural similarity: based on the strong correlation between the pixels of the images, the structural similarity between each generated image and the real image is calculated, and the average value is taken as the structural similarity measure of the generated image. The top N generated images with the highest average values ​​are taken as part of the segmentation model dataset in S3.

[0021] Preferably, in S3, the main modules of the self-supervised segmentation model include:

[0022] Preliminary task: Using unlabeled cervical spine MRI images, including both real and generated images, input them into a UNet-based Encoder feature extraction network for self-supervised learning to obtain a pre-trained model for feature extraction. Then, the parameters of this model are fed into the Encoder module of the downstream Att-UNet-URPC model.

[0023] Downstream tasks: Import the pre-trained model parameters, initialize the Encoder part of the Att-UNet-URPC network, and feed labeled cervical spine MRI images into the Att-UNet-URPC network for parameter fine-tuning to enable more effective image feature extraction. Finally, segment the cervical vertebrae and intervertebral disc regions using the Decoder module of the Att-UNet-URPC model.

[0024] Preferably, in S3, the training data and training data labels in the pre-task of the self-supervised segmentation model include:

[0025] Using unlabeled real images and generated images, at a scale of approximately 1:10.

[0026] The cutout portion of the cervical spine MRI image is used as a label, that is, a part of the image is randomly deleted, and the remaining part is used to predict the part that was cut out.

[0027] Preferably, in S3, the Att-UNet-URPC network in the downstream task of the self-supervised segmentation model adds a soft-attention mechanism to each layer based on the traditional UNet Decoder; it also constructs multi-scale outputs based on the Encoder-Decoder and performs random dropout on these outputs, while adding noise to the output of the last layer. The constructed loss function combines supervised and unsupervised loss.

[0028] Preferably, in S3, the loss function of the downstream task Att-UNet-URPC network of the self-supervised segmentation model includes:

[0029] The network is trained using supervised loss, including cross-entropy loss and Dice loss, for labeled data; the cross-entropy and Dice cross-entropy are weighted and fused to obtain the segmentation loss for the labeled data.

[0030] For unlabeled data, a pyramid structure is used to generate multi-scale predicted values ​​for consistency constraints. The Euclidean distance between the predicted value at each scale and the average predicted value across all scales is used for consistency constraints, and KL divergence is used to measure the uncertainty between the predicted values. Then, the Euclidean distance and uncertainty are used to construct an uncertainty correction function, and an uncertainty minimization term is introduced to jointly construct the loss function for unlabeled data.

[0031] Compared with the prior art, the beneficial effects of this invention are:

[0032] This invention expands the dataset by generating images, alleviating the problem of difficulty in obtaining medical images and improving the robustness and accuracy of the segmentation model. It provides more accurate reference information for the clinical diagnosis of cervical spine diseases. Attached Figure Description

[0033] Figure 1 This is a schematic diagram of the process of the present invention.

[0034] Figure 2 This is a network structure diagram of the present invention.

[0035] Figure 3 This is a network structure diagram of Att-UNet-URPC for the downstream task of this invention.

[0036] Figure 4 Patient images provided for the implementation of this invention.

[0037] Figure 5 This is a diagram showing the segmentation result of the present invention. Detailed Implementation

[0038] 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 embodiments and accompanying drawings. Here, the illustrative embodiments and their descriptions are used to explain the invention, but are not intended to limit the invention.

[0039] The present invention specifically includes the following steps:

[0040] S1. Image preprocessing: Annotated and clearly visible slices are extracted from DICOM format cervical MRI data, converted to PNG format, and super-resolution reconstruction is performed using the score-MRI diffusion model;

[0041] S2, Image Generation: The real images preprocessed in S1 are divided into training sets, and their relevant features are extracted as conditional information for image generation. The improved denoising probability diffusion model DDPM is used to generate cervical spine MRI images, and the generated images are selected based on structural similarity.

[0042] S3. Self-supervised segmentation: In the pre-task, the generated image and the unlabeled real image are input into the Encoder feature extraction network of the UNet model for self-supervised training to obtain a pre-trained model for feature extraction. The parameters of the pre-trained model are assigned to the Encoder module of the Att-UNet-URPC model in the downstream task, and the parameters are fine-tuned using labeled real images. Finally, the cervical vertebrae and intervertebral disc regions are segmented through the Decoder module of the Att-UNet-URPC model.

[0043] The data processing in S1 mainly includes the following steps:

[0044] Slice selection: Cervical spine MRI images are multimodal, and only in the T2 modality are the intervertebral discs and vertebrae clearly visualized. A cervical spine MRI image in the T2 modality contains dozens of slices. Generally speaking, the two middle slices show the best visualization, and only these clearly visualized slices contain annotation information.

[0045] Super-resolution reconstruction: Selected cervical spine MRI slices are subjected to super-resolution reconstruction at the same resolution to obtain richer texture details.

[0046] The main modules of the diffusion model based on score-MRI include:

[0047] Forward process: Using stochastic differential equations (SDE), Gaussian noise is continuously added to the original image for a total of T1 times to transform the complex data distribution into a known Gaussian distribution and to build a time-based score model for training.

[0048] The reverse process involves slowly removing noise over a total of T1 times, transforming the Gaussian distribution back to the data distribution, and using the score model trained in the forward process to predict the score at the previous time step, thereby generating the image.

[0049] The image generation in S2 mainly includes the following steps:

[0050] Image generation: All cervical spine MRI images are fed into the improved denoising probability diffusion model DDPM, and new cervical spine MRI images are obtained by training with randomly generated Gaussian noise as labels.

[0051] The main modules of the improved denoising probability diffusion model DDPM include:

[0052] Diffusion process: A nonlinear cosine noise scheme is used to replace linear noise for better generation results. Gaussian noise is continuously added to the original image for a total of T1 times. A 5-layer UNet network is constructed. An attention mechanism is added to the second layer of the Encoder module of the UNet network to improve the ability to extract features from the noisy image, while only the fourth layer of the Decoder module retains the attention mechanism to reduce the training time. In summary, an improved Att UNet network is constructed to learn the noise distribution and obtain a more complete MRI cervical spine image.

[0053] Inverse diffusion process: Randomly generate Gaussian noise, use the Att UNet network trained during the diffusion process to predict the noise from the previous step, and continuously denoise it for a total of T2 times, finally obtaining a distribution similar to the original image.

[0054] Exponential moving averages were applied to the obtained model parameters to improve test metrics and increase model robustness.

[0055] Image selection: Gaussian noise is random and cannot guarantee that the quality of the intervertebral discs and vertebral masses in each generated cervical MRI image meets the segmentation requirements. Therefore, structural similarity is used to measure and select the quality of the generated cervical MRI images.

[0056] There is a strong correlation between the pixels in the image. Based on this similarity, high-quality cervical MRI images are selected, the structural similarity between each generated image and the real image is calculated, and the average value is taken as the structural similarity measure of the generated image. The top 280 generated images with the highest average values ​​are taken as part of the segmentation model dataset in S3.

[0057] The self-supervised segmentation in S3 mainly includes the following steps:

[0058] Preliminary task: Using unlabeled cervical spine MRI images, including both real and generated images, input them into a UNet-based Encoder feature extraction network for self-supervised learning. This yields a feature extractor whose parameters are then assigned to the Encoder module of the downstream Att-UNet-URPC model.

[0059] The ratio of the input images to the extractor—unlabeled real cervical MRI images and generated images from an improved denoising probability diffusion model—is approximately 1:10.

[0060] The input image label for the extractor is: a portion of the input image is randomly cut out, and the cut-out portion is used as the label. The remaining portion is then used to predict the removed portion.

[0061] Downstream task: Import the pre-trained model parameters mentioned above, initialize the Encoder part of the Att-UNet-URPC network, and feed the labeled cervical MRI images into the Att-UNet-URPC network for parameter fine-tuning to enable it to extract image features more effectively. Finally, the Decoder module of the Att-UNet-URPC model is used to segment the cervical vertebrae and intervertebral disc regions.

[0062] The training set is fed into the Att-UNet-URPC network. The Encoder network extracts image features from the medical images, and the Decoder transforms the extracted features to generate multi-scale outputs. Random dropout is applied to these outputs, and noise is added to the final layer output. A loss function is constructed that combines supervised and unsupervised loss methods.

[0063] The constructed loss function includes: using supervised loss to train the network for labeled data, and using a pyramid structure to generate multi-scale predictions for unlabeled data to enforce consistency. These two methods are then weighted and fused.

[0064] The supervised loss includes cross-entropy loss and Dice loss. The cross-entropy and Dice cross-entropy are weighted and fused to obtain the segmentation loss for the labeled data. The formula is:

[0065]

[0066] In the formula: L sup Let S represent the segmentation loss function for labeled data, and L represent the output layer number of the Decoder module. dice L represents the Dice loss function. ce p represents the cross-entropy loss function. i Let y represent the output of the s-th layer, and y represent the label.

[0067] The unsupervised loss includes an uncertainty correction function and an uncertainty minimization term, the formula of which is:

[0068]

[0069] In the formula: L unsup This represents the segmentation loss function for unlabeled data. The first term represents the uncertainty minimization term, the second term represents the uncertainty correction function, β represents the weights, S represents the number of output layers of the Decoder module, and D... s p represents the KL divergence between the average-scale output and the outputs at different scales. s p represents the output of the s-th layer of the Decoder module. avg This represents the average value of the output of layer s in the Decoder module. The exponential operation of KL divergence is used to represent the divergence.

[0070] Therefore, the loss function of the Att-UNet-URPC model can be expressed as:

[0071] L total =L sup +μ*L unsup .

[0072] In the formula: L total L represents the loss function of the Att-UNet-URPC model. sup This represents the segmentation loss function for labeled data, where μ represents the weights, and L... unsup This represents the segmentation loss function for unlabeled data.

[0073] In the experimental section, T2 sagittal slice images of the cervical spine were used. After super-resolution reconstruction using score-MRI, these slice images were input into a diffusion model to generate cervical spine MRI images. Then, 780 high-quality cervical spine MRI images were selected from the generated images based on structural similarity. The cervical spine MRI dataset consisted of 400 real images and 780 generated images. In the self-supervised model, the pre-training task used 50 real images and 500 generated images as the pre-training dataset; the downstream task used 280 real images and 280 generated images as the training set for the segmentation model. Of these, 56 real images (10% of the total training set) were labeled, 30 were used as the validation set, and 40 were used as the test set, all fed into the segmentation network. Figure 5 As shown, the URPC-based method is significantly superior to other traditional methods.

[0074] The above description is merely an embodiment of the present invention and is not intended to limit the invention. Various modifications and variations can be made to the present invention by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of the claims of the present invention.

Claims

1. A method for self-supervised segmentation of cervical MRI images using a diffusion model to generate data, characterized in that, Includes the following steps: S1. Image preprocessing: Annotated and clearly visible slices are extracted from DICOM format cervical MRI data, converted to PNG format, and super-resolution reconstruction is performed using the score-MRI diffusion model; The use of the score-MRI diffusion model for super-resolution reconstruction includes: Forward process: Using stochastic differential equations, Gaussian noise is continuously added to the original image for a total of T1 times to transform the complex data distribution into a known Gaussian distribution and to build a time-based score model for training. Reverse process: Slowly remove noise for a total of T1 times, transform the Gaussian distribution back to the data distribution, and use the score model trained in the forward process to predict the score of the previous time step, thereby generating the image; S2, Image Generation: The real images preprocessed in S1 are divided into training sets, and their relevant features are extracted as conditional information for image generation. The improved denoising probability diffusion model DDPM is used to generate cervical spine MRI images, and the generated images are selected based on structural similarity. Among them, the improved denoising probability diffusion model DDPM is used to generate cervical spine MRI images, including: Diffusion process: A nonlinear cosine noise scheme is used instead of a linear noise scheme. Gaussian noise is continuously added to the original image for a total of T2 times until the complex data distribution is transformed into a known Gaussian distribution. At the same time, an attention mechanism is added to the second layer of the Encoder module of the UNet network, while only one layer of attention mechanism is retained in the Decoder module. The noise distribution is learned through the improved Att-Unet network. Inverse diffusion process: Gaussian noise is randomly generated, and the Att-UNet network trained during the diffusion process is used to predict the noise from the previous step. A total of T2 noise reduction operations are performed to obtain the generated image. The images selected based on structural similarity include: For each generated image, calculate its similarity metric with all real images, and take the average as the structural similarity metric for that generated image; The top N generated images with the highest average values ​​are used as part of the segmentation model dataset in S3. S3. Self-supervised segmentation: In the pre-task, the generated image and the unlabeled real image are input into the Encoder feature extraction network of the UNet model for self-supervised training to obtain a pre-trained model for feature extraction. The parameters of the pre-trained model are assigned to the Encoder module of the Att-UNet-URPC model in the downstream task, and the parameters are fine-tuned using labeled real images. Finally, the cervical vertebrae and intervertebral disc regions are segmented through the Decoder module of the Att-UNet-URPC model. The self-supervised training process includes: Pre-task: Use unlabeled cervical MRI images as input to the Encoder part of the UNet model for self-supervised learning to obtain a feature extractor as a pre-trained model for the downstream task; Downstream tasks: Import the parameters of the pre-trained model, initialize the Encoder part of the Att-UNet-URPC network, and feed the labeled cervical MRI images into the Att-UNet-URPC network for feature extraction, so as to perform feature extraction and segmentation of cervical vertebrae and intervertebral disc regions. The unlabeled data used in the pre-task included real cervical spine MRI images and images generated by an improved denoising probability diffusion model. Cutout was used as an image augmentation method in self-supervised learning. The Att-UNet-URPC model adds a soft-attention mechanism to each layer of the traditional UNet Decoder, and constructs multi-scale outputs based on its Encoder-Decoder. Random dropout is applied to these outputs, and noise is added to the output of the last layer to construct the loss function. The loss function consists of supervised loss and unsupervised loss, and its formula is as follows: ； In the formula: This represents the loss function of the Att-UNet-URPC model. This represents the segmentation loss function for labeled data, which is the supervised loss. The segmentation loss function for unlabeled data is the unsupervised loss. Indicates weight; For the segmentation loss function of labeled data, supervised loss is used for network training. The supervised loss is a weighted fusion of cross-entropy loss and Dice similarity coefficient loss, and its formula is: ； In the formula: S represents the number of output layers of the Decoder module. Represents the Dice loss function. Represents the cross-entropy loss function. Let y represent the output of the s-th layer and y represent the label. For the segmentation loss function of unlabeled data, a pyramid structure is used to generate multi-scale predicted values. Consistency constraints and KL divergence are applied to these predicted values ​​to measure the uncertainty between them. Then, an uncertainty correction function is constructed using Euclidean distance and uncertainty minimization terms. Together, these are used to construct the loss function for unlabeled data, with the following formula: ； In the formula: the first term represents the uncertainty minimization term, the second term represents the uncertainty correction function, β represents the weight, and S represents the number of output layers of the Decoder module. This represents the KL divergence between the average-scale output and the outputs at different scales. This represents the output of the s-th layer of the Decoder module. This represents the average value of the output of layer s in the Decoder module. The exponential operation of KL divergence is used to represent the divergence.

Images