A medical image segmentation method based on a two-path self-supervised branch

Abstract

The application discloses a medical image segmentation method based on a two-path self-supervised branch, which comprises the following steps: providing a sample image sequence; network training; feeding labeled data into a first deep learning network for training, and calculating a first loss function of the first deep learning network according to labels of the labeled data; feeding unlabeled data into the first deep learning network to output a first feature map, and feeding the unlabeled data into a second deep learning network to output a second feature map; calculating pseudo labels according to the first feature map and the second feature map; calculating a second loss function of the first deep learning network and the second deep learning network according to the pseudo labels; and iteratively training until the weighted sum of the first loss function and the second loss function is less than a preset value. The technical scheme provided by the application alleviates the dependence of the existing deep learning neural network training on labeled data, and improves the training efficiency of the deep learning neural network.

Description

Technical Field

[0001] This invention relates to the field of digital image processing technology, and in particular to a medical image segmentation method based on dual-path self-supervised branching. Background Technology

[0002] In existing technologies, deep learning-based medical image segmentation methods, such as UNet and UNet++, have achieved good segmentation results. These methods all utilize deep learning neural networks trained using supervised training strategies, requiring a large amount of labeled data for training. The quality of the labeled data also affects the model's performance.

[0003] In the field of 3D medical imaging applications, the annotation of medical image data requires professional physicians, which is costly and time-consuming. Unannotated medical images, on the other hand, lack this necessary process. Therefore, mining general and robust features from unannotated medical images and using them to train deep learning neural networks can alleviate the dependence of segmentation methods on labeled data in 3D medical imaging applications. Summary of the Invention

[0004] (I) Purpose of the Invention

[0005] The purpose of this application is to provide a medical image segmentation method based on dual-path self-supervised branches, which uses labeled and unlabeled data for network training, thereby alleviating the dependence of existing technologies on labeled data and improving the training efficiency of the network.

[0006] (II) Technical Solution

[0007] According to some embodiments, the present invention provides a medical image segmentation method based on dual-path self-supervised branching, comprising: providing a sample image sequence, the sample image sequence containing labeled data and unlabeled data; network training, feeding labeled data into a first deep learning network for training, calculating a first loss function of the first deep learning network based on the labels of the labeled data; feeding unlabeled data into the first deep learning network and outputting a first feature map, feeding unlabeled data into a second deep learning network and outputting a second feature map; calculating pseudo-labels based on the first and second feature maps; calculating a second loss function of the first and second deep learning networks based on the pseudo-labels; updating the parameters of the first deep learning network based on the gradient of the first deep learning network; updating the parameters of the second deep learning network based on the gradient of the second deep learning network and the parameters of the first deep learning network; iterative training until the weighted sum of the first and second loss functions is less than a preset value; and image segmentation, feeding the medical image into the first or second deep learning network to obtain the segmentation result of the medical image.

[0008] In one embodiment, prior to network training, the method further includes at least one of the following operations: random cropping, random discarding, and shuffling of grayscale values ​​of the sample data.

[0009] In one embodiment, the step of calculating pseudo-labels based on the first feature map and the second feature map includes: calculating the average value of each pixel after adding the first feature map and the second feature map; the category with the highest probability in the average value is the pseudo-label of the pixel.

[0010] In one embodiment, the step of calculating the second loss function of the first deep learning network and the second deep learning network includes: calculating the feature distribution P of each pixel in the first feature map and the second feature map; calculating the feature distribution Q of the average value of each pixel after adding the first feature map and the second feature map; and using the Kullback-Leibler divergence between the feature distribution P and the feature distribution Q as the second loss function.

[0011] In one embodiment, the first deep learning network and the second deep learning network are 2D U-Net or 3DUNETR.

[0012] (III) Beneficial Effects

[0013] The above-mentioned technical solution of the present invention has the following beneficial technical effects: through self-supervised learning implemented by the first deep learning network and the second deep learning network, unlabeled data is fully utilized, which alleviates the dependence of the training of existing deep learning neural networks for medical image segmentation on labeled data and improves the training efficiency of deep learning neural networks.

[0014] It should be understood that the above general description and the following detailed description are exemplary and explanatory only, and are not intended to limit this disclosure. Attached Figure Description

[0015] To more clearly illustrate the specific embodiments of the present invention or the technical solutions in the prior art, the drawings used in the description of the specific embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are one embodiment of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0016] Figure 1 This is a schematic diagram of a medical image segmentation method based on dual-path self-supervised branching according to an embodiment of the present invention;

[0017] Figure 2 This is a schematic diagram of a dual-path self-supervised branch network according to an embodiment of the present invention. Detailed Implementation

[0018] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to specific embodiments and the accompanying drawings. It should be understood that these descriptions are merely exemplary and not intended to limit the scope of the invention. Furthermore, descriptions of well-known structures and techniques are omitted in the following description to avoid unnecessarily obscuring the concept of the invention.

[0019] In the description of this invention, it should be noted that the terms "first," "second," and "third" are used for descriptive purposes only and should not be construed as indicating or implying relative importance.

[0020] Figure 1 This is a schematic diagram of a medical image segmentation method based on dual-path self-supervised branching according to an embodiment of the present invention.

[0021] like Figure 1 As shown, this method embodiment includes the following three steps.

[0022] Step S1: Provide sample image sequences. Specifically, sample image sequences are provided before network training. These sequences contain both labeled and unlabeled data. It should be noted that the sample image sequences are clinical patient medical images. The labeled data is data annotated by physicians or professionals in related fields, primarily marking the FOI (Feature of Interest) in the medical images; the unlabeled data consists of medical images without FOI annotation.

[0023] In one instance, prior to network training, at least one of the following operations is performed on the sample data: random cropping, random discarding, and shuffling of grayscale values, thereby improving the generalization ability of the network training.

[0024] Step S2: Feed the sample sequences into the network for training. Specifically, labeled data is fed into the first deep learning network for training, and the first loss function of the first deep learning network is calculated based on the labels of the labeled data; unlabeled data is fed into the first deep learning network, and the first feature map is output; unlabeled data is fed into the second deep learning network, and the second feature map is output; pseudo-labels are calculated based on the first and second feature maps; based on the pseudo-labels, the second loss functions of the first and second deep learning networks are calculated; the parameters of the first deep learning network are updated based on the gradient of the first deep learning network; the parameters of the second deep learning network are updated based on the gradient of the second deep learning network and the parameters of the first deep learning network; training is iterated until the weighted sum of the first and second loss functions is less than a preset value.

[0025] Figure 2 This is a schematic diagram of a dual-path self-supervised branch network according to an embodiment of the present invention.

[0026] refer to Figure 2 F θ1 For the parameters of the first deep learning network, F θ2 These are the parameters of the second deep learning network. The first and second deep learning networks have the same structure, but their initialization parameters differ. Both networks can have their initialization parameters generated using random numbers.

[0027] like Figure 2 As shown, in one branch, labeled data is fed into the first deep learning network for training, generating a feature map P0. Based on the feature map P0 and the labels Y of the labeled data, the first loss function of the first deep learning network is calculated. The first loss function can be the Dice loss function or the cross-entropy loss function. Figure 2 In another branch of the network shown, unlabeled data is fed into the first deep learning network, which outputs the first feature map P1, and unlabeled data is fed into the second deep learning network, which outputs the second feature map P2. The pseudo-label Ym is calculated based on the first feature map P1 and the second feature map P2. The second loss function of the first deep learning network and the second deep learning network is calculated based on the pseudo-label Ym.

[0028] It should be noted that in the network branch trained on unlabeled data, a cross-pseudo-label strategy can be used for supervised training, that is, the pseudo-label corresponding to the second feature map P2 supervises the first feature map P1, and the pseudo-label corresponding to the first feature map P1 supervises the second feature map P2.

[0029] In a preferred embodiment, a hybrid pseudo-label strategy is employed for network training. First, pseudo-labels are calculated based on the feature maps output by the network. For example, the step of calculating pseudo-labels based on a first feature map and a second feature map includes: calculating the average value of each pixel after adding the first and second feature maps; the class with the highest probability in the average value is the pseudo-label of the pixel. The hybrid pseudo-label strategy can alleviate the technical problem of different optimization directions arising from cross-pseudo-label supervision strategies. It should be noted that each pixel is a high-dimensional vector; adding the first and second feature maps involves adding the high-dimensional vectors corresponding to each pixel, and the average value is the average of the two high-dimensional vectors.

[0030] During the network parameter update phase of training, the parameters of the first deep learning network are updated based on the gradient of the first deep learning network; the parameters of the second deep learning network are updated based on the gradient of the second deep learning network and the parameters of the first deep learning network. For example, the parameters of the second deep learning network can be updated based on the gradient of the second deep learning network and a weighted average of the parameters of the first deep learning network.

[0031] In one embodiment, the step of calculating the second loss function of the first deep learning network and the second deep learning network includes: calculating the feature distribution P of each pixel in the first feature map and the second feature map; calculating the feature distribution Q of the average value of each pixel after adding the first feature map and the second feature map; and using the Kullback-Leibler divergence between the feature distribution P and the feature distribution Q as the second loss function.

[0032] It should be noted that this implementation constructs the true distribution using the feature maps of the first and second deep learning networks, thereby constraining network training. Specifically, after the unlabeled data in the sample sequence undergoes forward propagation through the network, feature maps of different resolutions are generated on the decoders of the first and second deep learning networks. Each pixel in the feature map encodes the semantic information of a certain region through a high-dimensional feature vector. Each dimension of this high-dimensional feature vector can be regarded as a random vector. The feature distribution can be obtained by taking the softmax of all high-dimensional feature vectors, i.e., calculating the feature distribution P of each pixel in the first and second feature maps.

[0033] Subsequently, the average value of the feature distributions at corresponding positions of the feature maps output by the first and second deep learning networks is taken as the true distribution at that position. That is, the feature distribution Q of the average value of each pixel after adding the first and second feature maps is calculated.

[0034] The distance between the feature distributions of the decoders in the first and second deep learning networks and the feature distributions of Q is measured by the Kullback-Leibler divergence between feature distributions P and Q. This allows for constrained training using a second loss function to make P closer to Q.

[0035] For example, the formula for calculating the Kullback-Leibler divergence between characteristic distribution P and characteristic distribution Q is: Where KL represents the Kullback-Leibler divergence, and i represents the index of a pixel in the feature map.

[0036] In one embodiment, the first deep learning network and the second deep learning network are 2D U-Net or 3DUNETR.

[0037] Step S3: Feed the medical images into the network for segmentation. Specifically, feed the medical images into a first deep learning network or a second deep learning network to obtain the segmentation results of the medical images.

[0038] It should be understood that the specific embodiments described above are merely illustrative or explanatory of the principles of the invention and do not constitute a limitation thereof. Therefore, any modifications, equivalent substitutions, improvements, etc., made without departing from the spirit and scope of the invention should be included within the protection scope of the invention. Furthermore, the appended claims are intended to cover all variations and modifications falling within the scope and boundaries of the appended claims, or equivalent forms of such scope and boundaries.

[0039] Those skilled in the art will understand that embodiments of the present invention can be provided as methods, systems, or computer program products. Therefore, the present invention can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, the present invention can take the form of a computer program product embodied on one or more computer-usable storage media (including, but not limited to, disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0040] The steps in the method of this invention can be adjusted, combined, or deleted according to actual needs. Similarly, the modules in the system device of this invention can be combined, divided, or deleted according to actual needs.

Claims

1. A medical image segmentation method based on dual-path self-supervised branching, characterized in that, include: Provide a sample image sequence, which contains labeled data and unlabeled data; Network training, The labeled data is fed into the first deep learning network for training, and the first loss function of the first deep learning network is calculated based on the labels of the labeled data. The unlabeled data is fed into a first deep learning network to output a first feature map, and the unlabeled data is fed into a second deep learning network to output a second feature map. Calculate pseudo-labels based on the first feature map and the second feature map; Based on the pseudo-labels, calculate the second loss function of the first deep learning network and the second deep learning network; Update the parameters of the first deep learning network based on the gradient of the first deep learning network. Update the parameters of the second deep learning network based on the gradient of the second deep learning network and the parameters of the first deep learning network. Iterate through the training until the weighted sum of the first and second loss functions is less than the preset value; Image segmentation, The medical images are fed into the first deep learning network or the second deep learning network to obtain the segmentation results of the medical images.

2. The method according to claim 1, characterized in that, The steps for calculating pseudo-labels based on the first feature map and the second feature map include: Calculate the average value of each pixel after adding the first feature map and the second feature map; The category with the highest probability in the average value is the pseudo-label of the pixel.

3. The method according to claim 1, characterized in that, The steps for calculating the second loss function of the first deep learning network and the second deep learning network include: Calculate the feature distribution P of each pixel in the first feature map and the second feature map; Calculate the feature distribution Q of the average value of each pixel after adding the first feature map and the second feature map; The Kullback-Leibler divergence between the feature distributions P and Q is the second loss function.

4. The method according to any one of claims 1-3, characterized in that, The first deep learning network and the second deep learning network are 2D U-Net or 3D UNETR.

Images