Endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance

Abstract

The invention discloses an endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance. The method comprises the steps of obtaining a data set and performing preprocessing; carrying out model design on the cross-branch feature enhancement module and the uncertainty-guided prediction regularization module on the basis of a visual Transform structure segmentation cutting model, and fusing a trainable convolutional neural network encoder branch to obtain an endoscope polyp segmentation model; training the obtained endoscopic polyp segmentation model based on a cross entropy loss function; after a polyp image of a test set is segmented according to the trained endoscopic polyp segmentation model, a Sigmoid function is used for extracting a segmentation result, the performance of the model is verified, and a final endoscopic polyp segmentation model is obtained. The method can enhance the features in a cross-branch manner, thereby improving the diversity of the extracted polyp image features. Meanwhile, an uncertainty-guided prediction regularization module is introduced, so that the uncertainty in the training process is relieved, and the generalization ability of the model is enhanced.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image processing, and in particular to an endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance. Background Art

[0002] According to research reports, the incidence of colorectal cancer ranks third among all cancers, posing a serious threat to public health. As the disease continues to worsen, the survival rate of colorectal cancer patients will be significantly reduced, and the incidence and mortality of colorectal cancer will also rise sharply. In clinical practice, early screening of colorectal cancer is extremely important for the prevention of colorectal cancer. Automatic segmentation of polyps is a key means of assisting the diagnosis of colorectal cancer, which is of great significance for the effective intervention and timely treatment of colorectal cancer.

[0003] Many studies have explored the use of deep learning frameworks for polyp segmentation and achieved remarkable performance. Some methods, such as Polyp-PVT and SSFormer, use pyramid visual transformers as the backbone network for feature extraction and build an encoder-decoder structure for feature extraction; some methods, such as CFANet, design cross-layer feature aggregation networks to incorporate boundary information into the segmentation network. However, due to the limitation of model size, existing methods are still slightly lacking in feature representation and extraction capabilities, especially in dealing with the similarity between polyps and backgrounds and the complexity of polyp boundaries. In addition, the limited size of datasets will also limit the diversity and universality of existing methods. Therefore, how to improve the feature extraction and generalization capabilities of deep learning models, and how to make full use of limited datasets for polyp segmentation, remains an important research issue.

[0004] Recently, the SegmentAnythingModel (SAM) has shown great potential. SAM is pre-trained on the largest segmentation dataset SA-1B, and its significant segmentation ability in various downstream tasks has been demonstrated. This innovative approach, with its large model size and massive data volume, brings a new perspective to the field of polyp segmentation.

[0005] However, due to the domain gap between the training data and endoscopic images, the performance of SAM in the polyp segmentation task is not satisfactory. To address the problem of unsupervised domain adaptation, a series of semantic segmentation methods have been proposed. For example, MIC proposed a mask image consistency module to encourage learning the contextual relationship of the target domain image to enhance the unsupervised domain adaptation ability. In addition, the context-aware domain adaptation method utilizes a new cross-attention mechanism to promote the transferability of contextual information in domain adaptation. However, for the uncertainty of the target domain, further research is still needed. Summary of the invention

[0006] The purpose of the present invention is to overcome the shortcomings of the prior art. To achieve the above purpose, an endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance is adopted to solve the problems raised in the above background technology.

[0007] An endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance comprises the following steps:

[0008] Step S1, obtaining a data set and performing preprocessing;

[0009] Step S2, designing a cross-branch feature enhancement module and an uncertainty-guided prediction regularization module based on the segmentation model of the visual Transformer structure, and integrating a trainable convolutional neural network encoder branch to obtain an endoscopic polyp segmentation model;

[0010] Step S3, training the obtained endoscopic polyp segmentation model based on the cross entropy loss function;

[0011] Step S4: After the polyp images of the test set are segmented according to the trained endoscopic polyp segmentation model, the segmentation results are extracted using the Sigmoid function, and the performance of the model is verified to obtain the final endoscopic polyp segmentation model.

[0012] As a further solution of the present invention: the specific steps in step S1 include:

[0013] Step S11, obtaining a data set, and dividing the data set to obtain a training set and a test set;

[0014] Step S12: perform data enhancement on the training set data, including random rotation, random flipping, and random scaling operations.

[0015] As a further solution of the present invention: the specific steps in step S2 include:

[0016] Step S21, cross-branch feature enhancement module:

[0017] The structure of the segmentation model was modified based on the visual Transformer structure, removing its prompt input and prompt encoding parts, while retaining its visual Transformer encoder and mask decoder parts;

[0018] Introduce a parallel convolutional neural network branch and add an additional multi-head cross-branch attention module to fuse the feature maps F from the convolutional neural network encoder. c and the feature map F of the visual Transformer encoder v, then the CBA formula of the multi-head cross-branch attention module is as follows:

[0019]

[0020] Where Q = F v W Q , K=V=F c W K , d is F v The number of feature dimensions of each head;

[0021] Fusion of deep contextual information and shallow boundary information from the encoder with the decoder features of the split-everything model to enhance output features;

[0022] Step S22, uncertainty-guided prediction regularization module:

[0023] Selectively activate LayerNorm in the visual Transformer encoder to enhance the generalization ability of the model;

[0024] At the same time, during the training process, the true labels are used as "hints" to correct the confidence of the model and fine-tune the layer normalization layer of the visual Transformer encoder;

[0025] Assess model uncertainty by splitting all model outputs into a confidence score;

[0026] During the training process, when the confidence is low, the true label is used as a "hint" to help the model learn more accurate feature representation; the confidence of the segmentation model is regarded as the image confidence c i , the calculation method of pixel confidence is:

[0027]

[0028] Where P represents the prediction of the model, and the final confidence is calculated as the image-level confidence c i and pixel-level confidence c p The sum of is expressed as:

[0029]

[0030] Among them, the confidence is determined by Bernoulli distribution whether to use the true label as a "hint";

[0031] If the confidence is low enough, the model needs a specific answer hint to learn the correct prediction; otherwise, there is no need to provide a hint. The weight of the hint is determined by the confidence c, and the prediction is expressed as:

[0032] P′=c·P+(1-c)·Y

[0033] And the confidence loss is introduced to supervise the confidence c. When c→0, the confidence loss will increase. The confidence loss is defined as:

[0034]

[0035] As a further solution of the present invention: the specific steps in step S3 include:

[0036] The common cross entropy loss function is used as the loss function for segmenting all models, and the confidence loss is added. λ is a hyperparameter. The final loss function is:

[0037]

[0038] As a further solution of the present invention: the specific steps in step S4 include:

[0039] The polyp images in the test set were segmented, and the Sigmoid function was used to extract the binary segmentation results from the output probability map of the model. The average Dice and average IoU of the predicted values ​​and the true values ​​were calculated to verify the performance of the model and obtain the final endoscopic polyp segmentation model.

[0040] Compared with the prior art, the present invention has the following technical effects:

[0041] By adopting the above technical solution and designing an endoscopic polyp segmentation method, it is possible to divide the training set and test set for training and verification according to the public polyp segmentation dataset. By fusing multi-scale feature maps and multi-level feature maps, features can be enhanced across branches, thereby improving the diversity of extracted polyp image features. At the same time, an uncertainty-guided prediction regularization module is introduced to alleviate the uncertainty in the target domain training process, thereby enhancing the generalization ability of the model. Moreover, compared with previous polyp segmentation methods and methods based on segmentation-everything models, significant performance improvements have been achieved in polyp segmentation tasks, showing significant practicality. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] The specific embodiments of the present invention are described in detail below in conjunction with the accompanying drawings:

[0043] Figure 1 A schematic diagram of the steps of the embodiment disclosed in this application;

[0044] Figure 2 This is a result diagram of polyp segmentation according to an embodiment disclosed in the present application. DETAILED DESCRIPTION

[0045] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0046] Please refer to Figure 1 and Figure 2 In an embodiment of the present invention, an endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance includes the following steps:

[0047] Step S1: Obtain a data set and perform preprocessing. The specific steps include:

[0048] Step S11, obtaining a data set, and dividing the data set to obtain a training set and a test set;

[0049] In the specific implementation steps, multiple public polyp datasets including Kvasir-SEG, CVC-ClinicDB, CVC-ColonDB, ETIS and EndoScene are used. The training set includes 900 samples from Kvasir-SEG and 550 samples from CVC-ClinicDB, and the test set includes all data samples outside the training set.

[0050] Step S12: perform data enhancement on the training set data, including random rotation, random flipping, and random scaling operations.

[0051] Step S2: Based on the visual Transformer structure, the segmentation model designs the cross-branch feature enhancement module and the uncertainty-guided prediction regularization module, and integrates the trainable convolutional neural network encoder branch to obtain an endoscopic polyp segmentation model. The specific steps include:

[0052] Step S21, cross-branch feature enhancement module:

[0053] The structure of the segmentation model was modified based on the visual Transformer structure, removing its prompt input and prompt encoding parts, while retaining its visual Transformer encoder and mask decoder parts;

[0054] In this embodiment, MSCAN-L is selected as the convolutional neural network encoder branch, the segmentation everything model is selected as the basic model, ViT-H is selected as its backbone network, its prompt input and prompt encoding parts are deleted, and its visual Transformer encoding and mask decoding parts are retained.

[0055] In this embodiment,

[0056] Introduce a parallel convolutional neural network branch and add an additional multi-head cross-branch attention module to fuse the feature maps F from the convolutional neural network encoder. c and the feature map F of the visual Transformer encoder v , then the CBA formula of the multi-head cross-branch attention module is as follows:

[0057]

[0058] Where Q = F v W Q , K=V=F c W K , d is F v The number of feature dimensions of each head;

[0059] Fusion of deep contextual information and shallow boundary information from the encoder with the decoder features of the split-everything model to enhance output features;

[0060] Step S22, uncertainty-guided prediction regularization module:

[0061] Selectively activate LayerNorm in the visual Transformer encoder to enhance the generalization ability of the model;

[0062] Specifically, only a very small number of trainable parameters are introduced to improve the generalization of the model.

[0063] At the same time, during the training process, the true labels are used as "hints" to correct the confidence of the model and fine-tune the layer normalization layer of the visual Transformer encoder;

[0064] Assess model uncertainty by splitting all model outputs into a confidence score;

[0065] During the training process, when the confidence is low, the true label is used as a "hint" to help the model learn more accurate feature representation; the confidence of the segmentation model is regarded as the image confidence c i , the calculation method of pixel confidence is:

[0066]

[0067] Where P represents the prediction of the model, and the final confidence is calculated as the image-level confidence c i and pixel-level confidence c p The sum of is expressed as:

[0068]

[0069] Among them, the confidence is determined by Bernoulli distribution whether to use the true label as a "hint";

[0070] If the confidence is low enough, the model needs a specific answer hint to learn the correct prediction; otherwise, there is no need to provide a hint. The weight of the hint is determined by the confidence c, and the prediction is expressed as:

[0071] P′=c·P+(1-c)·Y

[0072] And the confidence loss is introduced to supervise the confidence c. When c→0, the confidence loss will increase. The confidence loss is defined as:

[0073]

[0074] Step S3: training the obtained endoscopic polyp segmentation model based on the cross entropy loss function, the specific steps include:

[0075] The common cross entropy loss function is used as the loss function for segmenting all models, and the confidence loss is added. λ is a hyperparameter. The final loss function is:

[0076]

[0077] In this embodiment, the image is input into the model. For the convolutional neural network branch, the image size is set to 320×320. At the same time, for the visual Transformer encoder, the image size is set to 1024×1024. Finally, the size of the output segmentation result is 256×256.

[0078] The specific implementation of the case experiment uses Python language, and uses the deep learning library PyTorch to build and train the model. The experimental platform is NVIDIA GeForce RTX 3090GPU, using AdamW optimizer, learning rate 0.0001, batch size 4, and training cycle 40 cycles.

[0079] Among them, the hyperparameter settings in the loss function, the weight of the cross entropy loss function is 1, and the weight of the confidence loss is initialized to 0.1. If the confidence loss value is greater than 0.3, its weight is increased, otherwise it is reduced.

[0080] Step S4: After the polyp images of the test set are segmented according to the trained endoscopic polyp segmentation model, the segmentation results are extracted using the Sigmoid function, and the performance of the model is verified to obtain the final endoscopic polyp segmentation model. The specific steps include:

[0081] The polyp images in the test set were segmented, and the Sigmoid function was used to extract the binary segmentation results from the output probability map of the model. The average Dice and average IoU of the predicted values ​​and the true values ​​were calculated to verify the performance of the model and obtain the final endoscopic polyp segmentation model.

[0082] In the specific implementation steps, the polyp images in the test set are segmented, and the Sigmoid function is used to obtain the probability map from the output. The probability threshold is preset to 0.5 to obtain the binary segmentation result. The average Dice similarity coefficient and the average IoU intersection-over-union ratio of the evaluation prediction value and the true value are calculated as segmentation metrics to verify the performance of the model.

[0083] In this example, the average Dice similarity coefficient on all test sets is 0.890, and the average IoU intersection-over-union ratio is 0.817, which achieves significant results in the polyp segmentation task. The qualitative segmentation results are shown in the following table.

[0084] Each column in the table is the original image, the true label, the polyp segmentation method PraNet, the polyp segmentation method based on the segmentation model SAM-Adapter, MedSAM, SAMUS, and the segmentation prediction results of the endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance. The test scores are shown in the following table (the bold in the table indicates the highest score, and the underline indicates the second highest score):

[0085]

[0086] In summary, the method of the present invention has richer feature extraction capabilities, can effectively improve the performance of all segmentation models in polyp segmentation, and has significant practicality.

[0087] Although the embodiments of the present invention have been shown and described, it will be appreciated by those skilled in the art that various changes, modifications, substitutions and variations may be made to the embodiments without departing from the principles and spirit of the present invention. The scope of the present invention is defined by the appended claims and their equivalents, and all should be included within the scope of protection of the present invention.

Claims

1. An endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance, characterized in that: The following steps are involved: Step S1, obtaining a data set and performing preprocessing; Step S2, designing a cross-branch feature enhancement module and an uncertainty-guided prediction regularization module based on the segmentation model of the visual Transformer structure, and integrating a trainable convolutional neural network encoder branch to obtain an endoscopic polyp segmentation model; Step S3, training the obtained endoscopic polyp segmentation model based on the cross entropy loss function; Step S4: After the polyp images of the test set are segmented according to the trained endoscopic polyp segmentation model, the segmentation results are extracted using the Sigmoid function, and the performance of the model is verified to obtain the final endoscopic polyp segmentation model.

2. According to claim 1, the endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance is characterized in that: The specific steps in step S1 include: Step S11, obtaining a data set, and dividing the data set to obtain a training set and a test set; Step S12: perform data enhancement on the training set data, including random rotation, random flipping, and random scaling operations.

3. According to claim 1, the endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance is characterized in that: The specific steps in step S2 include: Step S21, cross-branch feature enhancement module: The structure of the segmentation model was modified based on the visual Transformer structure, removing its prompt input and prompt encoding parts, while retaining its visual Transformer encoder and mask decoder parts; Introduce a parallel convolutional neural network branch and add an additional multi-head cross-branch attention module to fuse the feature maps F from the convolutional neural network encoder. c and the feature map F of the visual Transformer encoder v , then the CBA formula of the multi-head cross-branch attention module is as follows: Where Q = F v W Q , K=V=F c W K , d is F v The number of feature dimensions of each head; Fusion of deep contextual information and shallow boundary information from the encoder with the decoder features of the split-everything model to enhance output features; Step S22, uncertainty-guided prediction regularization module: Selectively activate LayerNorm in the visual Transformer encoder to enhance the generalization ability of the model; At the same time, the real labels are used as "hints" during the training process to correct the confidence of the model and fine-tune the layer normalization layer of the visual Transformer encoder; Assess model uncertainty by splitting all model outputs into a confidence score; During training, when the confidence is low, the true label is used as a "hint" to help the model learn more accurate feature representations; the confidence of the segmentation model is regarded as the image confidence c i , the calculation method of pixel confidence is: Where P represents the prediction of the model, and the final confidence is calculated as the image-level confidence c i and pixel-level confidence c p The sum of is expressed as: Among them, the confidence level is determined by Bernoulli distribution whether to use the true label as a "hint"; If the confidence is low enough, the model needs a specific answer hint to learn the correct prediction; otherwise, there is no need to provide a hint. The weight of the hint is determined by the confidence c, and the prediction is expressed as: P′=c·P+(1-c)·Y And the confidence loss is introduced to supervise the confidence c. When c→0, the confidence loss will increase. The confidence loss is defined as:

4. According to claim 1, the endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance is characterized in that: The specific steps in step S3 include: The common cross entropy loss function is used as the loss function for segmenting all models, and the confidence loss is added. λ is a hyperparameter. The final loss function is:

5. The endoscopic polyp segmentation method based on cross-branch feature enhancement and uncertainty guidance according to claim 1, characterized in that: The specific steps in step S4 include: The polyp images in the test set were segmented, and the Sigmoid function was used to extract the binary segmentation results from the output probability map of the model. The average Dice and average IoU of the predicted values ​​and the true values ​​were calculated to verify the performance of the model and obtain the final endoscopic polyp segmentation model.

Images