Medical image automatic segmentation method of online clustering attention mechanism based on K neighbor graph guidance

Abstract

The invention discloses a medical image segmentation method (GCA-Net) based on graph-oriented clustering attention. According to the method, a pyramid vision Transform (PVT) backbone network is adopted to extract multi-scale features, and the features are recovered in combination with a U-Net decoder. In order to reduce the calculation complexity of a traditional Transform and avoid feature collapse, an end-to-end clustering framework is introduced, a K-nearest neighbor (KNN) graph is dynamically constructed by using image features, and a semantic clustering center is learned through a graph convolutional network (GCN) under the condition of not depending on an external clustering algorithm. On the basis, a clustering attention mechanism compression key and value are designed, so that the calculation complexity is reduced to be linear. In the decoding stage, a lightweight convolution attention module is fused to enhance local details, and a dynamic foreground enhancement module is introduced to enhance foreground and background boundary comparison. Experiments on a plurality of medical image data sets show that the method realizes higher segmentation precision while ensuring efficiency, is especially suitable for segmentation of lesions with irregular shapes, and has clinical auxiliary diagnosis value.

Description

Technical Field

[0001] This invention relates to the fields of computer vision and medical image processing technology, and is an automatic medical image segmentation method based on graph-guided clustering attention and dynamic foreground enhancement. Background Technology

[0002] With the development of imaging technologies such as CT, MRI, endoscopy, and dermoscopy, medical imaging has become a crucial basis for clinical diagnosis, efficacy evaluation, and surgical planning. Medical image segmentation, through the detailed delineation of organs, tissues, and lesion regions, provides key support for the quantitative analysis of the structure and function of the heart, brain, and tumors, as well as the detection and evaluation of skin lesions and colon polyps, significantly improving the accuracy and repeatability of disease screening and follow-up monitoring. However, medical images generally suffer from problems such as complex organ structures, variable lesion morphologies, blurred boundaries, and low contrast between lesions and normal tissues, making the construction of segmentation models that combine accuracy, efficiency, and generalization ability a significant challenge in the field of medical image analysis.

[0003] Early medical image segmentation was mainly dominated by convolutional neural networks (CNNs) and their encoder-decoder structures (such as UNet). CNNs can efficiently extract multi-scale features through local convolutional kernels and downsampling operations, but their receptive field is limited by the network depth and convolutional kernel size. They are essentially biased towards local modeling and lack the ability to characterize global dependencies between distant regions. They often encounter performance bottlenecks when facing medical scenarios with complex anatomical structures and discrete lesion distributions.

[0004] To overcome the limitations of CNNs, Visual Transformer (ViT) has been introduced into the field of medical image segmentation (e.g., TransUNet). Thanks to its self-attention mechanism, ViT performs exceptionally well in global semantic modeling. However, these methods have significant drawbacks: First, they generally rely on patch representations with regular grid partitioning, naturally ignoring the semantic structure and scale differences of images, treating all regions equally, and making it difficult to adaptively enhance the expression of organ boundaries and small lesions; second, the computational complexity of self-attention increases quadratically with the number of tokens (O(N...). 2 In high-resolution medical image scenarios, this leads to huge computational and storage overhead. In addition, on medical datasets with limited labeled samples, Transformer is prone to feature collapse, that is, the attention matrix gradually becomes uniform, making it difficult for deep networks to learn discriminative long-range dependencies.

[0005] To address the aforementioned complexity issue, existing research has attempted various improvement strategies, but shortcomings remain: Window Attention: Limiting computation to a local window reduces computational cost, but it also severs global connections and restricts the scope of perception. Sparse Attention: This strategy typically removes some tokens by pruning, which can easily disrupt the spatial structure of the original features, leading to performance degradation and difficulty in generalization. Memory-compressed attention reduces overhead by downsampling keys and values, but it is still based on a regular grid and cannot differentiate based on semantic importance.

[0006] In summary, existing technologies struggle to achieve both computational efficiency and precise segmentation of irregularly shaped and structurally complex organs and lesions. Therefore, there is an urgent need for an image segmentation method that can overcome the limitations of regular grids, dynamically aggregate information based on semantic features, and effectively prevent feature collapse while balancing global modeling with local details. Summary of the Invention

[0007] The purpose of this invention is to provide an automatic medical image segmentation method based on graph-guided clustering attention and dynamic foreground enhancement. Addressing the problems of high computational complexity, susceptibility to feature collapse, and difficulty in capturing details of irregular lesions in existing Vision Transformers for medical image segmentation, this invention proposes an end-to-end clustering attention network.

[0008] The technical solution to achieve the objective of this invention is as follows: Based on the Pyramid Visual Transformer (PVT) backbone network, an encoder-decoder segmentation network is constructed, which includes a graph-guided clustering attention module, a convolutional attention module, and a dynamic foreground enhancement module.

[0009] Step 1: Based on existing medical image datasets (such as skin lesions, polyps, multi-organ CT or cardiac MRI), preprocess the data, including image size normalization and data augmentation operations.

[0010] Step 2: Divide the processed data into training set, test set and validation set.

[0011] Step 3: Construct a medical image segmentation network framework based on the training set: Step 3.1: Construct an encoder using the Pyramid Visual Transformer (PVT) as the backbone network to extract multi-scale feature maps.

[0012] Step 3.2: Construct the Graph-Guided Clustering Attention (GCA) module: In the deep feature space, image features are treated as nodes, and a K-Nearest Neighbor (KNN) graph is dynamically constructed based on feature similarity; the node features are updated using a Graph Convolutional Network (GCN) and an expanded neighborhood aggregation strategy to learn semantically representative cluster centers; the probability of image features belonging to cluster centers is calculated, and cluster center tokens are generated as compressed keys and values, which are then used to perform attention calculations with the query generated from the original features to obtain global semantic features with long-distance dependencies.

[0013] Step 3.3: Construct the decoder and feature fusion module: Use U-Net-style skip connections to fuse the multi-scale features of the encoder with the upsampled features of the decoder; introduce a lightweight convolutional attention module during the fusion process to enhance the expression of local detail features through channel attention and spatial attention mechanisms.

[0014] Step 3.4: Construct the Dynamic Foreground Enhancement (DFE) module: At the end of the decoder, a dynamic clustering prototype is generated using the feature map. The prototype is then projected back into the pixel space to weight the feature response, adaptively enhancing the contrast between the foreground object and the background.

[0015] Step 3.5: Construct a hybrid loss function based on the cross-entropy loss function and the Dice loss function, calculate the difference between the network prediction results and the true labels, and update the network parameters using the gradient descent algorithm until the model converges.

[0016] Step 4: Repeat step 3 to train the segmentation network model.

[0017] Step 5: Use the trained segmentation network to test the test data and output the final medical image segmentation results.

[0018] Compared with the prior art, the significant advantages of the present invention are: (1) Reduce computational complexity while preserving global dependencies: The graph-guided clustering attention mechanism proposed in this invention learns a small number of semantic cluster centers to replace the full graph tokens in attention interactions, reducing the computational complexity from O(N) to O(N). 2 The time complexity is reduced to linear O(NM) (M is much smaller than N), which significantly reduces memory usage and computation while effectively retaining the advantage of Transformer in capturing long-range semantic dependencies.

[0019] (2) Solving feature collapse and adapting to irregular structures: Traditional gridded attention is prone to feature convergence (feature collapse). This invention can flexibly model irregular non-Euclidean structures (such as lesions and organ edges) in medical images by dynamically constructing a K-nearest neighbor graph structure, and through end-to-end clustering learning, the attention mechanism focuses on semantically related regions, thereby extracting more discriminative features.

[0020] (3) Multi-scale fusion and detail enhancement: It combines the multi-scale feature extraction capability of PVT, the convolutional attention module in the decoding stage, and the dynamic foreground enhancement module at the end. The convolutional attention compensates for the Transformer's shortcomings in local details, while the dynamic foreground enhancement module significantly improves the boundary distinction between the foreground and the background, making the model perform better when dealing with small lesions or blurred boundaries. Attached Figure Description

[0021] Figure 1 This is a schematic diagram of the overall process and network architecture of the medical image automatic segmentation method (GCA-Net) based on graph-guided clustering attention and dynamic foreground enhancement of the present invention.

[0022] Figure 2 This is a diagram illustrating the internal structure of the Graph-Guided Clustering Attention (GCA) module and the details of graph convolution operations in this invention.

[0023] Figure 3 This is a qualitative comparison of the method of this invention with several existing mainstream medical image segmentation methods (such as TransUNet, Swin-UNet, PVT-CASCADE, etc.) on skin, polyp, and multi-organ datasets.

[0024] Figure 4 This figure shows the experimental results comparing the method of this invention with several existing mainstream medical image segmentation methods on skin, polyp, and multi-organ datasets, using multiple quantitative indicators. Figure 5 The accompanying drawings are a summary of the invention and describe the overall flow logic of the invention. Detailed Implementation

[0025] The invention will now be described in further detail with reference to examples and accompanying drawings: This invention proposes a medical image segmentation method (GCA-Net) based on graph-guided clustering attention and dynamic foreground enhancement. This embodiment uses skin lesion segmentation (ISIC 2018 dataset) and multi-organ segmentation (Synapse dataset) as examples to illustrate the implementation process of this invention in detail.

[0026] The specific implementation process of this invention mainly includes four stages: data preprocessing, network model construction, model training, and inference testing.

[0027] I. Data Preprocessing Stage Step 1: Obtain the medical image dataset.

[0028] This embodiment selects the ISIC 2018 dermoscopy image dataset (containing 2694 images) and the Synapse multi-organ CT dataset (containing 30 abdominal CT scans). For ISIC 2018, all RGB images and their corresponding segmentation masks are uniformly resized. pixels. For the Synapse dataset, crop the abdominal window and resize the slice to... Step 2: Data augmentation and partitioning.

[0029] To improve the model's robustness and prevent overfitting, online data augmentation was performed on the input images during the training phase. Augmentation operations included: random horizontal and vertical flipping; random rotation (90°, 180°, 270°); and normalization: normalizing pixel values ​​to the [0, 1] interval. The processed dataset was then randomly divided into training and test sets in a 7:3 ratio.

[0030] Step 3.1: Construct a pyramid feature extraction encoder A pre-trained PVT-v2-B2 (Pyramid Vision Transformer) network is used as the backbone network. Input image. After four stages of feature extraction, four feature maps of different scales are generated. . Number of channels As the layer increases, the spatial resolution is halved with each layer (representing the input size). To construct a global semantic context, Upsampling and The features are concatenated along the channel dimension to form an aggregated feature tensor, which is then fed into the GCA module.

[0031] Step 3.2: Construct the Graph-Guided Clustering Attention (GCA) module This module is the core of this invention, aiming to solve the computational complexity problem in long-distance dependency modeling. The specific implementation steps are as follows: (1) Dynamic Graph Construction: Input feature sequence (where N is the number of pixels / tokens) are considered as vertices of an undirected graph. Set. Calculate the Euclidean distance between nodes in the feature space, and dynamically select K nearest neighbors (K-Nearest Neighbors, K=9 in this embodiment) for each node to construct the graph structure. .

[0032] (2) Dilated Graph Convolution Feature Update: To expand the receptive field and avoid overly smooth features, an expanded KNN strategy with an expansion rate of d is adopted. For each vertex v, its nth vertex is selected. The nearest neighbors constitute the expanded neighborhood. .

[0033] Feature updates employ the Max-Relative GCN algorithm: in, This represents element-wise feature difference max pooling, an operation that can keenly capture gradient changes at irregular boundaries in medical images.

[0034] (3) Semantic clustering and key-value compression: The features updated after graph convolution are passed through a... Convolution and Softmax functions generate clustering assignment matrices. , where M is the preset number of cluster centers.

[0035] Use this matrix to analyze the original input. We perform weighted aggregation to obtain M semantic cluster center features. : These cluster centers It represents a set of regions in an image that have similar semantics (such as "lesion core", "normal skin", "background").

[0036] (4) Linear complexity attention calculation: Original features Linear mapping to query Cluster centers Linear mapping to key Sum Perform multi-head self-attention calculation: Therefore, the computational complexity of the attention mechanism is reduced from the traditional Reduce to This enables efficient global modeling.

[0037] Step 3.3: Constructing the decoder and enhancing convolutional attention A U-Net-style decoder structure is adopted. The deep semantic features output by the GCA module are processed. Restore the multi-scale structure and upsample layer by layer. Feature fusion: Combine the upsampled features with the corresponding features of the backbone network. The data is then stitched together. Convolutional attention enhancement: A lightweight convolutional attention module (CBAM variant) is introduced at each decoding stage. This module generates channel attention weights by computing max pooling and average pooling in parallel, followed by a shared MLP; spatial attention weights are then generated through spatial convolution. This step is specifically designed to recover minute lesion edge details lost during downsampling.

[0038] Step 3.4: Construct the Dynamic Foreground Enhancement (DFE) module At the end of the decoder (the highest resolution layer), a DFE module is introduced to adaptively distinguish between foreground and background. The feature map... Mapping to the prototype space, the clustering prototype is obtained by calculating the autocorrelation matrix. Then, regarding the prototype... Compression and Softmax normalization are performed to obtain the global semantic descriptor. Finally, the descriptor is projected back into the pixel space and element-wise multiplied with the original features: This step is similar to a "dynamic gating" mechanism that automatically suppresses background noise response and significantly enhances the feature intensity of foreground lesions.

[0039] Step 4: Loss function and parameter optimization.

[0040] Loss function: A hybrid loss function is used. .in This is the binary cross-entropy loss, used for pixel classification; The Dice coefficient loss is used to address the extreme imbalance between positive and negative samples (lesions and background) in medical images. In this embodiment, weights are set... .

[0041] Training was performed using the AdamW optimizer, with an initial learning rate set to... Weight decay is set to The batch size was set to 24, and 200 epochs were trained on a single NVIDIA RTX 3090 GPU.

[0042] Step 5: Performance evaluation.

[0043] The trained model was validated on the ISIC 2018 test set, and the evaluation metrics included mIoU (mean crossover ratio) and DSC (Dice similarity coefficient).

[0044] Comparison results: Compared with the current mainstream Transformer method Swin-UNet, the DSC of the proposed method GCA-Net is improved from 87.79% to 90.59%; compared with the latest Mamba architecture model CCViM, the mIoU is improved from 81.92% to 82.79%.

[0045] Ablation experiment analysis: Experiments show that removing the GCA module leads to a decrease in DSC of approximately 1.8%, and removing the DFE module leads to a decrease in DSC of approximately 0.5%. This demonstrates the effectiveness of graph-guided clustering in capturing global context and foreground enhancement in refining boundaries.

[0046] Efficiency Analysis: In the Synapse multi-organ segmentation task, the computational cost of the method in this invention is only 7.65 GFLOPs, which is much lower than 48.93 GFLOPs of ScaleFormer, but it is 1.27% higher in the average Dice index, which proves that the graph clustering strategy can effectively improve the feature representation capability while significantly reducing the computational cost.

[0047] Qualitative results (see attached diagram): As shown in the accompanying drawings, when processing blurred boundaries (ISIC data) caused by uneven illumination or adhesion boundaries (Synapse data) caused by low contrast between organs, the segmentation results generated by this invention have smoother contours and can accurately delineate the fine structure of irregular lesions.

Claims

1. An automatic medical image segmentation method based on an online clustering attention mechanism guided by K-nearest neighbor graphs, characterized by the following steps: Step 1: Based on the existing medical image dataset, preprocess the data, including image size normalization and data augmentation operations. Step 2: Divide the processed data into training set, test set and validation set. Step 3: Design a medical image segmentation network framework based on the training set: Step 3.1: Construct an encoder using the Pyramid Visual Transformer (PVT) as the backbone network to extract multi-scale feature maps. Step 3.2: Construct a graph-guided clustering attention (GCA) module in the deep feature space: Treat image features as graph nodes and dynamically construct a K-nearest neighbor (KNN) graph based on feature similarity; Update node features using graph convolutional networks and expanded neighborhood aggregation strategies to learn semantic cluster centers; Use the cluster centers as compressed keys and values ​​to perform attention calculation with the query generated from the original features to obtain features containing global semantic dependencies. Step 3.3: Construct a decoder with skip connections to fuse the multi-scale features of the encoder with the upsampled features of the decoder; introduce a convolutional attention module during the fusion process to enhance local detail features through channel attention and spatial attention mechanisms. Step 3.4: Construct the Dynamic Foreground Enhancement (DFE) module: At the end of the decoder, a dynamic clustering prototype is generated using the feature map. The prototype is then projected back into the pixel space to weight the feature response, thereby enhancing the contrast between the foreground object and the background. Step 3.5: Calculate the difference between the network prediction and the true label using a hybrid loss function based on cross-entropy and Dice loss functions, and update the network parameters using gradient descent until the model converges. Step 4: Repeat Step 3 to train the segmentation network model. Step 5: Use the trained segmentation network to test the test data and output the final medical image segmentation results.

2. The automatic medical image segmentation method based on K-nearest neighbor graph-guided online clustering attention mechanism according to claim 1, characterized in that: The Graph-Directed Clustering Attention (GCA) module described in step 3.2 is implemented as follows: First, dynamic graph construction is performed, treating the input feature sequence as vertices of an undirected graph. The Euclidean distance between nodes is calculated, and the K nearest neighbors are selected to construct the graph structure. Second, dilated graph convolutional feature updates are performed, using a dilated KNN strategy with an dilation rate of d to expand the receptive field, and the maximum relative graph convolution algorithm is used to capture boundary gradient changes. Third, semantic clustering and key-value compression are performed, generating a clustering assignment matrix through convolution and the Softmax function. This matrix is ​​then used to weight and aggregate the original input to obtain M semantic cluster center features. Finally, linear attention computation is performed, mapping the original features to queries and the cluster centers to keys and values. Multi-head self-attention computation is then executed, reducing the computational complexity to a linear level.

3. The automatic medical image segmentation method based on K-nearest neighbor graph-guided online clustering attention mechanism according to claim 1, characterized in that: The convolutional attention module described in step 3.3 is specifically implemented as follows: the deep global semantic features extracted by online clustering attention guided by the K-nearest neighbor graph are added to the backbone network features of the back-connected and skip connections, and then a lightweight convolutional attention module is introduced in each decoding stage; this module first calculates the max pooling and average pooling of the feature map in parallel, and then generates channel attention weights through a shared multilayer perceptron. Spatial attention weights are then generated through convolutional operations in the spatial dimension; the features are weighted using channels and spatial attention weights, which enhances the model's ability to perceive the details of small lesion edges.

4. The automatic medical image segmentation method based on K-nearest neighbor graph-guided online clustering attention mechanism according to claim 1, characterized in that: The dynamic foreground enhancement module described in step 3.4 is specifically implemented by mapping the feature map to the prototype space at the highest resolution layer of the decoder and calculating the clustering prototype through the autocorrelation matrix. The cluster prototype is compressed and normalized using Softmax to obtain a global semantic descriptor. The global semantic descriptor is then projected back into the pixel space and element-wise multiplied with the original features. A dynamic gating mechanism is used to suppress background noise response and enhance foreground lesion features.

Images