Self-supervised contrastive learning and frequency domain feature enhancement for angiography image segmentation

Abstract

The application discloses a method for angiogram image segmentation based on self-supervised contrast learning and frequency domain feature enhancement. The method takes UNet as the basis, inputs the original angiogram image into the symmetrical encoder-decoder structure, integrates the discrete wavelet transform path in the encoding stage, injects the high-frequency texture component into the down-sampling feature of the encoder, realizes the detail compensation of the slender branch of the blood vessel, and connects the projection layer at the end of the bottleneck layer of the encoder to perform feature mapping. In the first stage, the decoder is frozen, self-supervised pre-training is performed by minimizing the contrast loss, and the feature capturing ability of the encoder is strengthened. In the second stage, the pre-training weight is loaded and the decoder is unfrozen, the Dice similarity coefficient and the binary cross entropy loss are combined for supervised fine-tuning, and accurate blood vessel segmentation is realized. The method can deeply fuse the global semantic topology and the local detail information, significantly improve the segmentation performance and edge accuracy, and maintain high calculation speed and low resource consumption.

Description

Technical Field

[0001] This invention belongs to the field of image processing technology, and relates to medical image segmentation, specifically to a method for segmenting angiography images using self-supervised contrastive learning and frequency domain feature enhancement. Background Technology

[0002] Computer-aided medical image analysis plays a crucial role in assisting clinical diagnosis and optimizing treatment planning. In particular, medical image segmentation technology, as a core pillar of the computer-aided diagnostic system, can achieve high-precision identification and anatomical localization of lesion areas, helping clinicians to quickly locate targets.

[0003] In the field of medical image segmentation, deep learning-based models have demonstrated excellent segmentation capabilities and robustness. Among them, U-Net, based on traditional convolutional neural networks (CNNs), employs a typical symmetric encoder-decoder structure and significantly improves segmentation capabilities by fusing deep semantic and shallow detail information through skip connections. While the existing U-Net architecture performs well for most medical image segmentation scenarios, for more complex images and with less labeled data, U-Net's downsampling leads to loss of detail, high dependence on labeled data, and poor generalization ability.

[0004] To address the issue of information loss, some methods introduce wavelet transform on top of the U-Net network, using wavelet transform to replace traditional pooling operations. This fuses extracted high-frequency details with encoder features, resolving the information loss problem. However, since wavelet downsampling is essentially a linear transform, it lacks the selective filtering capability of max pooling. This weakens the local translation invariance and robustness obtained by convolutional neural networks through nonlinear pooling to some extent. Consequently, when processing medical images containing complex artifacts or high-frequency background noise, the model becomes overly sensitive to interference signals, struggles to effectively suppress responses in non-critical regions, and suffers from poor segmentation accuracy. Summary of the Invention

[0005] To address the shortcomings of existing technologies, this invention proposes a self-supervised contrastive learning and frequency domain feature enhancement method for angiography image segmentation. By introducing a multi-level wavelet decomposition layer in the encoding stage, high-frequency information is effectively preserved, and a self-supervised contrastive learning module is constructed to enhance the discriminative expression of features. This reduces the dependence on labeled data while improving segmentation accuracy, solving the problems of insufficient capture of medical image details and excessive model dependence on labeled data.

[0006] A self-supervised contrastive learning and frequency domain feature enhancement method for angiography image segmentation includes the following steps:

[0007] Step 1: Use angiography images as samples and labeled real segmentation mask images as labels to construct a segmentation training set.

[0008] Step 2: Construct a U-shaped neural network CW-UNet based on contrastive learning and wavelet transform. On the basis of the traditional symmetric encoder-decoder structure of U-Net, a parallel wavelet feature extraction module is introduced at the encoder end, and a contrastive learning projection layer is introduced at the bottleneck layer.

[0009] Step 3: Compare the CW-UNet constructed in Step 2 with the CW-UNet constructed in Step 2 and perform two-stage training and parameter optimization.

[0010] In the first stage, the wavelet feature extraction module and decoder are frozen, and two different random data augmentations are performed on the same input sample in the training set to generate positive sample pairs. After multi-level downsampling by the encoder, the bottleneck layer features are projected to a low-dimensional space by contrastive learning of the projection head of the projection layer, and the contrastive loss is calculated. This forces the encoder to learn the feature invariance of the samples in the case of no label, and obtains pre-trained encoder weights with strong discriminative power.

[0011] In the second stage, the wavelet feature extraction module and decoder are unfrozen, while the contrastive learning projection layer is frozen. Training is performed using a segmentation training set. A recursive discrete wavelet transform is then performed through the wavelet feature extraction module to fuse the extracted high-frequency texture components with the corresponding layer's encoded features in real time. The segmentation loss is calculated, and the network parameters are adjusted.

[0012] Step 4: Input the angiography image to be processed into the CW-UNet network trained in two stages, process it through the activation function at the end of the decoder, and output a high-precision pixel-level blood vessel segmentation mask.

[0013] The present invention has the following beneficial effects:

[0014] 1. By using contrastive learning to pre-train the encoder, the model already possesses strong feature recognition capabilities and robustness against noise interference before entering the downstream segmentation task.

[0015] 2. A wavelet feature extraction module parallel to encoder downsampling was designed to extract sample texture and edge details from the high-frequency components obtained by wavelet transformation. This effectively compensates for the irreversible feature loss in the continuous downsampling process of traditional U-Net and preserves the original pooling layer. This allows the network to learn and select effective signals autonomously. While maintaining the integrity of the blood vessel topology, it significantly enhances the ability to suppress background interference and achieves a balance between detail recovery and model robustness.

[0016] 3. By adopting the above technical solution, the network is built on the lightweight and high-performance U-Net benchmark framework. It continues the classic symmetric encoder-decoder design, ensuring the efficiency and stability of the system. At the same time, it can deeply integrate global semantic topology and local detail information, which significantly improves segmentation performance and edge accuracy while maintaining extremely high computing speed and low resource consumption. Attached Figure Description

[0017] Figure 1 This is a schematic diagram of the CW-UNet network structure constructed in the embodiment.

[0018] Figure 2 This is a schematic diagram of the wavelet feature extraction module.

[0019] Figure 3 A schematic diagram of the projection layer structure is shown for comparative study.

[0020] Figure 4 For comparison, a learning and training flowchart is provided.

[0021] Figure 5 The flowchart for training the segmentation task. Detailed Implementation

[0022] The present invention will be further explained and described below with reference to the accompanying drawings. It should be understood that the present invention covers any substitutions, modifications, equivalent methods, and solutions made within the spirit and scope of the invention as defined in the claims. Furthermore, to provide the public with a better understanding of the present invention, certain specific details are described in detail below. However, those skilled in the art will fully understand the present invention even without these detailed descriptions.

[0023] A self-supervised contrastive learning and frequency domain feature enhancement method for angiography image segmentation includes the following steps:

[0024] Step 1: This embodiment uses a DSA angiography dataset as the original data source. Considering the complex backgrounds, thin and elongated vessels, and varying contrast in angiography images, the original images undergo quality enhancement and size normalization processing:

[0025] s1.1 Read the original DSA image and the corresponding expert-annotated mask, and convert it to single-channel grayscale mode.

[0026] s1.2. Randomly divide the original dataset into training, validation and test sets, and fix the random number seed to ensure the reproducibility of the experiment.

[0027] s1.3. Perform data augmentation operations on the training set images, including brightness shift, contrast change, saturation shift, and Gaussian blur, to generate positive sample pairs. Combine different images into negative sample pairs, which are used as training data for the first stage of contrastive learning. Continue to use spatial transformation data augmentation methods on the samples and masks as training data for the second stage of segmentation training.

[0028] Step 2, the construction of this embodiment is as follows: Figure 1 The U-shaped neural network CW-UNet shown here is based on contrastive learning and wavelet transform. This network adopts a symmetrical encoder-decoder architecture, and introduces a parallel wavelet feature extraction module at the encoder end, while introducing a contrastive learning projection layer at the bottleneck layer.

[0029] s2.1 The encoder performs a 5-layer encoding operation on the input sample, extracting multi-scale semantic features through downsampling convolution. Addressing the characteristics of long, thin blood vessels and easily lost edges in the DSA dataset, a wavelet feature extraction module is introduced in layers 2-4 for feature fusion, injecting high-frequency wavelet components into the encoding path to supplement details.

[0030]

[0031] Where m represents the hierarchical index of the encoding operation. Indicates the first Output characteristics of the layer encoder This represents the input sample. Indicates downsampling convolution. This represents the high-frequency detail feature components output by the wavelet feature extraction module at the corresponding level. Tensor splicing representing the channel dimension. For activation function, For batch normalization, This represents a convolution operation. The encoder can sequentially output five layers of encoded feature maps from shallow to deep, providing rich local details and global semantic information for subsequent segmentation.

[0032] like Figure 2 As shown, the wavelet feature extraction module uses discrete wavelet transform to perform a three-level recursive decomposition of the input samples. At each level, the image features are decomposed into low-frequency approximation (LL) and high-frequency details. The extracted high-frequency details include horizontal details (LH), vertical details (HL), and diagonal details (HH) components, which are used to capture the texture components of the image. After convolution, these components are input into the corresponding level of the encoder as high-frequency detail feature components. The low-frequency approximation component then proceeds to the next level of recursion:

[0033]

[0034] in, This represents the activation function.

[0035] s2.2, A contrastive learning projection layer is inserted into the bottleneck layer at the encoder end for contrastive learning training, such as... Figure 3 As shown, the bottleneck layer feature map is first processed using adaptive pooling. Transform into global feature vector Subsequently, it is mapped to a low-dimensional vector through a multilayer perceptron incorporating random deactivation and batch normalization. This is used to calculate the contrastive loss to enhance the discriminative power of the features:

[0036]

[0037]

[0038]

[0039]

[0040] in, ( ) indicates a flattened layer. ( ) indicates a random deactivation operation. () represents the ReLU() activation function. ( ) represents a pooling layer. ( ) represents a linear layer. This represents the eigenvector.

[0041] s2.3 The decoder performs 5 layers of upsampling on the bottleneck layer feature map and fuses the multi-scale features provided by the encoder through skip connections. The concatenated feature map undergoes convolution, normalization, and transposed convolution operations, finally outputting a predicted blood vessel segmentation mask map with the same dimension as the input sample. .

[0042] Step 3: Perform two-stage sequential training on the CW-UNet network constructed in Step 2. The specific process is as follows:

[0043] s3.1, such as Figure 4 As shown, the wavelet feature extraction module and decoder are frozen, the encoder and contrastive learning projection layer are activated, and unlabeled enhanced images are used to... Self-supervised training is performed while the decoder is frozen, by calculating and minimizing the InfoNCE contrastive loss. To enhance the encoder's ability to capture invariance to vascular features:

[0044]

[0045] in, This represents the low-dimensional representation of the positive sample pairs output by the projection head. This represents the low-dimensional representation of the negative sample pairs output by the projection head. This is a temperature hyperparameter used to control the smoothness of the similarity distribution; This indicates the calculation of cosine similarity. Indicates the number of samples.

[0046] After training is complete, the pre-trained encoder weights are output and saved.

[0047] s3.2, such as Figure 5 As shown, the encoder weights trained in S3.1 are loaded, and the decoder and wavelet feature extraction module are activated. The contrastive learning projection layer is frozen. Labeled images are used to segment and fine-tune the input network, and the Dice loss is calculated. With BCE loss Joint loss function Fine-tuning of all network parameters:

[0048]

[0049]

[0050]

[0051] in, , For the corresponding loss weights, The actual label pixel value, denoted as the network prediction probability value, and N is the number of pixels in the input sample.

[0052] Step 4: Compare the predicted segmentation mask of the angiography image output by the CW-UNet network with the corresponding real segmentation mask image of the sample, and perform segmentation processing on the next image. After all images have been processed, perform the next round of prediction. After all rounds are completed, the final prediction model is obtained.

[0053] For the FS-CAD and Kaggle Annotated X-Ray Angiography Datasets, after 100 training epochs in both the first and second stages, the average Dice scores were approximately 78.1 and 78.2, respectively. This indicates an average overlap of 78.1% and 78.2% between the predicted and ground truth images on the two datasets. The average intersection-union ratios were 64.2 and 64.3, suggesting that the intersection of the predicted and ground truth regions accounted for 64.2% and 64.3% of their union, respectively, on the two datasets.

[0054] Furthermore, based on the traditional U-Net, a wavelet transform-based U-shaped neural network is established to replace the max-pooling operation in the encoder, and comparative experiments are conducted. Using the same original dataset and data-augmented dataset, the network is trained for 200 epochs respectively, with average Dice scores of approximately 76.3 and 76.8, and average intersection-over-union ratios of 61.7 and 62.4. This performance degradation indicates that directly replacing max-pooling in the encoder with wavelet transform may lack the flexibility of feature selection due to the linear transformation nature of its fixed weights, leading to the loss of edge and detail information and diluting core semantic features. In contrast, this method retains the traditional downsampling mechanism while utilizing high-frequency components to supplement lost detail information, enabling the model to accurately capture vascular features even under complex background noise interference.

[0055] The embodiments of the present invention have been described in detail above with reference to the accompanying drawings, but the present invention is not limited to the described embodiments. For those skilled in the art, various changes, modifications, substitutions, and variations can be made to these embodiments, including components, without departing from the principles and spirit of the present invention, and these variations still fall within the protection scope of the present invention.

Claims

1. A self-supervised contrastive learning and frequency domain feature enhancement method for segmenting angiographic images, using a U-Net network to segment vascular masks from angiographic images. The U-Net network includes an encoder consisting of L downsampling layers and a decoder consisting of L upsampling layers, as well as a bottleneck layer connecting the encoder and decoder. There are skip connections between corresponding layers of the encoder and decoder. The method is characterized by: Parallel wavelet feature extraction modules are introduced in layers 2 to L-1 of the encoder. After downsampling the input features, the encoder fuses them with the high-frequency detail feature components output by the wavelet feature extraction modules as the encoding output of this layer. A contrastive learning projection layer is introduced in the bottleneck layer to map the bottleneck layer features to a low-dimensional space. Design a two-stage training and parameter optimization approach; In the first stage, positive sample pairs are constructed using data augmentation images of the same sample, and negative sample pairs are constructed using different samples. The encoder performs multi-level downsampling on the sample pairs, and then projects them to a low-dimensional space through a contrastive learning projection layer. The contrast loss between the features of the sample pairs is calculated, and the encoder is trained to learn the feature invariance of the samples in the absence of labels, thus obtaining the pre-trained encoder weights. In the second stage, the wavelet feature extraction module and decoder are unfrozen, and training continues using the segmentation training set based on the pre-trained encoder weights. The high-frequency texture components extracted by the wavelet feature extraction module are fused with the corresponding level's encoded features in real time, the segmentation loss is calculated, and the network parameters are adjusted.

2. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 1, characterized in that: The data enhancements include brightness shift, contrast variation, saturation shift, and Gaussian blur.

3. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 1, characterized in that: The encoder comprises five consecutive encoding operations, extracting multi-scale semantic features through downsampling convolution: Where m represents the hierarchical index of the encoding operation; Indicates the first Output characteristics of the layer encoder Indicates the input sample; Indicates downsampling convolution. This represents the high-frequency detail feature components output by the wavelet feature extraction module at the corresponding level; Tensor splicing representing the channel dimension. For activation function, For batch normalization, This represents the convolution operation.

4. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 3, characterized in that: The wavelet feature extraction module uses discrete wavelet transform to perform a three-level recursive decomposition of the input samples. At each level, the image features are decomposed into low-frequency approximate components and high-frequency detail components. The high-frequency detail components are then concatenated and convolved before being input into the corresponding level of the encoder as high-frequency detail feature components. The low-frequency approximation component then proceeds to the next level of recursion: in, , , These represent the high-frequency horizontal detail components, vertical detail components, and diagonal detail components obtained from the recursive decomposition in section m-1, respectively. This represents the activation function.

5. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 1, characterized in that: The contrastive learning projection layer first uses adaptive pooling. ( ) will feature map of the bottleneck layer Transform into global feature vector It is then mapped to a low-dimensional vector through a multilayer perceptron. : in, ( ) indicates a flattened layer. ( ) indicates a random deactivation operation. () represents the ReLU() activation function. ( ) represents a linear layer. This represents the eigenvector.

6. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 1, characterized in that: Set the loss function for the first stage as follows: : in, A low-dimensional representation of positive sample pairs. A low-dimensional representation of negative sample pairs; This refers to temperature hyperparameters. This indicates the calculation of cosine similarity. Indicates the number of samples.

7. The angiography image segmentation method based on self-supervised contrastive learning and frequency domain feature enhancement as described in claim 1, characterized in that: Set the loss function for the second stage as follows: : in, This indicates Dice's loss. Indicates BCE loss, , This represents the corresponding loss weight.

8. The self-supervised contrastive learning and frequency domain feature enhancement method for angiography image segmentation as described in any one of claims 1 to 7, characterized in that: The angiography image is input into the trained network, and the high-frequency detail feature components are decomposed by the wavelet feature extraction module. These components are then input into the encoder and fused with the downsampled features. Finally, the decoder is used to gradually recover the segmented blood vessel mask.

9. A computer-readable storage medium having a computer program stored thereon, which, when executed in a computer, causes the computer to perform the method of any one of claims 1 to 7.

Images