Colon polyp image segmentation system and method based on UNet

Abstract

The invention discloses a UNet-based colon polyp image segmentation system and method, and the method comprises the steps: constructing and training a colon polyp segmentation model, inputting a to-be-segmented colorectal polyp image into the trained colon polyp segmentation model, and obtaining a polyp segmentation image; the colon polyp segmentation model takes UNet as a main network, and comprises a cascade convolution module, a convolution attention module, a feature fusion module and a boundary refining module; according to the model provided by the invention, the segmentation accuracy of the colon polyp edge can be effectively improved, and meanwhile, the robustness of the network is improved.

Description

A UNet-based image segmentation system and method for colon polyps Technical Field

[0001] This invention relates to the field of medical image processing technology, and in particular to a UNet-based system and method for segmenting colon polyp images. Background Technology

[0002] Colorectal cancer is one of the most common malignant tumors worldwide, ranking third in cancer incidence and second in cancer mortality, making it a major challenge that urgently needs to be addressed in the global public health field. In China, the number of colorectal cancer patients has been rising steadily in recent years. Studies show that approximately 95% of colorectal cancers develop from colorectal adenomatous polyps, which represent an early stage of colorectal cancer development. Accurate identification of adenomatous polyps is crucial for the early diagnosis and treatment of colorectal cancer.

[0003] Colonoscopy is currently the most effective method for detecting polyps, clearly displaying the location and appearance of lesions, thus assisting doctors in timely detection and removal before malignant transformation. However, automated segmentation of colonic polyps faces several challenges. First, polyps vary significantly in morphology and size. Polyps from different patients can range in size from small lesions of a few millimeters to large lesions of several centimeters, and their shapes can be either convex structures with thin stalks or flat, plaque-like lesions closely adhering to the mucosa. This uncertainty in morphology and size significantly increases the difficulty of segmentation. Furthermore, in endoscopic images, the color and texture of polyps are sometimes almost indistinguishable from the surrounding healthy tissue, especially flat polyps, whose blurred boundaries and low distinction from the background make traditional segmentation methods prone to missed detections (mistaking polyps for normal mucosa) or false detections (mistaking normal folds for polyps). Second, endoscopic imaging can be affected by various interfering factors, such as uneven illumination, motion blur, and obstruction by intestinal contents, further complicating the segmentation task. Therefore, research on computer-aided diagnostic algorithms for automatic polyp identification is of great significance for improving the diagnostic accuracy and efficiency of colonoscopy.

[0004] Deep learning algorithms, especially convolutional neural network models, can automatically learn rich feature representations of colon polyps through multi-layer combinations of convolutional and pooling layers, thus greatly improving segmentation accuracy and achieving significant progress in medical image segmentation tasks. Compared with traditional manual methods, deep learning-based colon polyp image segmentation methods have the advantages of low cost, high efficiency, and high accuracy. However, existing segmentation algorithms still face many challenges, such as high computational complexity, insufficient generalization ability, and the need to improve the ability to identify polyps under different imaging conditions. These problems restrict their widespread application in clinical practice. Therefore, existing technologies have shortcomings and need improvement. Summary of the Invention

[0005] The purpose of this application is to provide a UNet-based method for segmenting colon polyp images. The UNet network is used as the backbone feature extraction network, and a cascaded convolution module, a convolutional attention module, a feature fusion module, and a boundary refinement module are designed to improve the accuracy and robustness of colon polyp image segmentation.

[0006] To achieve the above objectives, the technical solution of this invention is implemented as follows: A UNet-based method for segmenting colonic polyps images, comprising the following steps: S1: Obtaining a colonic polyp dataset and dividing it into a training set, a validation set, and a test set; S2: Constructing a UNet-based colonic polyp image segmentation network; Using UNet as the backbone network, performing multi-stage feature extraction on the input image, outputting four-level feature maps (F1, F2, F3, and F4) to provide multi-dimensional feature support for subsequent modules; Designing and using a cascaded convolutional module, which effectively compensates for the edge detail information lost in deep features by fusing adjacent deep and shallow features, and further strengthens the relevance of contextual information by combining a convolutional attention module, thereby comprehensively improving the accuracy of colonic polyp segmentation and the accuracy of boundary localization; Designing and using a convolutional attention module, employing a dual attention mechanism of channel attention and spatial attention. The channel attention module aims to highlight meaningful channels, while the spatial attention module focuses on highlighting key parts of spatial features. By dynamically weighting and optimizing the fused features, the saliency expression of key channels and spatial regions is further strengthened, thereby improving the contextual relevance and semantic consistency of feature representation.

[0007] The design employs a feature fusion module to connect shallow and deep features. By adjusting the scale (downsampling / upsampling), the spatial dimensions of features at different levels are made consistent. Then, features are fused through splicing and convolution operations to enhance the network's ability to perceive the boundaries, shape, and internal structure of polyps. The design also employs a boundary refinement module to apply wavelet transform and edge detection operators to the input image to extract boundary features, thereby improving the network's ability to capture boundary details.

[0008] S3: Train the colon polyp detection model based on the UNet network using the training set, validate it using the validation set during training, and evaluate the trained model using the test set.

[0009] A joint loss function, constructed by combining weighted IoU and weighted BCE, is used to balance the ratio of positive to negative samples. Minimizing the loss function guides model parameter updates. The preprocessed training set is input into the network for iterative parameter updates. A validation set is introduced during training to monitor model performance in real time and dynamically adjust the training strategy to avoid overfitting. Dice coefficient, mean intersection-over-union ratio (mIoU), and accuracy (ACC) are used as core evaluation metrics to comprehensively measure the model's segmentation performance.

[0010] The beneficial effects of this invention are as follows: This invention uses UNet as the backbone network of the encoder, which has a stronger feature representation than other neural networks, and maintains high computational efficiency through optimized calculation and parameter usage, making it suitable for a variety of complex visual tasks.

[0011] The design utilizes a cascaded convolution module and introduces a convolutional attention module. By splicing and fusing features from adjacent layers, it significantly enhances the ability to express shallow detailed features while effectively compensating for the loss of semantic information in deep features, thereby significantly improving the problem of blurred polyp segmentation boundaries.

[0012] The design utilizes a feature fusion module, which fully leverages shallow and deep information to provide the network with more comprehensive image background information and relationships between segmentation targets, thereby improving the network's ability to perceive polyp boundaries, shapes, and internal structures.

[0013] The design utilizes a boundary refinement module, applying wavelet transform and edge detection operators to enhance the ability to capture details of polyp boundaries and improve the network's performance when handling complex structures.

[0014] The model proposed in this invention can effectively improve the segmentation accuracy of colon polyp edges, while also enhancing the robustness of the network. Attached Figure Description

[0015] Figure 1 is a schematic diagram of the structure of the colon polyp image segmentation method of the present invention; Figure 2 is a schematic diagram of the structure of the cascaded convolution module of the present invention; Figure 3 is a schematic diagram of the structure of the convolutional attention module of the present invention; Figure 4 is a schematic diagram of the structure of the feature fusion module of the present invention; Figure 5 is a schematic diagram of the structure of the boundary refinement module of the present invention. Detailed Implementation

[0016] The present invention will be described in detail below with reference to specific embodiments.

[0017] This invention provides a UNet-based method for detecting colon polyps. A colon polyp segmentation model is constructed and trained. The polyp image to be segmented is input into the trained colon polyp segmentation model to obtain a polyp segmentation map. The colon polyp segmentation model includes a cascaded convolution module, a convolutional attention module, a feature fusion module, and a boundary refinement module.

[0018] The training process of the colorectal polyp segmentation model includes the following steps: S1: Obtain the colorectal polyp dataset and divide it into a training set, a validation set, and a test set; Specifically, the present invention uses the ClinicDB dataset as the original colorectal polyp image training dataset.

[0019] The ClinicDB dataset focuses on the detection and research of colorectal cancer and contains 612 colonoscopy images with a uniform size of 388 × 284 pixels. For experimental training, this invention uses 400 images for network training and 212 images for network testing.

[0020] S2: Construct a UNet-based image segmentation network for colon polyps, including a cascaded convolution module, a convolutional attention module, a multi-scale feature fusion module, and a boundary refinement module. Through the synergistic effect between modules, the accuracy and robustness of colon polyp image segmentation are improved.

[0021] Specifically, in the cascaded convolution module, the resolution and number of channels of the encoded features of adjacent high and low layers are aligned first, and the channels are spliced ​​and fused. After the feature splicing and fusion is completed, an attention mechanism is embedded to perform dynamic weighted optimization, thereby achieving effective complementarity of cross-layer features and enhancement of key region features, and improving the discriminativeness and contextual relevance of features.

[0022] First, the model receives features from two adjacent encoding layers, performs downsampling on the shallow features, then extracts features using 3×3 convolution, performs batch normalization for stable training, and then applies ReLU to filter and activate the batch-normalized feature maps to further ensure the model's training stability and ability to extract complex features. Finally, the processed shallow features are concatenated and fused with the deep features to compensate for the lost details in the deep features.

[0023] Specifically, the convolutional attention module employs a dual attention mechanism: channel attention and spatial attention. The channel attention module aims to highlight meaningful channels, while the spatial attention module focuses on highlighting key parts of spatial features.

[0024] First, the input feature map passes through a channel attention module, which generates global feature descriptions in two channel dimensions through average pooling and max pooling operations, resulting in two 1×1×C vectors. These vectors are then processed by a multilayer perceptron (MLP) to learn the dependencies between channels. The results are summed and weighted by the sigmoid function to obtain the channel weighting coefficients M. C (F); Multiply this coefficient by the input feature map F to obtain the channel-weighted feature map F'.

[0025] M C The formulas for (F) and F' are shown in (1) and (2) respectively.

[0026] (1) (2) In formula (1), the input feature map is denoted as F, AvgPool() and MaxPool() represent average and max pooling respectively, MLP represents multilayer perceptron, σ(·) refers to the Sigmoid function, M C (F) is the one-dimensional channel attention feature map obtained after weighting; in formula (2), ② represents element-level multiplication, and F' represents the feature map after channel weighting.

[0027] Subsequently, the channel-weighted feature map F' is input into the spatial attention module, where average pooling and max pooling are performed on each feature point along the channel dimension. The results are concatenated and processed by 7×7 convolution and adaptively weighted by the Sigmoid activation function to obtain a spatial weighted vector. This vector is then multiplied by F' to obtain a new weighted feature map F'' that considers both channel and spatial attention.

[0028] The combined unit is executed twice on the weighted feature map. This combined unit first extracts features using a 3×3 convolution, then performs batch normalization for stable training. Next, ReLU is applied to filter and activate the batch-normalized feature map, further ensuring model training stability and complex feature extraction capabilities. Finally, the refined features are processed in two branches: one part is passed to the same-layer decoder for feature fusion; the other part is fed into a cascaded convolution module for fusion with features from adjacent deep coding layers.

[0029] Specifically, in the feature fusion module, F1' and F3' are downsampled and upsampled respectively to adjust to the same spatial size as F2'. Then, feature fusion is performed through concatenation and convolution operations to generate a preliminary prediction map BPM. Subsequently, F4' is upsampled to adjust to the same spatial size as BPM. Then, feature fusion is performed through concatenation and convolution operations to generate an output feature map.

[0030] Specifically, in the boundary refinement module, wavelet transform and edge detection operators are applied to the input image to extract boundary features and improve the network's ability to capture boundary details.

[0031] First, the module employs edge detection operators in eight directions (0°, 22.5°, 45°, 67.5°, 90°, 112.5°, 135°, and 157.5°) to capture edge gradient information in different directions. Simultaneously, wavelet transform is used to decompose the input image into sub-images of different frequencies and directions to extract edge features. Next, the two are added together to fully fuse multi-scale, multi-directional, and multi-frequency polyp boundary information. Subsequently, the fused result is multiplied with the input image to obtain the final output feature map.

[0032] S3: Train the UNet-based colon polyp detection model using the training set and validate it using the validation set during training. Evaluate the trained model using the test set and obtain the detection and classification results of colorectal polyps.

[0033] Specifically, a loss function for the colorectal polyp segmentation model is constructed based on attention prediction features, feature maps at all levels, and polyp segmentation map labels corresponding to the original colon polyp images. The parameters of the colorectal polyp segmentation model are updated with the minimum loss function as the optimization objective, thus completing the training of the colorectal polyp segmentation model.

[0034] The network uses a joint loss function constructed by adding weighted IoU and weighted BCE. The weighted IoU and weighted BCE are shown below.

[0035]

[0036]

[0037] Where N represents the total number of pixels, and yi represents the label value of the i-th pixel. i represents the predicted value of the i-th pixel. wi represents the weight coefficient of the i-th pixel.

[0038] In the results output module, to comprehensively evaluate the overall performance of the proposed model in medical image segmentation, three commonly used evaluation metrics are employed: Dice coefficient, mean intersection-over-union ratio (mIoU), and accuracy (ACC). These metrics provide a comprehensive evaluation of the model's performance in image segmentation tasks. The specific calculation formulas are as follows:

[0039] TP, TN, FP, and FN represent true positive, true negative, false positive, and false negative, respectively.

[0040] The data from ClinicDB was input into the network for training and testing. The experimental results of the method of the present invention in this embodiment are shown in Table 1.

[0041] Table 1 It should be understood that those skilled in the art can make improvements or modifications based on the above description, and all such improvements and modifications should fall within the protection scope of the appended claims.

Claims

1. A UNet-based method for segmenting colon polyp images, characterized in that, A colon polyp segmentation model is constructed and trained. The colon polyp images to be segmented are input into the trained colon polyp segmentation model to obtain polyp segmentation maps. The colon polyp segmentation model uses UNet as the backbone network and includes a cascaded convolutional module, a convolutional attention module, a feature fusion module, and a boundary refinement module. The training process of the colon polyp segmentation model includes the following steps: S1: Obtain a colon polyp dataset and divide it into a training set, a validation set, and a test set; S2: Construct a UNet-based colon polyp image segmentation network, which includes a cascaded convolutional module, a convolutional attention module, a feature fusion module, and a boundary refinement module; S3: Train the colon polyp segmentation model using the training set, validate it using the validation set during training, and evaluate the model using the test set after training to obtain the segmentation results of the colon polyps.

2. The colon polyp image segmentation method based on improved UNet according to claim 1, characterized in that, In step S2: Multi-stage feature extraction is performed on the input image, and four-level feature maps of F1, F2, F3 and F4 are output to provide multi-dimensional feature support for subsequent modules.

3. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, The working process of the cascaded convolution module includes: (1) receiving the features of the two adjacent coding layers, performing downsampling on the shallow features, then extracting features through 3×3 convolution, performing batch normalization, and then applying the ReLU activation function to filter and activate the feature map; (2) splicing and fusing the processed shallow features with the deep features to compensate for the lost details of the deep features; (3) embedding an attention mechanism to dynamically weight and optimize the fused features, realizing effective complementarity of cross-layer features and enhancement of key region features, and improving the discriminativeness and contextual relevance of features.

4. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, The working process of the convolutional attention module includes: (1) Channel attention module processing: The input feature map F is processed by average pooling and max pooling to generate global feature descriptions in two channel dimensions respectively. The global feature descriptions are learned by multilayer perceptron (MLP) and then added together. The channel weighting coefficients M are then obtained by the Sigmoid function. C (F), multiply the channel weighting coefficients with the input feature map F element-wise to obtain the channel-weighted feature map F'; (2) Spatial attention module processing: perform average pooling and max pooling on the channel dimension of the feature map F' respectively, concatenate the pooling results and process them with 7×7 convolution and Sigmoid activation function to obtain the spatial weighting vector, multiply the spatial weighting vector with the feature map F' to obtain the weighted feature map F'' that considers both channel and spatial attention; (3) Repeat the feature processing unit twice on the weighted feature map F'', the feature processing unit includes 3×3 convolution, batch normalization and ReLU activation function operation in sequence; (4) Divide the refined features into two branches for output: one branch is passed to the same layer decoder for feature fusion, and the other branch is sent to the cascaded convolution module to fuse with the features of the adjacent deep coding end.

5. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, The working process of the multi-scale feature fusion module includes: (1) downsampling and upsampling the F1' and F3' feature maps output by the encoder respectively, adjusting them to the same spatial size as the F2' feature map, and performing feature fusion through splicing and convolution operations to generate a preliminary prediction map BPM; (2) upsampling the F4' feature map, adjusting it to the same spatial size as the preliminary prediction map BPM, and performing feature fusion through splicing and convolution operations to generate the final output feature map.

6. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, The working process of the boundary refinement module includes: (1) using edge detection operators in eight directions (0°, 22.5°, 45°, 67.5°, 90°, 112.5°, 135° and 157.5°) to capture edge gradient information in different directions; (2) using wavelet transform to decompose the input image into sub-images of different frequencies and directions, and extracting edge features; (3) adding the features extracted by the edge detection operators with the features extracted by the wavelet transform to fully integrate the polyp boundary information of multiple scales, multiple directions and multiple frequencies; (4) multiplying the fusion result with the input image to obtain the final output feature map.

7. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, In step S1: The dataset used is the ClinicDB dataset, which focuses on the detection and research of colorectal cancer and contains 612 colonoscopy images with a uniform size of 388 × 284 pixels. According to the experimental training requirements, in this invention, 400 images are used for network training and 212 images are used for network testing.

8. The UNet-based colon polyp image segmentation method according to claim 1, characterized in that, In step S3: the loss function for model training is a joint loss function constructed by adding weighted IoU and weighted BCE, and the model parameters are updated with the goal of minimizing the loss function; the metrics used for model evaluation include Dice coefficient, mean intersection-over-union ratio (mIoU) and accuracy (ACC); during training, a validation set is introduced to monitor model performance in real time and dynamically adjust the training strategy to avoid overfitting.

9. A colon polyp image segmentation system, characterized in that, The colon polyp segmentation model includes a UNet backbone network, a cascaded convolutional module, a convolutional attention module, a feature fusion module, and a boundary refinement module. The backbone network is used to extract multi-scale features of the image. The cascaded convolutional module is used to fuse adjacent shallow and deep features and enhance key region features. The convolutional attention module is used to optimize feature representation through a dual attention mechanism. The feature fusion module is used to establish the correlation between shallow and deep features. The boundary refinement module is used to extract boundary features. The model is trained using the method described in any one of claims 1-8.

Images