Unbalanced medical image segmentation method and system based on self-supervised learning and adaptive weight

Abstract

The invention relates to an unbalanced medical image segmentation method and system based on self-supervised learning and adaptive weight, and the method comprises the steps: dividing a data set, collecting data used for self-supervised pre-training and supervised training, and dividing a three-dimensional medical CT image used for supervised training into a training set and a test set; preprocessing the data set; aiming at the segmentation task, performing targeted self-supervised pre-training, and pre-training using three proxy tasks to obtain a pre-training weight; wherein the three proxy tasks comprise comparative learning, image restoration and rotation; loading the pre-training model to initialize an encoder of a subsequent training network, and training a downstream image segmentation task by using labeled data to obtain a trained model; and after a to-be-segmented medical image is preprocessed, reasoning is performed by using the trained model to obtain a segmented medical image. The method can effectively improve the segmentation precision of the segmentation task and reduce the calculation cost.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image segmentation, and particularly to an unbalanced medical image segmentation method and system based on self-supervised learning and adaptive weights. Background Art

[0002] Medical image segmentation methods play an increasingly important role in the field of clinical medicine. Assisting doctors in diagnosis after segmenting lesions can significantly shorten the diagnosis time. However, in clinical practice, for large-scale and diverse medical images, the labor cost and time cost required for image segmentation are very high, and the accuracy and precision of segmentation are also limited by the professional level of the processing personnel. When medical resources are scarce and time is pressing, the accuracy of manual segmentation is even more limited, which restricts the quality and efficiency of doctors' diagnosis. Therefore, in clinical practice, efficient and reliable medical image segmentation methods are becoming increasingly indispensable.

[0003] Neural network technologies rely on a large number of labeled datasets for fully supervised training. Commonly used segmentation networks include fully convolutional networks, U-Net, U-Net+, etc. U-Net is specifically designed for medical image segmentation tasks. Currently, U-Net variant networks with the U-Net design concept as the core have gradually become the mainstream in the field of medical image segmentation. Deep learning technologies have been widely adopted in the current field of medical image segmentation for the following reasons:

[0004] In terms of content, although there are certain rules among different medical images, there may be significant differences between different cases, between different stages of the same case, or even between different organs of the same case at the same stage. In medical images, position and size are quite important information, and the imbalance of medical images requires sufficient attention to small samples during the training process. Therefore, it is necessary to balance this pair of contradictions to achieve the highest overall segmentation accuracy, which poses new requirements for the feature extraction ability of neural networks.

[0005] In terms of acquisition difficulty, since medical images involve patient privacy, the acquisition and use of data not only require the consent of patients or their families but also comply with relatively strict industry reviews. Therefore, publicly available medical image datasets are not abundant. In addition, the annotation cost of medical images is high. Their diagnosis and region segmentation are quite professional medical problems and can only be annotated by people with rich medical knowledge such as clinical doctors. Therefore, publicly available annotated medical image datasets for research are quite precious and the available data samples are few.

[0006] Neural network technologies all rely on large amounts of labeled datasets for training. However, there are few labeled medical images. Therefore, self-supervised learning has good application prospects in the field of medical images: Self-supervised learning designs the use of pretext tasks to apply unlabeled data to training, greatly increasing the data samples available for training and reducing the difficulty of obtaining effective data. Currently, the mainstream self-supervised application method is staged training. Use unlabeled datasets for self-supervised pre-training, and then, for the target task, fine-tune with labeled data on the basis of the pre-training to finally obtain a neural network model with good segmentation accuracy. Summary of the Invention

[0007] Aiming at the deficiencies of the prior art, the present invention provides a segmentation method for unbalanced medical images to solve the problem of image segmentation with different sizes and large case differences commonly encountered in clinical practice and improve the segmentation accuracy of images under unbalanced samples.

[0008] Based on the SwinUNETR network, the present invention improves the loss function for segmentation application examples, assigns weights to different targets of different sizes according to the volume ratio, and designs a new weight generation function to enable different targets to obtain good segmentation accuracy; Secondly, targeted self-supervised pre-training is carried out for the segmentation task field, enabling the training network to learn the deep features of the data, saving training computing resources, shortening the training time, and improving the accuracy of the network.

[0009] The present invention also provides an unbalanced medical image segmentation system based on self-supervised learning and adaptive weights.

[0010] To achieve the above object, the technical solution of the present invention is as follows:

[0011] An unbalanced medical image segmentation method based on self-supervised learning and adaptive weights, including:

[0012] Divide the dataset, collect data for self-supervised pre-training and supervised training, and divide the three-dimensional medical CT images for supervised training into a training set and a test set;

[0013] Preprocess the dataset, including: performing normalization processing and cutting the data into three-dimensional sub-images of the same size;

[0014] Carry out targeted self-supervised pre-training for the segmentation task. Three pretext tasks are used in the pre-training to obtain pre-training weights; Among them, the three pretext tasks include: contrast learning, image inpainting, and rotation; In contrast learning, the pre-training model is trained by calculating the similarity between image pairs. In image inpainting, the network's feature representation ability is enhanced by restoring some missing image information; In the rotation task, the pre-training model predicts the rotation angle of the input image.

[0015] Load the pre-trained model to initialize the encoder of the subsequent training network, and use the annotated data to train the downstream image segmentation task to obtain the trained model;

[0016] After the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

[0017] According to the preferred embodiment of the present invention, the loss function L of contrastive learning cpntrast is Equation (1):

[0018]

[0019] where x i and x j are a pair of positive and negative samples, t is the measure of the normalized temperature scale, N is the number of negative samples, sim is the similarity function, and exp is the exponential function;

[0020] The loss function L of image inpainting inpaint is Equation (2):

[0021]

[0022] where Y is the original target, and is the reconstructed output after masking;

[0023] The loss function L of rotation rot is Equation (3):

[0024]

[0025] where R represents the rotation angle, represents the rotation category probability after Softmax normalization, and z r is the true value.

[0026] According to the preferred embodiment of the present invention, the pre-trained model includes an encoder and a decoder;

[0027] The encoder includes an image patch partition (Patch Partition) and a Swin Transformer module (Block);

[0028] The Swin Transformer module includes a window self-attention mechanism W-MSA, a shifted window self-attention mechanism SW-MSA, and a merging module. The merging module includes channel connection, layer normalization (LN), and linear transformation;

[0029] Take the preprocessed data as input, perform image patch partitioning, and send it to the Swin Transformer module for downsampling; the image patch partitioning divides the input image into local regions of a fixed size, which are used as the input sequence of the Swin Transformer module. Each image patch is transformed into an embedding vector through linear projection, and positional encoding is added to retain spatial information;

[0030] The Swin Transformer module performs layer normalization on the input data, extracts features through W-MSA, and then sends them to the fully connected layer; subsequently, the data undergoes sliding window self-attention feature extraction via SW-MSA to complete one stage of downsampling. After downsampling, the length, width, and height of the image are all half of the input;

[0031] While the feature map obtained by downsampling is passed to the lower layer for another stage of downsampling feature extraction, it is also skip-connected to the decoder after passing through the residual block;

[0032] The smallest size feature map obtained after multiple stages of downsampling; the smallest size feature map is sent to the decoder for decoding, and the original size is gradually restored through each upsampling; the decoder directly obtains the feature map extracted by the same layer encoder. The decoder splices and fuses this feature map with the output of the lower layer decoder and then outputs it to the upper layer decoder.

[0033] Further preferably, the decoder adds a Squeeze-and-Excitation module, namely the SE module. The SE module obtains the global information of the feature map through global average pooling, generates weights for each channel through the fully connected layer, and scales the feature map through the weights to strengthen the expression of important features.

[0034] Further preferably, the Swin Transformer module adds relative positional encoding and further expands the receptive field. The calculation formula for the self-attention score Attention(Q, K, V) with relative positional encoding is shown in Equation (4):

[0035]

[0036] In Equation (4), Q, K, and V respectively represent the query, key, and value, and B is a learned parameter representing the relative positional bias;

[0037] The encoder is used to extract local features and fuse the extracted local features with the global features extracted by the Swin Transformer module, which are expressed as Equation (5), Equation (6), Equation (7), and Equation (8):

[0038] f1 = ConvBlock12kernel=(3,3,3)9 (5);

[0039] f2 = ConvBlcok22 with kernel=(3,3,3) 9 (6);

[0040] Y1 = LeakyReluEIN2f1(X) 9 H (7);

[0041] Y = LeakyReluEx + IN2f2(Y1) 9 H (8);

[0042] Among them, f1 and f2 are two different convolutional layers, both using a 3×3×3 convolutional kernel. X and Y1 are the outputs of the SwinTransformer module, and Y is the final output of the encoder; kernel refers to the convolutional kernel, ConvBlock refers to the convolutional layer, IN refers to instance normalization, and LeakyRelu is the activation function.

[0043] Further preferably, in the compression stage, the SE module performs global average pooling on the input feature map to obtain global statistical information along the channel dimension, and then uses a fully connected layer and an activation function to model the global statistics, generating weights for each channel, expressed as Equations (9), (10), and (11):

[0044]

[0045] s i = σ2W2δ(W1z i ) 9 (10);

[0046]

[0047] Equation (9) is the compression process, Equation (10) is the excitation process, Equation (11) is the excitation process. The compression process performs global average pooling on each channel to compress the spatial dimension information. The excitation process learns the non-linear relationship between channels through two fully connected layers and the activation function sigmoid; weighting multiplies the weights back to the original feature map channel by channel to obtain the final output; s i is the reweighting parameter, is the output after weighting in Equation (11), z i is the output after compression in Equation (9), x ijNO represents the value of the i-th channel at position (j,k,l), H, W, and D respectively represent the height, width, and depth of the feature map, W1 and W2 are the weight matrices of the fully connected layers, δ(·) represents the activation function Relu, and σ(·) represents the sigmoid function.

[0048] According to the preference of the present invention, the loss function L of the pre-trained model is expressed as Equations (12) and (13):

[0049]

[0050] w ` = Softmax(log 1i (1 / σ ` )) (13);

[0051] where X i and Y i represent the predicted probability and the true probability for each class i, I is the total number of classes, and σ ` is the total prime number.

[0052] A computer device includes a memory and a processor. The memory stores a computer program, and when the processor executes the computer program, the steps of an unbalanced medical image segmentation method based on self-supervised learning and adaptive weights are implemented.

[0053] A computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the steps of an unbalanced medical image segmentation method based on self-supervised learning and adaptive weights are implemented.

[0054] An unbalanced medical image segmentation system based on self-supervised learning and adaptive weights includes:

[0055] A dataset partitioning module, configured to: partition a dataset, collect data for self-supervised pre-training and supervised training, and divide the three-dimensional medical CT images for supervised training into a training set and a test set;

[0056] A preprocessing module, configured to: preprocess the dataset, including: performing normalization processing and cutting the data into three-dimensional sub-images of the same size;

[0057] A self-supervised pre-training module, configured to: perform targeted self-supervised pre-training for the segmentation task. Three proxy tasks are used for pre-training to obtain pre-training weights. Among them, the three proxy tasks include: contrast learning, image inpainting, and rotation. Contrast learning trains the pre-training model by calculating the similarity between image pairs. Image inpainting enhances the feature representation ability of the network by restoring some missing image information. In the rotation task, the pre-training model predicts the rotation angle of the input image;

[0058] A training module, configured to: load the pre-training model to initialize the encoder of the subsequent training network, and use the labeled data to train the downstream image segmentation task to obtain a trained model;

[0059] An image segmentation module, configured to: after the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

[0060] The beneficial effects of the present invention are:

[0061] The present invention proposes a method for directly segmenting original three-dimensional CT images. Based on a self-supervised pre-training and an adaptive-weight loss function, an improved SwinUNETR model is utilized to enhance the ability to segment unbalanced targets, which can be directly applied to medical images used in clinical practice.

[0062] In the pre-training stage of the present invention, multiple proxy tasks are applied to improve the ability to obtain features, and targeted pre-training is designed for the actual tasks faced to improve the training effectiveness; the network structure is improved, and an SE module is added to the network to increase channel attention; a weighted loss function is designed to strengthen the attention to unbalanced targets. This multi-faceted method can effectively improve the segmentation accuracy of the segmentation task and reduce the computational cost. Brief Description of the Drawings

[0063] Figure 1 It is a schematic flowchart of two stages of self-supervised pre-training and segmentation task training of the present invention;

[0064] Figure 2 It is a schematic diagram of the network structure including SwinTransformer layers and SE modules;

[0065] Figure 3 It is a schematic diagram of the weight effect of the loss function with adaptive weights;

[0066] Figure 4 It is a schematic diagram of the encoder extracting local features and feature fusion. Detailed Embodiments

[0067] In order to better implement the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights proposed by the present invention, the following describes the embodiments of the present invention in detail in combination with specific embodiments and the accompanying drawings of the specification. It should be noted that the embodiments of the present invention are not limited to the examples, and any equivalent alternative solutions and improvements should be included within the protection scope of the present invention.

[0068] It should be pointed out that the following detailed description is exemplary and is 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 of ordinary skill in the technical field to which the present invention belongs.

[0069] Embodiment 1

[0070] The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights includes:

[0071] Divide the dataset and collect data for self-supervised pre-training and supervised training, including 20 CT scans of COVID-19 lungs and lesion areas for supervised training of the segmentation task, and the TCIA COVID-19 dataset containing 771 CT scans for self-supervised pre-training. Divide the 3D medical CT images for supervised training into a training set and a test set according to a ratio of 4:1.

[0072] Preprocess the dataset, including: performing normalization processing and cutting the data into 3D sub-images of the same size.

[0073] Conduct targeted self-supervised pre-training for the segmentation task. Three proxy tasks are used in the pre-training to obtain pre-training weights. Among them, the three proxy tasks include: contrastive learning, image inpainting, and rotation. Contrastive learning trains the pre-training model by calculating the similarity between image pairs. The core idea is to calculate the similarity of each pair of positive and negative sample pairs and adjust the convergence direction of the network based on maximizing the similarity between positive samples. Image inpainting enhances the feature representation ability of the network by restoring some missing image information. In the rotation task, the pre-training model predicts the rotation angle of the input image. By predicting the rotation angle of the input image, the network learns the structural features of the image itself. After constructing the pre-training method using the three proxy task methods, the network pre-trains the encoder through this method, and exports the information in the encoder as the pre-training model after the pre-training is completed.

[0074] Load the pre-training model to initialize the encoder of the subsequent training network, and use the labeled data to train the downstream image segmentation task to obtain the trained model.

[0075] After the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

[0076] Example 2

[0077] According to the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights described in Example 1, the difference is that:

[0078] The loss function L of contrastive learning contrast Is Equation (1):

[0079]

[0080] Where x i And x j Are a pair of positive and negative samples, t is a measure of the normalized temperature scale, N is the number of negative samples, sim is the similarity function, and exp is the exponential function;

[0081] The loss function L of image inpainting inpaint Is Equation (2):

[0082]

[0083] Among them, Y is the original target, and is the reconstructed output after masking (M);

[0084] The rotation loss function L rot is Equation (3):

[0085]

[0086] where R represents the rotation angle, represents the rotation category probability after Softmax normalization, and z r is the true value.

[0087] The pre-trained model includes an encoder and a decoder;

[0088] The encoder includes an image patch partition (Patch Partition) and a Swin Transformer module (Block);

[0089] The Swin Transformer module includes a window self-attention mechanism W-MSA, a shifted window self-attention mechanism SW-MSA, and a merging module. The merging module includes channel connection, layer normalization (LN), and linear transformation; the merging module is a hierarchical feature downsampling mechanism used to gradually reduce the spatial resolution and increase the channel dimension at different stages of the model, thereby constructing a hierarchical feature representation. The learnable linear transformation retains more semantic information modules; layer normalization (LN) is used to reduce data offset and stabilize training.

[0090] Taking the preprocessed data as input, perform image patch partitioning and send it to the Swin Transformer module for downsampling; the image patch partitioning divides the input image into local regions of a fixed size (such as 16×16 pixel blocks), which are used as the input sequence of the Swin Transformer module. Each image block is transformed into an embedding vector through linear projection, and position encoding is added to retain spatial information;

[0091] The Swin Transformer module performs layer normalization on the input data. After extracting features through W-MSA, the data is fed into the fully connected layer. The fully connected layer is one of the core components in a neural network, characterized by each neuron being connected to all neurons in the previous layer, forming a dense connection structure. The fully connected layer is often located at the end of the network and is responsible for integrating the local features extracted by the previous layer, transforming them into global features, and performing high-level reasoning. Subsequently, the data undergoes sliding window self-attention feature extraction via SW-MSA, completing downsampling for one stage. After downsampling, the length, width, and height of the image are all half of the input.

[0092] While the feature map obtained from downsampling is passed to the lower layer for another stage of downsampling feature extraction, it is also skip-connected to the decoder through a residual block. Preparation is made for concatenating or feature fusion with the image to be upsampled and restored by the decoder. The residual block is a key structure in a deep neural network, used to address the problems of gradient vanishing and network degradation during the training of deep networks. Its core idea is to directly add the input to the output after passing through several convolutional layers, activation functions, and batch normalization through a skip connection.

[0093] The smallest-sized feature map obtained after multiple stages of downsampling. After each encoder in the input data downsamples, the size becomes half of the original. After the last encoder samples the data, the smallest-sized feature map is obtained. The smallest-sized feature map is fed into the decoder based on a convolutional neural network for decoding and gradually restores the original size through each upsampling. Due to the existence of skip connections, the decoder directly obtains the feature map extracted by the same-layer encoder. The decoder splices and fuses this feature map with the output of the lower-layer decoder and then outputs it to the upper-layer decoder. The main structure of the network consists of an encoder, a decoder, and the skip connections connecting the two. The encoder part of the network consists of five layers, with each layer corresponding to an encoder. The number of decoder layers corresponds to the encoder. That is, there is one encoder and one decoder for each layer. All the encoders combined form the encoder part of the network, and all the decoders combined form the decoder part of the network.

[0094] As Figure 2 shown, the encoder is on the left and the decoder is on the right. The encoder consists of a shifted window Transformer and a Unetr encoder. The detailed structure of the shifted window Transformer layer is shown in the blue dashed box on the far left. On the right is the decoder, which consists of a Unetr decoder and a squeeze-excitation module.

[0095] The decoder adds a Squeeze-and-Excitation module, i.e., the SE module. The SE module obtains the global information of the feature map through global average pooling, generates weights for each channel through a fully connected layer, and scales the feature map by the weights to strengthen the expression of important features.

[0096] When training for the annotated medical image segmentation task, the pre-trained model randomly slices the input three-dimensional data into cube images of 96×96×96 as the data used in training, and then inputs this cube image into the Swin Transformer module after PatchPartition.

[0097] The Swin Transformer module adds relative position encoding on the basis of the traditional Transformer and further expands the receptive field. The calculation formula of the self-attention score Attention(Q, K, V) with relative position encoding is shown in Equation (4):

[0098]

[0099] In Equation (4), Q, K, and V represent the query, key, and value respectively, and B is a learned parameter representing the relative position bias;

[0100] Softmax can normalize a numerical vector into a probability distribution vector, and the sum of all probabilities is 1, which is specifically expressed as follows:

[0101]

[0102] Among them, z is a vector, z i and z j are the elements of z;

[0103] The pre-trained dataset enters the pre-trained network encoder to extract features after data augmentation, and then is transformed into an embedding vector through the mapping head and the loss is calculated through the loss function formula. Subsequently, according to the constraints of the loss function, backpropagation is used to update the parameters of the pre-trained network encoder. After the training is completed, the updated parameter matrix is exported by the encoder as the pre-trained model for downstream tasks.

[0104] There is an encoder based on a convolutional neural network after the Swin Transformer Block to extract local features, and the extracted local features are fused with the global features extracted by the Swin Transformer module, as Figure 4 shown, expressed as Equation (5), Equation (6), Equation (7), Equation (8):

[0105] f1 = ConvBlock12kernel=(3,3,3)9 (5);

[0106] f2 = ConvBlcok22 with kernel=(3,3,3), dilation=(6);

[0107] Y1 = LeakyReluEIN2f1(X), dilation=(7);

[0108] Y = LeakyReluEx + IN2f2(Y1), dilation=(8);

[0109] Among them, f1 and f2 are two different convolutional layers, both using a 3×3×3 convolutional kernel. X and Y1 are the outputs of the SwinTransformer module, and Y is the final output of the CNN-based encoder; kernel refers to the convolutional kernel, ConvBlock refers to the convolutional layer, IN refers to Instance Norm (Instance Normalization), and LeakyRelu is the activation function. The formula is:

[0110]

[0111] α is a small constant, usually taken as 0.01.

[0112] The data passes through the Swin Transformer Block and then is input into the UNETR encoder, and then makes a skip connection with the decoder. After going through downsampling and upsampling recovery, an image matching the input size is obtained.

[0113] At the backend of the decoder, in order to improve the network structure, the present invention introduces an SE module (Squeeze-and-Excitation Block). In the compression stage, the SE module performs global average pooling on the input feature map to obtain global statistical information along the channel dimension, and then uses a fully connected layer and an activation function to model the global statistics, generating weights for each channel, expressed as equations (9), (10), and (11):

[0114]

[0115] s i = σ(2W2δ(W1z i ), dilation=(10);

[0116]

[0117] Equation (9) is the compression process, equation (10) is the excitation process, and equation (11) is the excitation process. The compression process performs global average pooling on each channel to compress the spatial dimension information. The excitation process learns the non-linear relationship between channels through two fully connected layers and the activation function sigmoid; weighting multiplies the weights back to the original feature map channel by channel to obtain the final output; si is the reweighting parameter, is the output after weighting in Equation (11), z i is the output after compression in Equation (9), x ijNO represents the values of the i-th channels at position (j, k, l), H, W, and D respectively represent the height, width, and depth of the feature map, W1 and W2 are the weight matrices of the fully connected layers, δ(·) represents the activation function Relu, and σ(·) represents the sigmoid function.

[0118] The present invention provides a new method for designing the weights of the loss function, which allocates weights to multiple segmentation targets according to the volume relationship and optimizes the weight generation method to adapt to unbalanced medical images, such as Figure 3 shown, the loss function L of the pre-trained model is expressed as Equations (12) and (13):

[0119]

[0120] w ` = Softmax(log 1i (1 / σ ` )) (13);

[0121] Among them, X i and Y i represent the predicted probability and the true probability of each category i, I is the total number of categories, σ ` is the total number of voxels.

[0122] Example 3

[0123] According to the method for segmenting unbalanced medical images based on self-supervised learning and adaptive weights described in Example 1, the difference is:

[0124] Complete the CT medical image segmentation task, and the flowchart is as Figure 1 shown, and the specific implementation steps are as follows:

[0125] Collect the pre-training dataset, which contains 771 COVID-19 case data, and divide the dataset into a training set containing 722 CT images and a test set containing 49 CT images; then collect the training data for the segmentation task, including two groups of three-dimensional medical CT image datasets, namely the abdominal multi-organ CT image dataset and the COVID-19 dataset.

[0126] Before supervised training, the present invention uses self-supervised pre-training to improve the feature learning ability of the network, reduce the demand for training data, and accelerate the convergence speed of the network. The self-supervised pre-training includes three proxy tasks: contrast learning, image inpainting, and rotation prediction.

[0127] Contrastive learning trains the network by calculating the similarity between image pairs, calculating the similarity of each pair of positive and negative sample pairs;

[0128]

[0129] The image inpainting task enhances the feature representation ability of the network by restoring some missing image information;

[0130]

[0131] In the rotation task, the network needs to predict the rotation angle of the input image;

[0132]

[0133] After self-supervised pre-training is completed, the self-supervised training weights are used to initialize the supervised training network. The network enters the supervised training stage. In this stage, annotated 3D medical images are used for training, aiming to enable the network to accurately perform image segmentation.

[0134] This embodiment uses the BTCV dataset as the training data, including 30 CT images and annotations. 24 images are selected as the training set and 6 images are selected as the test set; the COVID-19 dataset is also used for training, including 20 CT images and annotations. 16 images are selected as the training set and 4 images are selected as the test set. This dataset contains medical images of different clinical cases, with inconsistent image sizes and large differences in the size and shape of the target objects. First, this dataset is preprocessed.

[0135] This embodiment is based on SwinUNETR as the backbone network for image segmentation. The Squeeze-and-Excitation module is added to the decoder part of the network to further improve the segmentation performance of the network, and targeted self-supervised pre-training is used to initialize the network well.

[0136] The training process includes:

[0137] The training network obtains data, and each 3D medical CT image is cut into 3D sub-images of the same size, with the size set to 96×96×96 voxels. This can reduce the computational complexity and make the size of each input sample consistent, facilitating network training.

[0138] The network takes the preprocessed data as input, performs Patch Partition, and sends it to the Swin TransformerBlock for downsampling;

[0139] The Swin Transformer Block performs layer normalization on the input data. After extracting features through W-MSA, the data is sent to a fully connected layer. Subsequently, the data undergoes sliding window self-attention feature extraction through SW-MSA to complete downsampling in one stage. After downsampling, the length, width, and height of the image are all half of the input.

[0140] While the feature map obtained from downsampling is passed to the lower layer for another stage of downsampling feature extraction, it is also skip-connected to the decoder through a residual block, preparing to be concatenated or feature-fused with the image upsampled and restored by the decoder.

[0141] The smallest feature map obtained through multiple stages of downsampling is sent to a decoder based on a convolutional neural network for decoding, gradually restoring the original size through each upsampling. The decoder at each stage needs to concatenate and fuse the feature map delivered by the skip connection of the same-level decoder with the feature map obtained from the input of the upsampled lower layer and then output it to the upper decoder.

[0142] In this embodiment, a Squeeze-and-Excitation module is added to the decoder. The SE module obtains the global information of the feature map through global average pooling, generates weights for each channel through a fully connected layer, and scales the feature map with these weights to enhance the expression of important features.

[0143] Specifically, the SE module is introduced into each layer of the decoder in this example, and the influence of the SE module on the performance of the decoder is studied through experiments, as shown in Table 1:

[0144] Table 1

[0145]

[0146] De_0 to De_4 respectively represent the positions of the decoder. The experiment is conducted on the COVID-19 dataset, and the result is the DSC score.

[0147]

[0148] Among them, X i and Y i respectively represent the true value and the predicted value.

[0149] To evaluate the performance of the model proposed in this embodiment, a variety of targeted self-supervised trainings are carried out on the COVID-19 dataset. The training uses CT images used in clinical practice. Except for cropping the lung region and normalizing the data when collecting the dataset, no other processing is performed on the data. The training and pre-training of the experiment are both completed on an RTX 3090 24G.

[0150] The pre-training is iterated 10,000 times, and the parameters are set as Batchsize = 2, learning rate 1×10 -4 , and the optimizer AdamW is selected; for the downstream task training, it is iterated 5,000 times, and the parameters are set as batchsize = 1, learning rate 1×10 -4 , feature size = 48, and the optimizer AdamW is selected.

[0151] Without self-supervised pre-training, only full-supervised training is carried out, and the average precision of each organ and lesion area is 85.0%; using a set of multi-organ general self-supervised training weights provided by the SwinUNETR network as network initialization and then training, the average precision is also 85.0%;

[0152] Using a targeted unlabeled dataset of lungs and COVID-19, three proxy tasks are adopted and pre-trained with the same weights. After training, the average precision is 86.7%. At this time, the loss function of the pre-trained network is:

[0153] L totaO = L contrast + L inpaint + L rot ;

[0154] Using two prediction-based proxy tasks, namely the image inpainting task and the rotation task, for pre-training, the average precision after training is 87.1%. At this time, the loss function of the pre-trained network is:

[0155] L totaO = L inpaint + L rot ;

[0156] Table 2 shows the results of the self-supervised experiment;

[0157] Table 2

[0158]

[0159] Example 4

[0160] A computer device includes a memory and a processor. The memory stores a computer program, and when the processor executes the computer program, it implements the steps of the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights described in Example 1 or 2.

[0161] Example 5

[0162] A computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, it implements the steps of the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights described in Example 1 or 2.

[0163] Example 6

[0164] An imbalanced medical image segmentation system based on self-supervised learning and adaptive weights, comprising:

[0165] A dataset division module, configured to: divide the dataset, collect data for self-supervised pre-training and supervised training, and divide the three-dimensional medical CT images for supervised training into a training set and a test set;

[0166] A preprocessing module, configured to: preprocess the dataset, including: performing normalization processing and cutting the data into three-dimensional sub-images of the same size;

[0167] A self-supervised pre-training module, configured to: perform targeted self-supervised pre-training for the segmentation task, and three proxy tasks are used in the pre-training to obtain pre-training weights; among them, the three proxy tasks include: contrast learning, image inpainting, and rotation; contrast learning trains the pre-training model by calculating the similarity between image pairs, and image inpainting enhances the feature representation ability of the network by restoring some missing image information; in the rotation task, the pre-training model predicts the rotation angle of the input image;

[0168] A training module, configured to: load the pre-training model to initialize the encoder of the subsequent training network, and use the labeled data to train the downstream image segmentation task to obtain a trained model;

[0169] An image segmentation module, configured to: after the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

Claims

1. An unbalanced medical image segmentation method based on self-supervised learning and adaptive weights, characterized in that, Including: Partition the dataset, collect data for self-supervised pre-training and supervised training, and divide the 3D medical CT images for supervised training into a training set and a test set; Preprocess the dataset, including: performing normalization processing and cutting the data into 3D sub-images of the same size; Conduct targeted self-supervised pre-training for the segmentation task. Three proxy tasks are used in the pre-training to obtain pre-training weights. Among them, the three proxy tasks include: contrast learning, image inpainting, and rotation. Contrast learning trains the pre-training model by calculating the similarity between image pairs. Image inpainting enhances the feature representation ability of the network by restoring some missing image information. In the rotation task, the pre-training model predicts the rotation angle of the input image; Load the pre-training model to initialize the encoder of the subsequent training network, and use the labeled data to train the downstream image segmentation task to obtain the trained model; After the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

2. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to claim 1, wherein The loss function L of contrastive learning contrast is given by Equation (1): where x i and x j are a pair of positive and negative samples, t is a measure of the normalized temperature scale, N is the number of negative samples, sim is a similarity function, and exp is an exponential function; Loss function L for image inpainting inpaint is given by Equation (2): where Y is the original target, and is the reconstructed output after masking; Rotating loss function L rot is given by Equation (3): where R represents the rotation angle, represents the rotation class probability after Softmax normalization, and z r is the true value.

3. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to claim 1, wherein The pre-training model includes an encoder and a decoder; The encoder includes image patch partitioning and Swin Transformer modules; The Swin Transformer module includes a window self-attention mechanism W-MSA, a shifted window self-attention mechanism SW-MSA, and a merging module. The merging module includes channel connection, layer normalization, and linear transformation; Take the preprocessed data as input, perform image patch partitioning, and send it to the Swin Transformer module for downsampling. Image patch partitioning divides the input image into local regions of a fixed size as the input sequence of the Swin Transformer module. Each image patch is transformed into an embedding vector through linear projection, and position encoding is added to retain spatial information; The Swin Transformer module performs layer normalization on the input data, extracts features through W-MSA, and then sends them to the fully connected layer; Subsequently, the data undergoes sliding window self-attention feature extraction through SW-MSA to complete one stage of downsampling. After downsampling, the length, width, and height of the image are all half of the input; While the feature map obtained by downsampling is passed to the lower layer for another stage of downsampling feature extraction, it is skip-connected to the decoder after passing through the residual block; The smallest-sized feature map obtained after multiple stages of downsampling; the smallest-sized feature map is sent to the decoder for decoding, and the original size is gradually restored through each upsampling; The decoder directly obtains the feature map extracted by the encoder of the same layer. The decoder splices and fuses this feature map with the output of the lower-layer decoder and then outputs it to the upper-layer decoder.

4. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to claim 3, wherein The decoder adds a Squeeze-and-Excitation module, i.e., the SE module. The SE module obtains the global information of the feature map through global average pooling, generates weights for each channel through the fully connected layer, and scales the feature map through the weights to strengthen the expression of important features.

5. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to claim 3, wherein The Swin Transformer module adds relative position encoding and further expands the receptive field. The calculation formula of the self-attention score Attention(Q, K, V) with relative position encoding is shown in Equation (4): In Equation (4), Q, K, and V represent the query, key, and value respectively, and B is a learned parameter representing the relative position bias; The encoder is used to extract local features and fuse the extracted local features with the global features extracted by the Swin Transformer module, which are expressed as Equation (5), Equation (6), Equation (7), Equation (8): f1 = ConvBlock12(kernel=(3, 3, 3)) (5); f A = ConvBlock22(kernel=(3, 3, 3))(6); Y1 = LeakyRelu(EIN2(f1(X))) (7); Y = LeakyReluEx + IN2f A (Y1))) (8); Among them, f1 and f A are two different convolutional layers, both using a 3×3×3 convolutional kernel. X and Y1 are the outputs of the Swin Transformer module, and Y is the final output of the encoder; kernel refers to the convolutional kernel, ConvBlock refers to the convolutional layer, IN refers to instance normalization, and LeakyRelu is the activation function.

6. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to claim 3, characterized in that In the compression stage, the SE module performs global average pooling on the input feature map to obtain global statistics along the channel dimension, and then uses a fully connected layer and an activation function to model the global statistics to generate weights for each channel, which are expressed as Equation (9), Equation (10), Equation (11): s i = σ2W A δ(W1z i )) (10); Equation (9) is the compression process, Equation (10) is the excitation process, and Equation (11) is the excitation process. The compression process performs global average pooling on each channel to compress the spatial dimension information, and the excitation process learns the non-linear relationship between channels through two fully connected layers and the activation function sigmoid; Weighting multiplies the weights channel by channel back to the original feature map to obtain the final output; s i is the reweighting parameter, is the output after weighting in Equation (11), z i is the output after compression in Equation (9), x ijNO represents the value of the i channels at the position (j, k, l), and H, W, D represent the height, width, and depth of the feature map respectively, and W1 and W A are the weight matrices of the fully connected layer, δ(·) represents the activation function Relu, and σ(·) represents the sigmoid function.

7. The unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to any one of claims 1-6, characterized in that The loss function L of the pre-trained model is expressed as Equation (12), Equation (13): w ` = Softmax(log 1i (1 / σ ` )) (13); where X i and Y i represent the predicted probability and the true probability for each class i, I being the total number of classes, and σ ` is the overall prime number.

8. A computer device, comprising a memory and a processor, the memory storing a computer program, characterized in that, When the processor executes the computer program, it implements the steps of the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to any one of claims 1-7.

9. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by the processor, it implements the steps of the unbalanced medical image segmentation method based on self-supervised learning and adaptive weights according to any one of claims 1-7.

10. An imbalanced medical image segmentation system based on self-supervised learning and adaptive weights, characterized in that, Including: A dataset division module, configured to: divide the dataset, collect data for self-supervised pre-training and supervised training, and divide the three-dimensional medical CT images for supervised training into a training set and a test set; A preprocessing module, configured to: preprocess the dataset, including: performing normalization processing and cutting the data into three-dimensional sub-images of the same size; A self-supervised pre-training module, configured to: perform targeted self-supervised pre-training for the segmentation task, and obtain pre-training weights using three proxy tasks; among them, the three proxy tasks include: contrast learning, image inpainting, and rotation; contrast learning trains the pre-trained model by calculating the similarity between image pairs, and image inpainting enhances the feature representation ability of the network by restoring some missing image information; in the rotation task, the pre-trained model predicts the rotation angle of the input image; A training module, configured to: load the pre-trained model to initialize the encoder of the subsequent training network, and use the labeled data to train the downstream image segmentation task to obtain the trained model; An image segmentation module, configured to: after the medical image to be segmented is preprocessed, use the trained model for inference to obtain the segmented medical image.

Images