Classification assistance-based U-shaped network few-sample image segmentation training method and system
By employing a classification-assisted U-shaped network training method, which utilizes image-level labels for pre-training and a small number of pixel-level labels for fine-tuning, the problem of insufficient model training and poor generalization ability caused by data scarcity in medical image segmentation of U-shaped networks is solved, achieving high-precision and efficient image segmentation.
Patent Information
- Application Number
- CN202512009111.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-12-29
- Publication Date
- 2026-02-24
AI Technical Summary
Existing U-shaped networks suffer from insufficient model training and poor generalization ability in medical image segmentation due to the scarcity of pixel-level labeled data.
By constructing a U-shaped network model, pre-training with image-level classification labels, and combining a multi-scale feature fusion classification module to optimize model parameters, the model is then fine-tuned using a small number of pixel-level segmentation labels. A self-supervised consistency loss and dynamic learning rate scheduling strategy are adopted to improve the model's segmentation performance and generalization ability.
With very limited pixel-level labeled data, the model significantly improved segmentation accuracy and generalization ability, reduced the risk of overfitting, and achieved efficient medical image segmentation.
Smart Images

Figure CN121564516A_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of image segmentation technology, specifically to a method and system for training few-shot image segmentation using a classification-assisted U-shaped network. Background Technology
[0002] In existing technologies, U-shaped network structures, represented by U-Net, have become the mainstream method in medical image segmentation due to their encoder-decoder architecture and skip connection design, which effectively fuses multi-scale features. However, the training of such models heavily relies on a large number of accurate pixel-level segmentation labels. In practical applications such as medical imaging, obtaining such finely labeled data is costly and time-consuming, resulting in an extremely limited number of high-quality segmentation samples available for training.
[0003] To address the challenge of scarce labeled data, existing technologies have explored various approaches. One approach is to directly improve the U-Net network structure, such as optimizing the encoder-decoder or skip connections, to enhance feature extraction and fusion capabilities. However, this does not fundamentally reduce the demand for pixel-level labeled data. Another approach is to introduce paradigms such as few-shot learning, semi-supervised learning, or weakly supervised learning, attempting to utilize unlabeled, incomplete, or inaccurate labeled data. For example, few-shot or semi-supervised learning primarily utilizes a large number of unlabeled images; while weakly supervised learning focuses on utilizing more readily available weak labels, such as image-level classification labels (e.g., labeling only whether an image contains lesions), bounding boxes, doodles, or point annotations. Some studies have attempted to combine this type of weakly supervised information with a small number of pixel-level segmentation labels for multi-task learning, hoping to utilize these more readily available supervisory signals to assist in segmentation tasks.
[0004] However, existing methods, when using classification labels to assist segmentation, often focus on complex network structure design to achieve feature interactions between tasks. This may increase the risk of overfitting in scenarios with scarce data. Finding a simple and effective training paradigm that fully pre-trains the segmentation network backbone using a large number of classification labels, enabling it to achieve strong generalization ability with only a very small number of segmentation labels for fine-tuning, remains a practically valuable challenge. Summary of the Invention
[0005] In view of the above problems, this application provides a classification-assisted U-shaped network few-shot image segmentation training method and system to solve the technical problems of insufficient model training and poor generalization ability caused by the scarcity of pixel-level labeled data when performing image segmentation under few-shot conditions.
[0006] To achieve the above objectives, in a first aspect, a few-shot image segmentation training method based on a classification-assisted U-shaped network is provided, the method comprising:
[0007] S1: Construct a U-shaped network model, which includes an encoder and a decoder;
[0008] The encoder is used to extract multi-scale features from the input image;
[0009] The decoder is connected to the encoder via a skip connection and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales.
[0010] The decoder's outputs at at least two different scales are also connected to a multi-scale feature fusion classification module;
[0011] S2: Obtain the first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes the first sample image and the category label corresponding to the first sample image.
[0012] During the training of the U-shaped network model, the multi-scale feature fusion classification module performs the following operations:
[0013] S21: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the k-th scale output by the decoder to obtain the first feature vector, where k is an integer greater than 1;
[0014] S22: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the (k-1)th scale output by the decoder to obtain an intermediate feature vector, and perform a linear transformation on the intermediate feature vector to obtain a second feature vector;
[0015] S23: Concatenate the first feature vector with the second feature vector, and perform a linear transformation on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector;
[0016] S24: The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector;
[0017] S25: Input the fused feature vector into a fully connected classifier, map it to the category number dimension through a fully connected layer, and then output the classification prediction result of the first sample image through the SoftMax activation function;
[0018] S3: Obtain a second training set, where the second training set includes second sample images and labels for each pixel in the second sample images indicating the semantic class to which it belongs, and generate pixel-level segmentation labels. Among them, the second sample images include lesion regions, and the number of first sample images in the second training set is less than that in the first training set; Use the second training set to optimize the encoder and decoder parameters of the U-shaped network model trained in step S2 by minimizing an image segmentation loss function to obtain a final image segmentation model.
[0019] Further, in step S2, after obtaining the first training set and before using the first training set to train the U-shaped network model, the method further includes:
[0020] S201: Perform black bar insertion and scaling on the first sample images. The black bar insertion and scaling specifically include:
[0021] Obtain the height H and width W of the first sample image;
[0022] If H > W, insert black bars with a width of (H - W) / 2 on the left and right sides of the first sample image respectively to make the first sample image into a square;
[0023] If H < W, insert black bars with a height of (W - H) / 2 on the top and bottom sides of the first sample image respectively to make the first sample image into a square;
[0024] Scale the obtained square first sample image to a preset size P×P by bicubic interpolation to obtain a preprocessed image;
[0025] S202: Perform random rotation on the preprocessed image, and the rotation angle is randomly selected within a preset angle range.
[0026] Further, the image classification loss function includes a cross-entropy loss function L cls , and the calculation formula of the cross-entropy loss function L cls is:
[0027] ;
[0028] Among them, G is the true class label of the first sample image, and M cls is the classification prediction result of the first sample image output by mapping to the class number dimension through the fully connected layer and then passing through the SoftMax activation function.
[0029] Further, the image segmentation loss function includes a Dice loss function L dice , and the calculation formula of the Dice loss function L dice is as follows:
[0030] ;
[0031] Among them, M seg G represents the lesion segmentation result output by the decoder. seg Let |·| represent the number of elements in the set, and ∩ represent taking the intersection.
[0032] Furthermore, the method also includes:
[0033] After training is completed in step S2, the U-shaped network model that has been trained at this time is saved as a prior model;
[0034] In the optimization process of step S3, the U-shaped network model to be optimized is initialized with the parameters of the prior model, and used as the current model for any second sample image in the second training set, the following operations are performed:
[0035] The second sample image is simultaneously input into the current model and the prior model;
[0036] Using the decoders of the current model and the prior model, class activation maps corresponding to the second sample image are generated respectively, and the self-supervised consistency loss L between the two class activation maps is calculated. ss ;
[0037] The image segmentation loss function is supervised by segmentation loss L. sup And self-supervised consistency loss L ss The weighted summation yields the image segmentation loss function L. seg The calculation formula is as follows:
[0038] ;
[0039] in, A balance coefficient greater than 0 is used to adjust the strength of the self-supervised signal; the supervised segmentation loss L... sup This is the supervised segmentation loss calculated based on the pixel-level segmentation labels in the second training set.
[0040] Furthermore,
[0041] In steps S2 and / or S3, the AdamW optimizer is used to optimize the model parameters, and a dynamic learning rate scheduling strategy is employed. Specifically, the dynamic learning rate scheduling strategy includes: during the initial warm-up phase... After rounds of iteration, the learning rate rises to the target learning rate, and then decreases from the target learning rate to the target minimum learning rate according to the cosine function law;
[0042] The learning rate during the initial warm-up phase After several iterations, the learning rate rises to the target learning rate, and its calculation formula is as follows:
[0043] , ];
[0044] in, Indicates the first The learning rate for each iteration round. This represents the minimum learning rate. Indicates the target learning rate. This represents the total number of iterations during the warm-up phase.
[0045] The calculation formula for decreasing the learning rate from the target learning rate to the target minimum learning rate according to the cosine function law is as follows:
[0046] ];
[0047] in, This indicates the number of iterations remaining after the warm-up phase. The learning rate for each iteration round. This indicates the total number of training iterations.
[0048] Further in step S3, using the second training set, the encoder and decoder parameters of the U-shaped network model trained in step S2 are optimized by minimizing the image segmentation loss function to obtain the final image segmentation model, specifically including:
[0049] S31: Use the second training set as the current training set;
[0050] S32: Using the current training set, optimize the encoder and decoder of the U-shaped network model by minimizing the image segmentation loss function; after this round of optimization, use the U-shaped network model to perform forward inference on all second sample images in the current training set to obtain the pixel-level segmentation prediction probability map corresponding to each second sample image;
[0051] S33: For each second sample image in the current training set, perform the following operations:
[0052] S331: Based on the pixel-level segmentation prediction probability map, calculate the uncertainty score for each pixel position; the uncertainty score is calculated by the variance of the pixel-level segmentation prediction probability map under multiple random forward propagations, or by the information entropy of the pixel-level segmentation prediction probability map itself.
[0053] S332: According to the preset purification strategy, the pixel-level segmentation label corresponding to the second sample image is corrected using the uncertainty score;
[0054] S34: Update all the second sample images obtained after purification in step S33 and their corrected pixel-level segmentation labels to the new current training set;
[0055] S35: Repeat steps S32 to S34 until the preset convergence condition is met.
[0056] Furthermore, the purification strategy includes:
[0057] For pixels with uncertainty scores below the first score threshold, their original pixel-level segmentation labels are retained.
[0058] For pixels with uncertainty scores higher than the first score threshold but lower than the second score threshold, the original pixel-level segmentation label is replaced by the prediction category at the corresponding position in the pixel-level segmentation prediction probability map.
[0059] For pixels with uncertainty scores higher than the second score threshold, their original pixel-level segmentation labels are retained and they are marked as samples to be reviewed.
[0060] The preset convergence conditions include the number of iterations reaching a preset iteration threshold, or the proportion of pixels in the current training set marked as samples to be reviewed being lower than a preset proportion.
[0061] Furthermore, in step S32, the image segmentation loss function is an uncertainty-weighted loss function; using the current training set, the encoder and decoder of the U-shaped network model are optimized by minimizing the image segmentation loss function, specifically including:
[0062] Input the second sample image in the current training set into the current model to obtain the pixel-level segmentation prediction probability map and the corresponding pixel-level uncertainty score;
[0063] Based on the uncertainty score, a dynamic weight is assigned to each pixel, wherein pixels with higher uncertainty are given lower weights.
[0064] Based on the dynamic weights and the pixel-level segmentation prediction probability map, a weighted image segmentation loss function value is calculated, and the model parameters are updated accordingly; wherein, the weighted image segmentation loss function... Calculated using the following formula:
[0065] ;
[0066] in, This represents the index of all pixels in all second sample images within the current batch. Let be the class probability vector predicted by the model for the i-th pixel. Let be the one-hot encoded vector of the pixel-level segmentation label corresponding to the i-th pixel in the current training set. To calculate the basic loss function for the difference between the prediction and the label, This represents the dynamic weight corresponding to the i-th pixel.
[0067] In a second aspect, this application provides a classification-assisted U-shaped network few-shot image segmentation training system, the system being used to perform the method as described in the first aspect of this application, the system comprising:
[0068] A U-shaped network model construction module is used to construct a U-shaped network model, which includes an encoder and a decoder;
[0069] The encoder is used to extract multi-scale features from the input image;
[0070] The decoder is connected to the encoder via a skip connection and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales.
[0071] The decoder's outputs at at least two different scales are also connected to a multi-scale feature fusion classification module;
[0072] The first training module is used to acquire a first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes a first sample image and the category label corresponding to the first sample image.
[0073] During the training of the U-shaped network model, the multi-scale feature fusion classification module performs the following operations:
[0074] The decoded feature map at the k-th scale output by the decoder is subjected to convolutional dimensionality reduction, global average pooling, and flattening operations to obtain the first feature vector, where k is an integer greater than 1;
[0075] The (k-1)th scale decoded feature map output by the decoder is subjected to the convolutional dimensionality reduction, global average pooling, and flattening operations to obtain an intermediate feature vector, and the intermediate feature vector is subjected to a linear transformation to obtain a second feature vector.
[0076] The first feature vector is concatenated with the second feature vector, and a linear transformation is performed on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector.
[0077] The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector;
[0078] The fused feature vector is input into a fully connected classifier, mapped to the number of categories through a fully connected layer, and then the classification prediction result of the first sample image is output through the SoftMax activation function.
[0079] The second training module is used to obtain a second training set, which includes a second sample image and a label for each pixel in the second sample image to indicate its semantic category, generating pixel-level segmentation labels. The second sample image contains lesion regions, and the number of first sample images in the second training set is less than that in the first training set. Using the second training set, the encoder and decoder parameters of the U-shaped network model trained in step S2 are optimized by minimizing the image segmentation loss function to obtain the final image segmentation model.
[0080] Unlike existing technologies, the above-mentioned technical solution involves a few-shot image segmentation training method and system based on classification-assisted U-shaped network. The method first trains a U-shaped network model using a first training set with image-level category labels, optimizing parameters by minimizing the classification loss. During this process, a multi-scale feature fusion classification module fuses and classifies the multi-scale features output by the decoder, enabling the network to learn strong semantic features related to segmentation. Then, a small number of samples are selected from the first training set and labeled with pixel-level segmentation labels to form a second training set. Finally, with the classification module parameters fixed, the encoder and decoder parameters are optimized by minimizing the segmentation loss using only this small second training set, resulting in the final image segmentation model. This invention, through a two-stage training strategy and utilizing a large number of readily available classification labels for pre-training, effectively improves the model's segmentation performance and generalization ability when pixel-level labeled data is extremely limited.
[0081] The above description of the invention is merely an overview of the technical solution of this application. In order to enable those skilled in the art to better understand the technical solution of this application and to implement it based on the description and drawings, and to make the above-mentioned objectives and other objectives, features and advantages of this application easier to understand, the following description is provided in conjunction with the specific embodiments and drawings of this application. Attached Figure Description
[0082] The accompanying drawings are only used to illustrate the principles, implementation methods, applications, features, and effects of specific embodiments of this application and other related content, and should not be considered as limitations on this application.
[0083] In the accompanying drawings of the instruction manual:
[0084] Figure 1 This is a flowchart of the few-shot image segmentation training method based on classification-assisted U-shaped network according to the first exemplary embodiment of this application;
[0085] Figure 2 This is a flowchart of the few-shot image segmentation training method based on classification-assisted U-shaped network according to the second exemplary embodiment of this application;
[0086] Figure 3 This is a flowchart of the few-shot image segmentation training method based on classification-assisted U-shaped network according to the third exemplary embodiment of this application;
[0087] Figure 4 This is a flowchart of the few-shot image segmentation training method based on classification-assisted U-shaped network according to the fourth exemplary embodiment of this application;
[0088] Figure 5 This is a flowchart of the few-shot image segmentation training method based on classification-assisted U-shaped network according to the fifth exemplary embodiment of this application;
[0089] Figure 6 This is an overall architecture diagram of an image segmentation model according to an embodiment of this application;
[0090] Figure 7 This is a schematic diagram of a few-shot image segmentation training system based on a classification-assisted U-shaped network according to an exemplary embodiment of this application;
[0091] The reference numerals used in the above figures are explained as follows:
[0092] 10. Few-shot image segmentation training system; 101. U-shaped network model construction module; 102. First training module; 103. Second training module. Detailed Implementation
[0093] To explain in detail the possible application scenarios, technical principles, specific feasible solutions, and the objectives and effects that this application can achieve, the following detailed description is provided in conjunction with the listed specific embodiments and accompanying drawings. The embodiments described herein are only used to more clearly illustrate the technical solutions of this application, and are therefore only examples, and should not be used to limit the scope of protection of this application.
[0094] In the first aspect, such as Figure 1 As shown, this application provides a few-shot image segmentation training method based on a classification-assisted U-shaped network, the method comprising:
[0095] S1: Construct a U-shaped network model;
[0096] S2: Obtain the first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes the first sample image and the category label corresponding to the first sample image.
[0097] S3: Obtain a second training set, which includes a second sample image and a label for each pixel in the second sample image that indicates its semantic category, and generate pixel-level segmentation labels. The second sample image contains lesion regions, and the number of first sample images in the second training set is less than that in the first training set. Using the second training set, optimize the encoder and decoder parameters of the U-shaped network model trained in step S2 by minimizing the image segmentation loss function to obtain the final image segmentation model.
[0098] like Figure 6 As shown, the U-shaped network model includes an encoder and a decoder; wherein, the encoder is used to extract multi-scale features from the input image; the decoder is connected to the encoder via a skip connection, and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales; the output terminals of the at least two different scales of the decoder are also connected to a multi-scale feature fusion classification module.
[0099] Specifically, the encoder consists of multiple convolutional layers and downsampling layers. The input is a medical image (such as a grayscale image of a frontal chest X-ray). The encoder extracts low-, medium-, and high-dimensional multi-scale features of the image through layer-by-layer convolution. Low-dimensional features contain spatial information such as image details, while high-dimensional features contain contextual information such as lesion categories.
[0100] The decoder consists of multiple upsampling layers and convolutional layers, and is connected to the encoder via skip connections. On the one hand, it gradually restores the image to a preset size through upsampling, and on the other hand, it receives features from the corresponding layers of the encoder through skip connections, realizing the fusion of shallow spatial information and deep contextual information, and finally outputs at least two decoded feature maps at different scales (different scales correspond to different levels of feature abstraction and image size).
[0101] The multi-scale feature fusion classification module is used to extract multi-scale decoded features and perform fusion processing to achieve image classification tasks and provide auxiliary supervision for segmentation tasks.
[0102] like Figure 2 As shown, during the training of the U-shaped network model in step S2, the multi-scale feature fusion classification module performs the following operations:
[0103] S21: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the k-th scale output by the decoder to obtain the first feature vector;
[0104] S22: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the (k-1)th scale output by the decoder to obtain an intermediate feature vector, and perform a linear transformation on the intermediate feature vector to obtain a second feature vector;
[0105] S23: Concatenate the first feature vector with the second feature vector, and perform a linear transformation on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector;
[0106] S24: The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector;
[0107] S25: Input the fused feature vector into a fully connected classifier, map it to the category number dimension through a fully connected layer, and then output the classification prediction result of the first sample image through the SoftMax activation function.
[0108] The first training set is a dataset used for classification pre-training, containing a large number of medical first sample images and corresponding image-level category labels (no pixel-level annotation required). The category labels are used to indicate whether the image belongs to the lesion category or the normal category. Preferably, the first sample images in the first training set are medical grayscale images.
[0109] The second training set is a dataset used for fine-tuning few-shot segmentation. It includes semantic category labels (pixel-level segmentation labels) for each pixel to further segment lesion regions and non-lesion regions. The second training set contains far fewer samples than the first training set.
[0110] In step S21, a 1×1 convolution dimensionality reduction operation is performed on the decoded feature map of the k-th scale output by the decoder (k is an integer greater than 1, representing any non-lowest scale decoded feature map) to unify the number of channels of the high-dimensional features to a preset dimension, thereby reducing computational complexity. Then, a flattening operation is performed to convert the two-dimensional feature map into a one-dimensional first feature vector.
[0111] In step S22, the decoded feature map at the (k-1)th scale (the feature map at the next lower scale adjacent to the kth scale) output by the decoder is subjected to the same convolutional dimensionality reduction, global average pooling, and flattening operations as in step S21 to obtain an intermediate feature vector. Since the dimensions of features at different scales differ, the intermediate feature vector is linearly transformed (through a fully connected layer) to convert it into a second feature vector that matches the dimension of the first feature vector.
[0112] In step S23, the first feature vector and the second feature vector are concatenated to obtain a concatenated feature vector with a dimension equal to the sum of the two. The concatenated feature vector is then subjected to a linear transformation (through a fully connected layer) to compress it into a third feature vector with the same dimension as the first feature vector. This preserves multi-scale feature information while ensuring the consistency of feature dimensions, which facilitates feature fusion at subsequent levels.
[0113] In step S24, by repeating the operations of steps S21-S23 above until all scale output features of the decoder have been processed, the third feature vector from the highest scale decoding feature map is obtained. The highest scale decoding feature map corresponds to the feature with the highest resolution and the size closest to the original input image, containing the richest spatial detail information. Then, the encoded feature map output by the encoder after the last downsampling (containing the most abstract deep context information) is compressed so that its size is in the same dimension as the third feature vector. Then, the compressed encoded feature map and the third feature vector are fused (by means of element concatenation, element addition, etc.) to obtain a fused feature vector that combines multi-scale decoding features and deep encoding features.
[0114] In step S25, the fused feature vector is input to the fully connected classifier, mapped to the category number dimension through the fully connected layer, and then output as the classification prediction probability distribution (i.e., classification prediction result) of the first sample image through the SoftMax activation function.
[0115] Through classification pre-training, the encoder and decoder, under image-level label supervision, not only learn high-level semantic features but also spontaneously form spatial attention priors for lesion regions driven by the classification task. This allows the network to focus on boundary refinement directly based on the established semantic and spatial relationships during subsequent few-shot segmentation fine-tuning, without having to learn lesion localization from scratch. This mechanism significantly reduces the learning difficulty under few-shot conditions and improves the model's convergence speed and segmentation accuracy.
[0116] In step S3, the U-shaped network model trained in step S2 is fine-tuned using the second training set. By minimizing the image segmentation loss function, only the parameters of the encoder and decoder are optimized (while the parameters of the multi-scale feature fusion classification module are fixed). This further improves the pixel-level segmentation accuracy based on the model's already mastered feature extraction and localization capabilities, ultimately yielding the image segmentation model. Since the model has already acquired strong feature extraction and lesion localization capabilities through classification pre-training, it can converge quickly with only a small amount of pixel-level labeled data, achieving high-precision segmentation and effectively solving the overfitting problem in scenarios with few samples.
[0117] Specifically, for the i-th stage feature output by the decoder... ,in, For the stage Number of feature channels, For the stage The feature width and height are first adopted. Bottleneck convolution and global average pooling compress high-dimensional abstract features, reducing the number of channels and accelerating training convergence. The resulting flattened feature vectors are then used to generate feature vectors. ( (For the preset feature vector length). Then the feature vector Features compressed by the next stage fully connected layer Perform splicing to generate fused features Then, by compressing features through linear layers, a compressed vector of the same length as the output vector of each stage is generated. This enables the step-by-step fusion of multi-scale features.
[0118] The feature vectors obtained through multiple linear layer compressions are concatenated with the intermediate features of the U-shaped network (deep features output by the encoder) after compression by fully connected layers. This concatenation is then input into a classifier consisting of fully connected layers and a SoftMax activation function, and the output is the classification prediction result. .
[0119] The above scheme injects rich prior semantic knowledge into the segmentation task through classification pre-training, enabling the model to achieve performance close to or exceeding that of traditional methods trained with a large number of segmentation labels when fine-tuned with only a few segmentation labels, effectively overcoming overfitting. By forcibly fusing multi-scale features from the decoder and deep features from the encoder through a multi-scale feature fusion classification module, the encoder and decoder jointly learn discriminative features and localization capabilities during the pre-training phase, improving the feature representation ability of the entire network.
[0120] In some embodiments, such as Figure 3 As shown, in step S2, after obtaining the first training set and before training the U-shaped network model using the first training set, the method further includes:
[0121] S201: Perform black bar insertion and scaling on the first sample image;
[0122] S202: Perform random rotation processing on the preprocessed image, with the rotation angle randomly selected within a preset angle range.
[0123] In step S201, the black bar insertion and scaling process specifically includes:
[0124] Obtain the height H and width W of the first sample image;
[0125] If H > W (image height is greater than width), then black bars with a width of (H - W) / 2 are inserted on the left and right sides of the first sample image respectively (the black bars are zero-valued pixels and do not affect the original image information), so that the aspect ratio of the inserted image becomes 1:1 (i.e., a square); the size of the black bars is H × [(H - W) / 2] × 1 (consistent with the height and number of channels of the first sample image), ensuring the integrity of the inserted image.
[0126] If H < W (image width is greater than height), then black bars with a height of (W - H) / 2 are inserted on the top and bottom sides of the first sample image respectively to make the image a square; the size of the black bars is [(W - H) / 2] × W × 1, consistent with the width and number of channels of the first sample image.
[0127] For the obtained square first sample image, it is scaled to the preset size P × P using bicubic interpolation to obtain a preprocessed image; bicubic interpolation can maintain the detailed features of the image during the scaling process by comprehensively considering the gray information of adjacent pixels, reduce image blurring or distortion caused by scaling, and ensure the image quality input to the model.
[0128] In step S202, the preprocessed image obtained in step S201 (with a size of P × P) is randomly rotated: with the center of the image as the center point, a rotation angle is randomly selected within the preset angle range (such as [-30°, 30°]) to perform a rotation operation on the image; the rotated image still maintains the size of P × P, and the image size is ensured to remain unchanged by filling zero-valued pixels. This operation artificially increases the diversity of training samples without changing the semantic content of the image, simulates the image angle differences caused by slight changes in the patient's body position during medical image acquisition, forces the model to learn rotation-invariant features, improves the robustness of the model to image angle changes, and helps to alleviate the overfitting problem of the model.
[0129] In some other embodiments, the preprocessing operations further include brightness adjustment, contrast adjustment, Gaussian noise addition, etc., to further improve the generalization ability of the model.
[0130] The above solution standardizes the input image size (uniformly to P × P, and the value of P is preferably 256) through black bar insertion and scaling processing, while maximizing the retention of all information of the original image, especially protecting the lesion areas that may be located at the edges of the image, and avoiding the loss of key information caused by cropping or stretching. The diversity of training samples is enhanced through random rotation, enabling the model to learn more essential and general visual patterns during the pre-training stage rather than relying on specific spatial orientations, effectively improving the adaptability of the model to image angle changes.
[0131] In some embodiments, the image classification loss function includes the cross-entropy loss function L clsCross-entropy loss function L cls The calculation formula is:
[0132] ;
[0133] Where G is the true class label of the first sample image, and M cls The classification prediction result of the first sample image is output after mapping to the number of categories through a fully connected layer and then passing through the SoftMax activation function.
[0134] In this embodiment, the true class label (G) refers to the image-level class label corresponding to the first sample image in the first training set.
[0135] Classification prediction results (M) cls ) refers to the probability distribution of the first sample image category prediction output by the fully connected classifier. The dimension is N (consistent with the number of categories). Each element represents the probability that the sample belongs to the corresponding category. The sum of all elements is 1 (processed by the SoftMax activation function).
[0136] Specifically, for a single first sample image, the true class label G is a one-hot encoded vector, with only the position corresponding to the true class having a value of 1, and the rest having a value of 0. Therefore, the calculation of cross-entropy loss essentially only focuses on the probability value corresponding to the true class in the predicted probability distribution. That is, the closer the predicted probability is to 1, the closer the cross-entropy loss is to 0, and the closer the loss value is to 0, the more accurate the classification prediction is. When the predicted probability is closer to 0, the cross-entropy loss increases sharply, and the larger the loss value is, the more serious the classification prediction error is.
[0137] During the classification pre-training phase, the model outputs a classification prediction result M at each forward propagation. cls Then, calculate L. cls The value of the loss is used to calculate the partial derivative of the loss value with respect to all trainable parameters of the model through backpropagation, thus obtaining the gradient of each parameter. The optimizer (such as AdamW) updates the model parameters based on the gradient information, so that the predicted probability distribution M of the model output is optimized. cls By continuously approaching the true label distribution G, the model's classification ability is ultimately improved.
[0138] The above scheme, through cross-entropy loss, provides a clear and differentiable optimization objective for the classification pre-training stage. This efficiently drives the multi-scale feature fusion classification module and the entire U-shaped network backbone (encoder-decoder) to learn the correct class discrimination boundaries, enabling the model to quickly master image class discrimination capabilities. The gradient calculation of the cross-entropy loss function avoids the gradient saturation problem that may arise from mean squared error loss (especially when combined with the SoftMax activation function), providing a stable gradient flow in both the early and late stages of training, ensuring fast and stable model convergence.
[0139] Furthermore, to improve the model's generalization ability during the classification pre-training stage and prevent overfitting, label smoothing techniques can be used to process the category labels corresponding to the first sample images when training on the first training set. The smoothed labels are denoted as... The calculation formula is as follows:
[0140] ;
[0141] in, This represents the component of the c-th category in the smoothed label vector. This represents the smoothing coefficient (usually around 0.1). The component of the c-th class bit in the one-hot encoded label vector (the true class is 1, and the rest are 0). Let L represent the total number of categories, and c represent the category index. In the cross-entropy loss function L... cls At that time, based on the smoothed label And the classification prediction result M cls Perform the calculation.
[0142] In some embodiments, the image segmentation loss function includes the Dice loss function L. dice Dice loss function L dice The calculation formula is as follows:
[0143] ;
[0144] Among them, M seg G represents the lesion segmentation result output by the decoder. seg Let |·| represent the number of elements in the set, and ∩ represent taking the intersection.
[0145] In this embodiment, the Dice loss function (L dice ) is a loss function in medical image segmentation tasks, constructed based on the Dice similarity coefficient (DSC), used to measure the degree of regional overlap between the segmentation prediction result and the true segmentation label.
[0146] Lesion segmentation results (M) seg () refers to the pixel-level segmentation prediction result output by the decoder, which is a feature map with the same size as the input image. The value of each pixel represents the probability or category label of that pixel belonging to the lesion category.
[0147] Pixel-level segmentation labels (G seg ) refers to the true segmentation label corresponding to the first sample image in the second training set. It is a binary or multi-valued image with the same size as the input image. The value of each pixel indicates the semantic category to which the pixel belongs (e.g., lesion pixels are 1, background pixels are 0).
[0148] Specifically, when the segmentation prediction result M seg Compared with the true segmentation label G seg When fully overlapping (i.e.) ), DSC=1, L dice =0 indicates perfect segmentation; when the segmentation prediction result M seg Compared with the true segmentation label G seg The lower the degree of overlap, the smaller the DSC, and the lower the L. dice The larger the value, the greater the difference between the segmentation prediction and the actual situation. Unlike the cross-entropy loss, which calculates the loss pixel by pixel in classification tasks, the Dice loss is a region-level loss function. It focuses more on the consistency of the overall shape and position of the predicted region and the real region, and is relatively insensitive to the classification errors of individual pixels within the region. This makes it more suitable for the accuracy requirements of lesion region localization in medical image segmentation. In the few-shot segmentation fine-tuning stage, the model forward propagation outputs the lesion segmentation result M. seg Then, calculate L. dice The value of is used to obtain the gradient through backpropagation and update the parameters of the encoder and decoder, so that the segmentation prediction results of the model continuously approach the true segmentation label, thereby improving the pixel-level segmentation accuracy.
[0149] In medical images, lesion (foreground) regions typically occupy only a small portion of the image, while background pixels constitute the vast majority. Traditional pixel-wise loss functions tend to cause models to favor predicting the background category. Dice loss automatically balances the contributions of the foreground and background through region overlap calculations, allowing model optimization to focus more accurately on segmenting the target lesion region, effectively alleviating the training difficulties caused by the small number of foreground pixels. Because Dice loss focuses on the overall overlap of regions, the optimization process directly encourages the predicted region boundaries to converge towards the true boundaries, contributing to the generation of clearer and more complete segmentation results.
[0150] In some embodiments, the method further includes:
[0151] After training is completed in step S2, the U-shaped network model that has been trained at this time is saved as a prior model;
[0152] In the optimization process of step S3, the U-shaped network model to be optimized is initialized with the parameters of the prior model, and used as the current model for any second sample image in the second training set, the following operations are performed:
[0153] The second sample image is simultaneously input into the current model and the prior model;
[0154] Using the decoders of the current model and the prior model, class activation maps corresponding to the second sample image are generated respectively, and the self-supervised consistency loss L between the two class activation maps is calculated. ss ;
[0155] The image segmentation loss function is supervised by segmentation loss L. sup And self-supervised consistency loss L ss The weighted summation yields the image segmentation loss function L. seg The calculation formula is as follows:
[0156] ;
[0157] in, A balance coefficient greater than 0 is used to adjust the strength of the self-supervised signal; the supervised segmentation loss L... sup This is the supervised segmentation loss calculated based on the pixel-level segmentation labels in the second training set.
[0158] In this embodiment, the prior model refers to the U-shaped network model saved after the classification pre-training stage. It has been trained with a large amount of classification data and has obtained strong category discrimination ability and preliminary lesion localization ability. Its parameters serve as the initial parameters and prior knowledge source for subsequent fine-tuning.
[0159] Category activation maps (CAMs) are visual feature maps generated by the model's decoder that reflect the image regions on which the model makes classification decisions. These regions typically correspond to the location of lesions and contain coarse localization information about the lesions.
[0160] Cross-stage self-supervised signals refer to the supervision information constructed based on the category activation maps output by the prior model and the current model. No additional manual annotation is required. It is used to constrain the current model to maintain its ability to locate lesion areas during fine-tuning and is a form of self-supervised learning.
[0161] Self-supervised consistency loss (L ss ) refers to a loss function (such as mean squared error loss) used to measure the difference between the class activation maps output by the current model and the prior model, forcing the two activation maps to maintain consistency.
[0162] Supervision of segmentation loss (L) sup The segmentation loss (such as the Dice loss mentioned above) is calculated based on the pixel-level segmentation labels in the second training set, providing the main supervision signal for the fine-tuning stage.
[0163] Specifically, after the classification pre-training stage is completed, the trained U-shaped network model is saved as a prior model; the prior model is used to perform forward propagation processing on all first sample images in the first training set, and the prior localization knowledge learned by the prior model on the first training set is saved.
[0164] During the model optimization process in the few-shot segmentation fine-tuning stage, the U-shaped network model to be optimized is initialized with the parameters of the prior model and used as the current model. For each second sample image in the second training set, the following operations are performed: the second sample image is simultaneously input into the current model and the prior model; the decoders of the current model and the prior model generate class activation maps corresponding to the sample image, denoted as CAM. current and CAM prior The brighter areas in the category activation maps are the key regions (usually lesion areas) that the model uses to determine the category of the image. These category activation maps contain the model's learned coarse lesion localization knowledge, which is stored as a source of cross-stage self-supervised signals.
[0165] Then calculate CAM current With CAM prior The difference between them yields the self-supervised consistency loss L. ss Then, an image segmentation loss function is constructed. , where L sup The supervised segmentation loss (such as Dice loss) is calculated based on the pixel-level segmentation labels of the second training set, where λ is a balancing coefficient used to adjust L. ss The intensity of the influence (can be set to a value between 0.1 and 1.0 based on experimental results); with L seg To optimize the target, the encoder and decoder parameters of the current model are updated through backpropagation. While optimizing the segmentation accuracy, the class activation map of the current model is forced to be consistent with the prior model, ensuring that the model does not forget the lesion localization knowledge learned in the pre-training stage.
[0166] When fine-tuning on very limited data, the model is prone to overfitting to the segmentation task and forgetting the general semantic knowledge and lesion localization capabilities learned during pre-training. The above approach effectively preserves prior semantic knowledge by using self-supervised consistency loss as a regularization term. The class activation map itself contains coarse lesion localization information. Through consistency constraints, it is equivalent to providing the current model with a large number of weak localization supervision signals from unlabeled images, guiding the model to continuously focus on the correct semantic regions during fine-tuning. Especially when segmentation labels are scarce, this significantly improves segmentation accuracy compared to other fine-tuning methods that use mainstream regularization techniques (such as data augmentation and weight decay).
[0167] In some embodiments, in steps S2 and / or S3, the AdamW optimizer is used to optimize the model parameters, and a dynamic learning rate scheduling strategy is employed. Specifically, the dynamic learning rate scheduling strategy includes: during the initial warm-up phase... After rounds of iteration, the learning rate rises to the target learning rate, and then decreases from the target learning rate to the target minimum learning rate according to the cosine function law;
[0168] The learning rate during the initial warm-up phase After several iterations, the learning rate rises to the target learning rate, and its calculation formula is as follows:
[0169] , ];
[0170] in, Indicates the first The learning rate for each iteration round. This represents the minimum learning rate. Indicates the target learning rate. This represents the total number of iterations during the warm-up phase.
[0171] The calculation formula for decreasing the learning rate from the target learning rate to the target minimum learning rate according to the cosine function law is as follows:
[0172] ];
[0173] in, This indicates the number of iterations remaining after the warm-up phase. The learning rate for each iteration round. This indicates the total number of training iterations.
[0174] The AdamW optimizer combines the adaptive learning rate advantage of Adam with the regularization effect of weight decay: Adam dynamically adjusts the learning rate of each parameter by calculating the first moment estimate (momentum) and second moment estimate (adaptive learning rate) of the gradient, making it suitable for fast training of complex networks; weight decay, on the other hand, penalizes the model weight parameters to prevent overfitting caused by excessively large parameter values. Unlike the traditional Adam which integrates weight decay into gradient calculation, AdamW directly decays the weights in the parameter update step, correcting the unintended behavior when weight decay is combined with adaptive learning rate, making the regularization effect more accurate, especially suitable for the complex network structure of two-stage training in the method involved in this application. The optimal parameter settings for the AdamW optimizer are: decay rate 0.01, momentum 0.9, β2=0.999, β1=0.9. This parameter combination has been experimentally verified to achieve stable convergence in both the classification pre-training and segmentation fine-tuning stages.
[0175] During the preheating phase, after the initial... Round of iterative training (preferred) In the range of 5, the learning rate starts from the initial minimum (e.g., 10). -6 Linearly increase to the target learning rate (e.g., 10). -4Since the model parameters are initially random values, the gradient may be large and unstable. A low initial learning rate can avoid initial model oscillations and allow the model to gradually adapt to the training data. A linear boosting approach can smoothly transition to a normal training state, laying a stable foundation for subsequent optimization.
[0176] During the cosine annealing stage, after the warm-up, the learning rate decreases according to the law of the cosine function. The trend of the learning rate decrease in this stage is slow at first, then fast, and then slow again. The learning rate decreases slowly in the early stage, allowing the model to fully explore the parameter space. In the later stage, the learning rate decreases rapidly, helping the model to fine-tune the parameters. Then the learning rate decreases slowly again, converging to a flatter local optimum with stronger generalization ability.
[0177] In this embodiment, the total number of iterations for both phases of training is set to 100, and the learning rate eventually decreases to 10. -6 This ensures that the model can continue to update slowly in the later stages of training, avoiding premature stagnation.
[0178] In some embodiments, during the visualization stage of the decoder output, taking the segmentation of lesion regions as an example, a segmentation threshold can be set. Pixels with calculated pixel-level segmentation label probability values lower than the segmentation threshold are marked as non-lesion regions, and pixels with calculated pixel-level segmentation label probability values greater than or equal to the segmentation threshold are marked as lesion region pixels.
[0179] In other embodiments, such as Figure 4 As shown, in step S3, using the second training set, the encoder and decoder parameters of the U-shaped network model trained in step S2 are optimized by minimizing the image segmentation loss function to obtain the final image segmentation model, specifically including:
[0180] S31: Use the second training set as the current training set;
[0181] S32: Using the current training set, optimize the encoder and decoder of the U-shaped network model by minimizing the image segmentation loss function; after this round of optimization, use the U-shaped network model to perform forward inference on all second sample images in the current training set to obtain the pixel-level segmentation prediction probability map corresponding to each second sample image;
[0182] S33: For each second sample image in the current training set, perform the following operations: S331: Calculate the uncertainty score for each pixel position based on the pixel-level segmentation prediction probability map; the uncertainty score is calculated using the variance of the pixel-level segmentation prediction probability map under multiple random forward propagations, or using the information entropy of the pixel-level segmentation prediction probability map itself; S332: Correct the pixel-level segmentation label corresponding to the second sample image using the uncertainty score according to a preset cleanup strategy;
[0183] S34: Update all the second sample images obtained after purification in step S33 and their corrected pixel-level segmentation labels to the new current training set;
[0184] S35: Repeat steps S32 to S34 until the preset convergence condition is met.
[0185] A pixel-level segmentation prediction probability map is a probability distribution matrix output by the model after performing forward inference on the second sample image of the current training set. The matrix represents the probability distribution of each pixel belonging to each semantic category. The dimension is the same as the input image, and each element represents the probability of the corresponding pixel belonging to a certain category.
[0186] Uncertainty score is an indicator that measures the confidence of a model in its pixel-level segmentation prediction results. The higher the score, the more uncertain the model is about the class of that pixel. It can be calculated by the variance of multiple random forward propagations or the information entropy of the prediction probability map.
[0187] The purification strategy refers to the rules for correcting pixel-level segmentation labels based on uncertainty scores. The purpose is to remove label noise, supplement high-quality pseudo-labels, and improve the quality of training data.
[0188] Preset convergence conditions refer to the criteria for determining the termination of the iteration, including the number of iterations reaching a threshold or the proportion of samples to be reviewed falling below a threshold, to ensure a balance between training efficiency and effectiveness.
[0189] In step S32, the uncertainty score of each pixel position is calculated by using the variance of the pixel-level segmentation prediction probability map under multiple random forward propagations. Specifically, this includes performing multiple random forward propagations on the same image (e.g., repeating inference 10 times with the Dropout layer enabled) to obtain multiple segmentation prediction results, and calculating the variance of the prediction category for each pixel. The larger the variance, the higher the uncertainty.
[0190] The uncertainty score for each pixel position is calculated by using the information entropy of the pixel-level segmentation prediction probability map itself. Specifically, this includes: for the pixel's prediction probability distribution p = (p1, p2, ... p... m …, p n(where n is the number of categories, m < n, indicating the m-th pixel), the formula for calculating information entropy H is as follows: The larger the entropy value, the more uniform the probability distribution, and the less certain the model is in classifying the pixel.
[0191] In scenarios with few samples, the above approach generates high-quality pseudo-labels through high-confidence model prediction, effectively amplifying the supervisory signal and compensating for the lack of sufficient labeled data. Simultaneously, uncertainty score filtering eliminates noise from the original labels, improving the inherent consistency of the training data and providing a more reliable foundation for model optimization. The label cleansing strategy reduces the model's dependence on highly uncertain (potentially erroneous) labels, making the model learning process more robust and more tolerant of labeling errors, particularly suitable for scenarios where human error may exist in medical image annotation. Through iterative self-training, model performance gradually optimizes as the quality of training data improves, overcoming the performance limitations of initially limited labeled data and achieving segmentation accuracy superior to traditional fine-tuning. This is especially suitable for real-world applications where labels are scarce and of varying quality.
[0192] In some embodiments, the purification strategy includes:
[0193] For pixels with uncertainty scores below the first score threshold, their original pixel-level segmentation labels are retained.
[0194] For pixels with uncertainty scores higher than the first score threshold but lower than the second score threshold, the original pixel-level segmentation label is replaced by the prediction category at the corresponding position in the pixel-level segmentation prediction probability map.
[0195] For pixels with uncertainty scores higher than the second score threshold, their original pixel-level segmentation labels are retained and they are marked as samples to be reviewed.
[0196] The preset convergence conditions include the number of iterations reaching a preset iteration threshold, or the proportion of pixels in the current training set marked as samples to be reviewed being lower than a preset proportion.
[0197] The three-tiered processing mechanism described above differentiates labels based on model confidence. It leverages model capabilities to correct clearly identifiable label noise while employing a conservative strategy to prevent error propagation, thus balancing automation efficiency with risk control. Medium-confidence regions are typically the primary source of label inconsistencies or ambiguities. Replacing labels in these regions with model predictions significantly improves the intrinsic consistency of the training set, reduces noise interference during model training, and enhances segmentation accuracy. The preservation and labeling of high-uncertainty regions prevents the model from blindly learning in perplexing domains, avoids the risk of self-reinforcing errors, and ensures the robustness of the self-training process. Simultaneously, the labeling of samples awaiting review provides domain experts with clear review targets, improving the efficiency of the labeling workflow and enhancing the interpretability of the model training process.
[0198] In some embodiments, in step S32, the image segmentation loss function is a loss function based on uncertainty weighting; such as Figure 5 As shown, using the current training set, the encoder and decoder of the U-shaped network model are optimized in one round by minimizing the image segmentation loss function, specifically including:
[0199] S51: Input the second sample image in the current training set into the current model to obtain the pixel-level segmentation prediction probability map and the corresponding pixel-level uncertainty score;
[0200] S52: Assign a dynamic weight to each pixel based on the uncertainty score, wherein pixels with higher uncertainty are assigned lower weights;
[0201] S53: Based on the dynamic weights and the pixel-level segmentation prediction probability map, calculate the weighted image segmentation loss function value, and update the model parameters accordingly; wherein, the weighted image segmentation loss function... Calculated using the following formula:
[0202] ;
[0203] in, This represents the index of all pixels in all second sample images within the current batch. Let be the class probability vector predicted by the model for the i-th pixel. Let be the one-hot encoded vector of the pixel-level segmentation label corresponding to the i-th pixel in the current training set. To calculate the basic loss function for the difference between the prediction and the label, This represents the dynamic weight corresponding to the i-th pixel.
[0204] In step S52, after inputting the second sample image of the current training set into the model, not only is a pixel-level segmentation prediction probability map obtained, but also the uncertainty score of each pixel is calculated simultaneously (the calculation method is as described above). Then, dynamic weights w are assigned to each pixel based on the uncertainty score. i The weights are negatively correlated with the uncertainty score: the higher the uncertainty, the higher the weights. i The smaller, the better. i The larger the weight, the better. The core function of dynamic weights is to allow the model to focus more on pixels with low uncertainty and high reliability (high weight, large contribution) when updating parameters, while reducing interference from pixels with high uncertainty and potential errors (low weight, small contribution). The mapping between weights and uncertainty scores can be either exponential decay or linear decay.
[0205] In step S53, for each pixel, the basic loss L is calculated. base (p i ,y i (e.g., Dice loss or cross-entropy loss) measures the difference between the pixel prediction and the label. The base loss for each pixel is then multiplied by the corresponding dynamic weight w. i Multiplying these results in a weighted loss, highlighting the loss contribution of reliable pixels and downplaying the loss contribution of unreliable pixels. The weighted losses of all pixels are summed, then divided by the sum of the weights of all pixels (normalization process) to obtain the weighted loss value for the current batch, avoiding fluctuations in loss values caused by differences in weight distribution. Then, L... weight To optimize the target, gradients are calculated through backpropagation and the encoder and decoder parameters of the model are updated, so that the model automatically focuses on high-quality labeled regions during the learning process, thereby gradually improving the segmentation accuracy.
[0206] Through the above approach, the model will prioritize learning regions with low uncertainty and simple, clear characteristics (such as the core region of the lesion) in the early stages of training. As training progresses and the model's capabilities improve, it will gradually expand its learning focus to complex regions with moderate uncertainty (such as the boundary of the lesion). This aligns with the human learning pattern of "from easy to difficult," resulting in smoother convergence and better final performance.
[0207] In the second aspect, such as Figure 7 As shown, this application provides a few-shot image segmentation training system based on a classification-assisted U-shaped network. The system is used to perform the method described in the first aspect of this application. The few-shot image segmentation training system 10 includes:
[0208] U-shaped network model construction module 101 is used to construct a U-shaped network model, which includes an encoder and a decoder;
[0209] The encoder is used to extract multi-scale features from the input image;
[0210] The decoder is connected to the encoder via a skip connection and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales.
[0211] The decoder's outputs at at least two different scales are also connected to a multi-scale feature fusion classification module;
[0212] The first training module 102 is used to acquire a first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes a first sample image and the category label corresponding to the first sample image.
[0213] During the training of the U-shaped network model, the multi-scale feature fusion classification module performs the following operations:
[0214] The decoded feature map at the k-th scale output by the decoder is subjected to convolutional dimensionality reduction, global average pooling, and flattening operations to obtain the first feature vector, where k is an integer greater than 1;
[0215] The (k-1)th scale decoded feature map output by the decoder is subjected to the convolutional dimensionality reduction, global average pooling, and flattening operations to obtain an intermediate feature vector, and the intermediate feature vector is subjected to a linear transformation to obtain a second feature vector.
[0216] The first feature vector is concatenated with the second feature vector, and a linear transformation is performed on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector.
[0217] The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector;
[0218] The fused feature vector is input into a fully connected classifier, mapped to the number of categories through a fully connected layer, and then the classification prediction result of the first sample image is output through the SoftMax activation function.
[0219] The second training module 103 is used to obtain a second training set, which includes a second sample image and a label for each pixel in the second sample image to indicate its semantic category, generating pixel-level segmentation labels. The second sample image contains lesion regions, and the number of first sample images in the second training set is less than that in the first training set. Using the second training set, the encoder and decoder parameters of the U-shaped network model trained in step S2 are optimized by minimizing the image segmentation loss function to obtain the final image segmentation model.
[0220] like Figure 6 As shown, the encoder configuration consists of four stages, each containing two 3×3 convolutional layers, a ReLU activation layer, and one 2×2 max pooling layer (the last stage has no pooling). The number of feature channels are 64, 128, 256, and 512, respectively.
[0221] The decoder configuration consists of four stages, each containing one 2×2 transposed convolutional layer (upsampling), two 3×3 convolutional layers, and a ReLU activation layer, with the number of feature channels being 512, 256, 128, and 64 respectively.
[0222] The multi-scale feature fusion classification module includes a fully connected classifier, which consists of two fully connected layers (hidden layer dimension 2048, output layer dimension N) and a SoftMax activation layer. The feature vector length after 1×1 convolution dimensionality reduction is K=1024.
[0223] This application discloses a method and system for training a U-shaped network for few-shot image segmentation based on classification assistance. The method first trains a U-shaped network model using a first training set with image-level category labels, optimizing parameters by minimizing the classification loss. During this process, a multi-scale feature fusion classification module fuses and classifies the multi-scale features output by the decoder, enabling the network to learn strong semantic features related to segmentation. Then, a second training set containing a small number of sample images and pixel-level segmentation labels is obtained. By fixing the classification module parameters, the encoder and decoder parameters are optimized by minimizing the segmentation loss using only this second training set with fewer samples, resulting in the final image segmentation model. This invention, through a two-stage training strategy, utilizes a large number of readily available classification labels for pre-training, effectively improving the model's segmentation performance and generalization ability when pixel-level labeled data is extremely limited.
[0224] Finally, it should be noted that although the above embodiments have been described in the text and drawings of this application, this should not limit the scope of patent protection of this application. Any technical solutions that are based on the essential concept of this application and utilize the content described in the text and drawings of this application, resulting in equivalent structural or procedural substitutions or modifications, as well as the direct or indirect application of the technical solutions of the above embodiments to other related technical fields, are all included within the scope of patent protection of this application.
Claims
1. A few-shot image segmentation training method based on a classification-assisted U-shaped network, characterized in that, The method includes: S1: Construct a U-shaped network model, which includes an encoder and a decoder; The encoder is used to extract multi-scale features from the input image; The decoder is connected to the encoder via a skip connection and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales. The decoder's outputs at at least two different scales are also connected to a multi-scale feature fusion classification module; S2: Obtain the first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes the first sample image and the category label corresponding to the first sample image. During the training of the U-shaped network model, the multi-scale feature fusion classification module performs the following operations: S21: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the k-th scale output by the decoder to obtain the first feature vector, where k is an integer greater than 1; S22: Perform convolutional dimensionality reduction, global average pooling, and flattening operations on the decoded feature map of the (k-1)th scale output by the decoder to obtain an intermediate feature vector, and perform a linear transformation on the intermediate feature vector to obtain a second feature vector; S23: Concatenate the first feature vector with the second feature vector, and perform a linear transformation on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector; S24: The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector; S25: Input the fused feature vector into a fully connected classifier, map it to the category number dimension through a fully connected layer, and then output the classification prediction result of the first sample image through the SoftMax activation function; S3: Obtain a second training set, which includes a second sample image and a label for each pixel in the second sample image that indicates its semantic category, and generate pixel-level segmentation labels. The second sample image contains lesion regions, and the number of first sample images in the second training set is less than that in the first training set. Using the second training set, optimize the encoder and decoder parameters of the U-shaped network model trained in step S2 by minimizing the image segmentation loss function to obtain the final image segmentation model.
2. The method for training few-shot image segmentation using a classification-assisted U-shaped network as described in claim 1, characterized in that, In step S2, after obtaining the first training set and before training the U-shaped network model using the first training set, the method further includes: S201: Perform black bar insertion and scaling processing on the first sample image, wherein the black bar insertion and scaling processing specifically includes: Obtain the height H and width W of the first sample image; If H>W, then insert black bars with a width of (HW) / 2 on the left and right sides of the first sample image respectively, so that the first sample image becomes a square; If H < W, black bars with a height of (W - H) / 2 are inserted on the upper and lower sides of the first sample image respectively to make the first sample image into a square; The obtained square first sample image is scaled to a preset size P×P by bicubic interpolation to obtain a preprocessed image; S202: Perform a random rotation on the preprocessed image, and the rotation angle is randomly selected within a preset angle range.
3. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 1, characterized in that, The image classification loss function includes the cross-entropy loss function L. cls Cross-entropy loss function L cls The calculation formula is: ; Where G is the true class label of the first sample image, and M cls The classification prediction result of the first sample image is output after mapping to the number of categories through a fully connected layer and then passing through the SoftMax activation function.
4. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 1, characterized in that, The image segmentation loss function includes the Dice loss function L. dice Dice loss function L dice The calculation formula is as follows: ; Among them, M seg G represents the lesion segmentation result output by the decoder. seg Let |·| represent the number of elements in the set, and ∩ represent taking the intersection.
5. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 1, characterized in that, The method further includes: After the training in step S2 is completed, save the U-shaped network model that is trained at this time as a prior model; During the optimization process of step S3, initialize the U-shaped network model to be optimized with the parameters of the prior model as the current model. For any second sample image in the second training set, perform the following operations: Input the second sample image into the current model and the prior model simultaneously; Using the decoders of the current model and the prior model, class activation maps corresponding to the second sample image are generated respectively, and the self-supervised consistency loss L between the two class activation maps is calculated. ss ; The image segmentation loss function is supervised by segmentation loss L. sup And self-supervised consistency loss L ss The weighted summation yields the image segmentation loss function L. seg The calculation formula is as follows: ; in, A balance coefficient greater than 0 is used to adjust the strength of the self-supervised signal; the supervised segmentation loss L... sup This is the supervised segmentation loss calculated based on the pixel-level segmentation labels in the second training set.
6. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 1, characterized in that, In steps S2 and / or S3, the AdamW optimizer is used to optimize the model parameters, and a dynamic learning rate scheduling strategy is employed. Specifically, the dynamic learning rate scheduling strategy includes: during the initial warm-up phase... After rounds of iteration, the learning rate rises to the target learning rate, and then decreases from the target learning rate to the target minimum learning rate according to the cosine function law; The learning rate during the initial warm-up phase After several iterations, the learning rate rises to the target learning rate, and its calculation formula is as follows: , ]; in, Indicates the first The learning rate for each iteration round. This represents the minimum learning rate. Indicates the target learning rate. This represents the total number of iterations during the warm-up phase. The reduction from the target learning rate to the target minimum learning rate follows the cosine function rule, and its calculation formula is as follows: ]; in, This indicates the number of iterations remaining after the warm-up phase. The learning rate for each iteration round. This indicates the total number of training iterations.
7. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 1, characterized in that, In step S3, using the second training set, by minimizing the image segmentation loss function, optimize the encoder and decoder parameters of the U-shaped network model trained in step S2 to obtain a final image segmentation model, specifically including: S31: Use the second training set as the current training set; S32: Use the current training set to perform one round of optimization on the encoder and decoder of the U-shaped network model by minimizing the image segmentation loss function; after this round of optimization is completed, use the U-shaped network model to perform forward inference on all second sample images in the current training set to obtain a pixel-level segmentation prediction probability map corresponding to each second sample image; S33: For each second sample image in the current training set, perform the following operations: S331: Based on the pixel-level segmentation prediction probability map, calculate the uncertainty score for each pixel position; the uncertainty score is calculated by the variance of the pixel-level segmentation prediction probability map under multiple random forward propagations, or by the information entropy of the pixel-level segmentation prediction probability map itself; S332: According to a preset purification strategy, use the uncertainty score to correct the pixel-level segmentation label corresponding to the second sample image; S34: Update all the second sample images and their corrected pixel-level segmentation labels obtained after being purified in step S33 as the new current training set; 8. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 7, characterized in that, S35: Repeat steps S32 to S34 until a preset convergence condition is reached. The purification strategy includes: For pixels with an uncertainty score lower than the first score threshold, retain their original pixel-level segmentation labels; For pixels with an uncertainty score higher than the first score threshold but lower than the second score threshold, replace their original pixel-level segmentation labels with the predicted categories at the corresponding positions in the pixel-level segmentation prediction probability map; For pixels with an uncertainty score higher than the second score threshold, retain their original pixel-level segmentation labels and mark them as samples to be reviewed; The preset convergence condition includes that the number of iterations reaches a preset iteration number threshold, or the proportion of pixels marked as samples to be reviewed in the current training set is lower than a preset proportion.
9. The few-shot image segmentation training method based on classification-assisted U-shaped network as described in claim 7, characterized in that, In step S32, the image segmentation loss function is an uncertainty-weighted loss function; using the current training set, the encoder and decoder of the U-shaped network model are optimized by minimizing the image segmentation loss function, specifically including: Input the second sample image in the current training set into the current model to obtain the pixel-level segmentation prediction probability map and the corresponding pixel-level uncertainty score; Based on the uncertainty score, a dynamic weight is assigned to each pixel, wherein pixels with higher uncertainty are given lower weights. Based on the dynamic weights and the pixel-level segmentation prediction probability map, a weighted image segmentation loss function value is calculated, and the model parameters are updated accordingly; wherein, the weighted image segmentation loss function... Calculated using the following formula: ; in, This represents the index of all pixels in all second sample images within the current batch. Let be the class probability vector predicted by the model for the i-th pixel. Let be the one-hot encoded vector of the pixel-level segmentation label corresponding to the i-th pixel in the current training set. To calculate the basic loss function for the difference between the prediction and the label, This represents the dynamic weight corresponding to the i-th pixel.
10. A few-shot image segmentation training system based on a classification-assisted U-shaped network, characterized in that, The system is configured to perform the method as described in any one of claims 1 to 9, the system comprising: A U-shaped network model construction module is used to construct a U-shaped network model, which includes an encoder and a decoder; The encoder is used to extract multi-scale features from the input image; The decoder is connected to the encoder via a skip connection and is used to upsample and fuse the features extracted by the encoder, and output at least two decoded feature maps at different scales. The decoder's outputs at at least two different scales are also connected to a multi-scale feature fusion classification module; The first training module is used to acquire a first training set, train the U-shaped network model using the first training set, and optimize the model parameters by minimizing the image classification loss function. The first training set includes a first sample image and the category label corresponding to the first sample image. During the training of the U-shaped network model, the multi-scale feature fusion classification module performs the following operations: The decoded feature map at the k-th scale output by the decoder is subjected to convolutional dimensionality reduction, global average pooling, and flattening operations to obtain the first feature vector, where k is an integer greater than 1; The (k-1)th scale decoded feature map output by the decoder is subjected to the convolutional dimensionality reduction, global average pooling, and flattening operations to obtain an intermediate feature vector, and the intermediate feature vector is subjected to a linear transformation to obtain a second feature vector. The first feature vector is concatenated with the second feature vector, and a linear transformation is performed on the concatenated feature vector to obtain a third feature vector with the same dimension as the first feature vector. The third feature vector from the highest-scale decoded feature map is fused with the encoded feature map output by the encoder after the last downsampling to obtain a fused feature vector; The fused feature vector is input into a fully connected classifier, mapped to the number of categories through a fully connected layer, and then the classification prediction result of the first sample image is output through the SoftMax activation function. The second training module is used to obtain a second training set, which includes a second sample image and a label for each pixel in the second sample image to indicate its semantic category, generating pixel-level segmentation labels. The second sample image contains lesion regions, and the number of first sample images in the second training set is less than that in the first training set. Using the second training set, the encoder and decoder parameters of the U-shaped network model trained in step S2 are optimized by minimizing the image segmentation loss function to obtain the final image segmentation model.