A semi-supervised retinal segmentation method for OCT images based on uncertainty

Abstract

The present invention discloses a semi-supervised OCT image retinal segmentation method based on uncertainty. For labeled data, the feature map obtained by the first network is extracted to obtain an attention map, and the input image is subjected to attention enhancement, and then the attention is sent to the first network and the second network again for training; for unlabeled data, in order to reduce the error of pseudo-labels, cross-supervision of the first network and the second network is adopted, supplemented by uncertainty guidance; multi-scale features and uncertainty maps are obtained by the first network, and sent together to the uncertainty repair module to obtain the repaired pseudo-labels and then supervise the output of the first network to calculate the weighted uncertainty repair loss; after each round of iteration, the results are tested on the test set, and when the current Dice coefficient is higher than the result of the previous round of iteration, the model parameters are saved. The present invention greatly improves the generalization of the model, solves the problem of pseudo-label misleading, and improves the boundary segmentation accuracy of the pseudo-label.

Description

Technical Field

[0001] The present invention belongs to the field of image segmentation, and in particular relates to an uncertainty-based semi-supervised OCT image retinal segmentation method. Background Art

[0002] Optical coherence tomography (OCT) is a non-invasive, real-time, micro-resolution, non-contact imaging technology widely used for high-resolution retinal imaging. Many ophthalmic and other diseases (such as myopia, diabetic retinopathy, and AMD) can cause changes in the retinal structure. Automatic retinal segmentation techniques can better segment the retina from OCT images, allowing for better visualization of the features of the relevant regions and assisting in identifying the causes of these changes. Therefore, automated retinal segmentation in OCT images has significant application and research value.

[0003] Traditional retinal segmentation methods are mostly based on mathematical models and supervised learning. The former has very complex constraints and multiple iterations are time-consuming. With the rise of deep learning technology in recent years, convolutional neural networks have been proven to have powerful image segmentation capabilities. However, despite this, the development of automatic retinal segmentation technology is still immature, mainly due to the following defects and shortcomings: (1) The retinal boundary is easily affected by lesions, such as the retinal structure will be broken when leakage occurs; (2) Retinal OCT images are easily affected by speckle noise, the image resolution is blurred, and it is also affected by the acquisition machine, making it difficult to accurately segment the image; (3) Retinal OCT image data is scarce and the cost of manual annotation is high. Summary of the Invention

[0004] In view of the shortcomings of the existing technology, the present invention proposes a semi-supervised OCT image retinal segmentation method based on uncertainty.

[0005] To achieve the above object, the present invention comprises the following steps:

[0006] S1. The OCT image dataset is divided into a training set and a test set, where the training set consists of 10% labeled images and 90% unlabeled images. The training and test set images are preprocessed by cropping, rescaling, and unifying the data size. The training set is then fed into the model for training.

[0007] S2. For labeled data, the feature map obtained by the first network is extracted to obtain the attention map, the input image is focused, and then the attention map is sent to the first network and the second network for training again.

[0008] S3. For unlabeled data, in order to reduce the error of pseudo-labeling, cross-supervision of the first network and the second network is adopted, supplemented by uncertainty guidance.

[0009] S4. The multi-scale features and uncertainty maps are obtained through the first network and sent to the uncertainty repair module together. The repaired pseudo-labels are then used to supervise the output of the first network and calculate the weighted uncertainty repair loss.

[0010] S5. After each iteration, the results are tested on the test set to obtain the segmented OCT image. The test indicator is the Dice coefficient. When the current Dice coefficient is higher than the result of the previous iteration.

[0011] The first network is a U-net network model, and the second network is a Swin-Unet network model.

[0012] Furthermore, in the above-mentioned semi-supervised OCT image retinal segmentation method based on uncertainty guidance and restoration, step S2 specifically includes the following steps:

[0013] S21: Input the labeled data into the first network and the second network, and then calculate the first supervision loss with the corresponding labels using the Dice function.

[0014] S22: Select the feature map obtained after four downsamplings of the first network, use dimensionality reduction operation to generate multiple attention maps, and then randomly select an attention map.

[0015] S23: On the attention map selected in step S22, randomly generate multiple square black occlusion areas and superimpose them on the original image to obtain an enhanced image.

[0016] S24: The enhanced image obtained by S23 is input into the first network and the second network, and then the second supervised loss is calculated using the Dice function with the corresponding label. The first supervised loss and the second supervised loss are combined to obtain the total supervised loss function.

[0017] The first supervision loss and the second supervision loss calculated in step S2 can, on the one hand, ensure that the first network and the second network learn correct knowledge from the labels, and on the other hand, improve the robustness of the first network and the second network.

[0018] Furthermore, in the above-mentioned semi-supervised OCT image retinal segmentation method based on uncertainty guidance and restoration, step S3 specifically includes the following steps:

[0019] S31: Input the unlabeled data into the first network and the second network to obtain an image segmentation probability map, and then binarize the image segmentation probability map based on a threshold of 0.5 to obtain a corresponding image segmentation mask.

[0020] S32: The first network and the second network enable the dropout function, and perform random forward prediction multiple times to obtain multiple image segmentation probability maps and calculate the average probability map. Then, the information entropy calculation formula is used to calculate the corresponding uncertainty map.

[0021] S33: The uncertainty map obtained in S32 is combined with the MSE function to guide the calculation of the weighted uncertainty loss between the segmentation mask of the second network and the image segmentation probability map of the first network, and the segmentation mask of the first network and the image segmentation probability map of the second network.

[0022] The weighted uncertainty loss calculated in step S3 enables the first network and the second network to calculate the loss based on each other's segmentation probability map and segmentation mask, thereby better avoiding learning erroneous information and improving the utilization of unlabeled data.

[0023] Furthermore, in the above-mentioned semi-supervised OCT image retinal segmentation method based on uncertainty guidance and restoration, step S4 specifically includes the following steps:

[0024] S41: The decoder stage of the first network. Each stage is sampled to the same dimension as the next stage on the deconvolution network and then spliced, and so on to obtain a multi-scale feature map.

[0025] S42: The uncertainty map of the first network and the multi-scale feature map obtained by S41 are input into the uncertainty repair module to obtain a repaired pseudo-label with clearer boundaries.

[0026] S43: Use the Dice function to calculate the repaired pseudo-label obtained in S42, supervise the segmentation probability map obtained in step S31, and calculate the weighted uncertainty repair loss.

[0027] The weighted uncertainty repair loss calculated in step S4 can further improve the learning effect of the first network for unlabeled data.

[0028] Furthermore, the above-mentioned semi-supervised OCT image retinal segmentation method based on uncertainty guidance and repair, wherein: in the step S4, the uncertainty repair module is composed of three layers of convolutional layers and pooling layers of the same structure, and a fourth layer consisting only of convolutional layers. The convolutional layer is composed of 3x3 convolution kernels, and the pooling layer is composed of 3x3 kernels. In addition, a residual structure is adopted. Except for the first layer, the input of each layer is the concatenation of the output and input of the previous layer. The input of each layer passes through a 3x3 convolutional layer, a normalization layer and a Relu activation function to extract image features, and then the image is compressed through a 3x3 maximum pooling layer. After repeated operations, an image segmentation mask with clearer boundaries is obtained.

[0029] Beneficial effects of the present invention: First, the present invention takes into account the noise problem of the OCT image itself and the influence of the retinal lesion area, and proposes an attention-based image enhancement method, which can greatly improve the generalization of the model; secondly, due to the limitations of semi-supervised learning itself, considering the uncertainty of the model itself, the proposed uncertainty guidance can well solve the problem of pseudo-label misleading; finally, the uncertainty repair module can improve the boundary segmentation accuracy of the pseudo-label, and better enable the model to pay attention to the boundary area. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 Schematic diagram of the overall network structure;

[0031] Figure 2 This is a schematic diagram of the uncertain repair module structure;

[0032] Figure 3 This is a comparison chart of the results of the first network after the uncertainty repair module;

[0033] Figure 4 This is the result image after attention image enhancement;

[0034] Figure 5 The segmentation results of the baseline model and the present invention are shown in FIG. DETAILED DESCRIPTION

[0035] The technical solution of the present invention is further described in detail below with reference to the accompanying drawings:

[0036] Previously, a semi-supervised learning framework (CTCT) used CNN / Transformer cross-teaching as implicit consistency regularization, which can produce more stable and accurate pseudo-labels than explicit consistency regularization. This framework benefits from two different learning paradigms. CNN focuses on local information, and Transformers models long-range relationships. Therefore, cross-teaching helps to simultaneously learn a unified segmenter with these two properties, and has also been proven effective in some medical image segmentation. However, based on the data specificity of this task and the fact that the original network did not consider the impact of network uncertainty on pseudo-labels, this paper uses this as a baseline model to propose a semi-supervised OCT image retinal segmentation method based on uncertainty, which can better improve the accuracy of retinal segmentation.

[0037] The embodiment of the present invention includes the following steps:

[0038] S1, OCT image dataset is divided into training set and test set in a ratio of 7:3, where the training set consists of 10% labeled images and 90% unlabeled images. Figure 1Where L is a labeled image and U is an unlabeled image. The training and test set images are preprocessed by cropping, adjusting the scale, and unifying the data size. The training set is then fed into the model for training.

[0039] Since the original collected OCT images have different sizes, the left half of the OCT image is a scan line diagram, and the right half is the required image, which cannot be directly input into the network for training. First, the fundus image part is cropped, and then cropped and scaled to obtain 224×224 image data.

[0040] S2. For labeled data, the attention map obtained by the first network will enhance the attention of the image to obtain the image La, which is then sent to the network again for training.

[0041] Semi-supervised training includes labeled data and unlabeled data. In this step, the module only acts on labeled data. The labeled data is input into the first network, namely the U-net network model. After four convolution downsamplings, a compressed feature map can be obtained. The multi-channel feature map contains the information that the network pays attention to. Then, after the feature map undergoes 1x1 convolution, multiple corresponding attention maps can be obtained. Then, an attention map is selected according to the ratio of the mean of each attention map to the sum of the mean of all attention maps as the sampling probability. Then, six square blocks of the same size are randomly selected on the obtained attention map, and their pixel values ​​are set to 0. Multiplying the original image with the attention enhancement map will result in an enhanced map with multiple black areas, as shown in the following example. Figure 4 As shown, this image is sent to the first network and the second network again. The loss function used is the Dice function. The supervised loss function is the sum of the first supervised loss and the second supervised loss. The expression is as follows:

[0042] L sup =L sup1 +L sup2 +λ1(L att1 +L att2 )

[0043] Among them, L sup1 , L sup2 They are Figure 1 The loss loss calculated by the unenhanced labeled image L and the label GT of the first and second networks is the first supervised loss; L att1 , L att2 They are Figure 1 The enhanced images La of the first and second networks and the label GT calculate the loss loss, and the sum of the two is the second supervision loss; λ1 is the proportional weight.

[0044] This method performs corresponding image enhancement for the blurred retinal boundaries and faults in OCT images, which can greatly increase the generalization of the model and improve the segmentation accuracy of the model for such situations.

[0045] S3. For unlabeled data, in order to reduce the error of pseudo-labeling, cross-supervision of the first network and the second network is adopted, supplemented by uncertainty guidance.

[0046] This step uses Bayesian deep learning techniques to estimate the uncertainty of the model. The dropout function in existing models can be used to approximate the sampling of model parameters. For each input data, T forward propagations are performed to obtain the prediction results. Each time, random noise is randomly added to the input data. Each pixel has T prediction results, which can be used to obtain the average probability map. The calculation formula is as follows:

[0047]

[0048] Among them, u i is the average predicted probability value of the i-th pixel, t is the t-th prediction, is the predicted probability of the i-th pixel in the t-th prediction, and the corresponding uncertainty map can be further calculated:

[0049]

[0050] The first and second networks calculate their own uncertainty maps U. The uncertainty maps can be used to filter out unreliable predictions. Specifically, the predictions of the first and second networks are used, and the segmentation probability maps are binarized based on a threshold of 0.5 to obtain the corresponding image segmentation masks. The masks obtained by each network are then combined with the uncertainty maps to supervise the output probability maps of the other network. The specific calculation method is as follows:

[0051]

[0052] Among them, U i is the uncertainty value of the i-th pixel, H is the set uncertainty threshold, and returns 1 when the uncertainty value of the corresponding pixel is less than H, otherwise it returns 0. i ′ is the probability of the pixel corresponding to the network prediction result, y i is the value corresponding to the pixel of the segmentation mask of the other network. From this, we can calculate the weighted uncertainty loss with uncertainty constraint, that is, the sum of the weighted uncertainty loss of the first network and the second network, then:

[0053] L unsup =λ2(L un1 +L un2 )

[0054] Among them, L un1 and L un2 for Figure 1 The weighted uncertainty loss is calculated by cross-calculating the output results of the first network and the second network, and λ2 is the proportional weight.

[0055] S4. Obtain multi-scale features and uncertainties through the first network Figure 1 The two are then fed into the uncertainty repair module to obtain the repaired pseudo-labels to supervise the output of the first network.

[0056] The uncertainty repair module uses pixels with high certainty in the network model to repair adjacent uncertain pixels. In the decoder stage of the first network, considering that the multi-scale information of the model can better reflect the information of the image, the feature map of each stage is sampled to the same latitude as the next stage through the deconvolution network and then spliced. The multi-scale feature map that integrates multi-scale information is then obtained. The uncertainty map of the first network is then spliced ​​with the multi-scale feature map and input into the uncertainty repair module. Figure 2 As shown, the uncertainty restoration module consists of three layers of identical convolutional and pooling layers, with the fourth layer consisting solely of a convolutional layer. The convolutional layers are constructed with 3x3 convolution kernels, and the pooling layers are constructed with 3x3 kernels. Furthermore, the module employs a residual structure, where the input of each layer except the first is the concatenation of the previous layer's output and input, resulting in pseudo-labels Ur with clearer boundaries.

[0057] This pseudo label supervises the prediction result probability map output by the first network, further enabling the first network to learn useful information and calculate Figure 1 The weighted uncertainty restoration loss of the probability map of the pseudo label Ur and the prediction result of the first network, that is, λ3L reun , λ3 is the proportional weight, and the loss function is the Dice function.

[0058] S5. After each iteration, the results are tested on the test set to obtain the segmented OCT image. The test indicator is the Dice coefficient. When the current Dice coefficient is higher than the result of the previous iteration, the model parameters are saved.

[0059] In order to verify the beneficial effects of the present invention, Figure 4 The original image in the training phase and the enhanced image after the attention enhancement method can be found. Compared with the original image, most of the black patches appear in the boundary area of ​​the retina, which can better regularize and improve the robustness of the model. Figure 3 The third column is the initial segmentation probability map of the first network, the fourth column is the corresponding uncertainty map, and the fifth column is the segmentation probability map after repair. It can be found that the difference between the initial segmentation probability map and the labeled image is the black area in the uncertainty map. The repaired segmentation probability map eliminates this difference well. Figure 5The segmentation results of the present invention and the benchmark model are shown in FIG. 1 . It can be found that the segmentation results of the present invention are more complete, and the present invention can better segment the corresponding boundaries in the picture with unclear boundaries (e) and the picture with interference caused by lesions (f).

[0060] The embodiments of the present invention described above do not limit the scope of protection of the present invention. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention shall be included in the scope of protection of the claims of the present invention.

Claims

1. A semi-supervised OCT image retinal segmentation method based on uncertainty, characterized by The following steps are involved: S1. Divide the OCT image dataset into training set and test set; The training set and test set images are preprocessed by cropping, adjusting the scale, and unifying the data size, and then the training set is input into the model for training; S2. For labeled data, the feature map obtained by the first network is extracted to obtain the attention map, which is then sent to the first network and the second network for training again; S21, input the labeled data into the first network and the second network, and then calculate the first supervision loss with the corresponding labels using the Dice function; S22. Select the feature map obtained after four downsamplings of the first network, use dimensionality reduction operation to generate multiple attention maps, and then randomly select an attention map; S23. On the attention map selected in step S22, randomly generate multiple square black occlusion areas and superimpose them on the original image to obtain an enhanced image; S24, input the enhanced image into the first network and the second network, and then calculate the second supervision loss with the corresponding label using the Dice function; Combining the first supervised loss and the second supervised loss, we get the total supervised loss function; S3. For unlabeled data, cross-supervision of the first and second networks is adopted, supplemented by uncertainty guidance; S31, inputting unlabeled data into the first network and the second network to obtain an image segmentation probability map, and then binarizing the image segmentation probability map based on a threshold to obtain a corresponding image segmentation mask; S32: The first network and the second network enable the dropout function, and randomly perform forward prediction multiple times to obtain multiple image segmentation probability maps; Calculate the average probability map, and then use the information entropy calculation formula to calculate the corresponding uncertainty map; S33, combining the uncertainty map with the MSE function to guide the image segmentation mask of the second network and the image segmentation probability map of the first network, and the image segmentation mask of the first network and the image segmentation probability map of the second network, to calculate a weighted uncertainty loss; S4. The multi-scale features and uncertainty maps are obtained through the first network and sent to the uncertainty restoration module. The restored pseudo labels are then used to supervise the output of the first network and calculate the weighted uncertainty restoration loss. S41, the decoder stage of the first network, each stage is sampled to the same dimension as the next stage on the deconvolution network and then spliced ​​to obtain a multi-scale feature map; S42: Input the uncertainty map of the first network and the multi-scale feature map into the uncertainty restoration module to obtain a restored pseudo label with a clearer boundary; S43. Use the Dice function to calculate pseudo labels, supervise the segmentation probability map obtained in step S31, and calculate the weighted uncertainty restoration loss; S5. After each iteration, the results are tested on the test set to obtain the segmented OCT image. The test indicator is the Dice coefficient. When the current Dice coefficient is higher than the result of the previous iteration, the model parameters are saved. The first network is a U-net network model, and the second network is a Swin-Unet network model.

2. The uncertainty-based semi-supervised OCT image retinal segmentation method according to claim 1, characterized in that: The training set in step S1 consists of 10% labeled images and 90% unlabeled images.

3. The uncertainty-based semi-supervised OCT image retinal segmentation method according to claim 1, characterized in that: The uncertainty repair module in step S4 is composed of three layers of convolutional layers and pooling layers with the same structure, and a fourth layer consisting only of convolutional layers; The convolution layer is composed of 3x3 convolution kernels, and the pooling layer is composed of 3x3 kernels and adopts a residual structure. Except for the first layer, the input of each layer is the concatenation of the output and input of the previous layer. The input of each layer passes through a 3x3 convolution layer, a normalization layer and a ReLU activation function to extract image features, and then passes through a 3x3 maximum pooling layer to compress the image. After multiple repeated operations, an image segmentation mask with clearer boundaries is obtained.

4. The uncertainty-based semi-supervised OCT image retinal segmentation method according to claim 1, characterized in that: In step S22, an attention map is selected according to the ratio of the mean of each attention map to the sum of the means of all attention maps as the sampling probability.

5. The uncertainty-based semi-supervised OCT image retinal segmentation method according to claim 1, characterized in that: Step S32 performs random forward prediction multiple times, and each time randomly adds random noise to the input data.

Images