Multi-task medical image segmentation method and system

Abstract

The invention discloses a multi-task medical image segmentation method and system, and relates to the field of biomedical image recognition and segmentation. Aiming at the problems of high development cost, large computing resource consumption and the like of multi-task medical image segmentation, the method comprises the following steps: firstly, acquiring a medical image segmentation data set, and then constructing and training an improved SAM visual large model; on this basis, an affinity graph prediction network and a category prediction network are introduced to obtain a multi-task medical image instance segmentation model, and joint training is carried out; and finally, inputting a medical image to be separated into the trained multi-task medical image instance segmentation model, and outputting a final segmentation result. According to the method, the SAM model with strong feature extraction capability and generalization capability is improved, so that the SAM model has affinity graph prediction and category prediction functions, and an efficient and accurate multi-task medical image instance segmentation model is realized.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image processing, and in particular to a multi-task medical image segmentation method and system, and in particular to a multi-task medical image segmentation method and system based on a large visual model. Background Art

[0002] With the development of deep learning technology, convolutional neural network (CNN)-based methods have achieved remarkable results in medical image segmentation. For example, the U-Net model effectively captures contextual information through its unique encoder-decoder structure, achieving high-precision segmentation. However, most of these models are designed and trained for a single type of medical image segmentation task.

[0003] When faced with multiple medical image segmentation tasks, separate models are typically trained for each task. This presents numerous challenges: First, development costs are high, requiring significant manpower and time to collect data, design model structures, and train the models. Second, computational resources are intensive, as the resources required to run multiple models simultaneously increase exponentially, making them unsustainable for medical institutions or mobile devices with limited computing resources. Furthermore, the maintenance and management of multiple models is complex, making practical applications more challenging.

[0004] In recent years, some research has attempted to transfer the SAM (Segment Anything Model), a model with outstanding performance in natural image segmentation, to the medical imaging domain. For example, MedSAM, trained on data from 10 modalities and over 1.5 million masks, has enabled SAM to achieve general medical image segmentation capabilities. However, these direct transfer approaches inherit the shortcomings of the SAM model, namely, the requirement for a large number of prompts to complete the segmentation of all instances. In medical imaging scenarios, doctors or professional operators must spend a considerable amount of time and effort manually annotating prompts, significantly reducing segmentation efficiency and severely impacting the smoothness of diagnostic processes in time-sensitive scenarios such as clinical diagnosis. Even with automatically generated prompts, the model still takes a long time to complete segmentation based on these prompts, failing to meet the demand for fast and accurate segmentation results in practical applications. The recent UN-SAM model, which extends SAM to multi-domain cell nucleus segmentation, essentially performs only semantic segmentation and cannot distinguish adjacent objects, making it unsuitable for direct application to instance segmentation tasks.

[0005] In summary, the existing technology lacks efficient and universal methods and models for handling various medical image segmentation tasks, and innovative solutions are urgently needed. Summary of the Invention

[0006] In view of the shortcomings of the existing technology that multi-task medical image segmentation requires separate model training, resulting in high development costs, large consumption of computing resources, and low segmentation efficiency using the SAM model, the present invention proposes a multi-task medical image segmentation method and system. The present invention makes innovative improvements to the SAM visual large model and is committed to realizing the efficient and accurate processing of multiple medical image segmentation tasks using a single model.

[0007] In a first aspect, the present invention provides a multi-task medical image segmentation method, characterized in that it includes the following steps:

[0008] 1) obtaining a medical image dataset in the same field as the medical image to be segmented, wherein the dataset includes image data and annotated segmentation masks;

[0009] 2) preprocessing the data set;

[0010] 3) Build an improved SAM visual model that introduces task tokens. Initialize the parameters of the original SAM model, use a combination of Dice loss and Focal loss as the loss function, and perform the first phase of training using only randomly selected pixels in the region indicated by the ground truth mask as cues.

[0011] 4) Based on the improved SAM visual model, an affinity map prediction network and a category prediction network are introduced. The affinity map prediction network uses the UNETR model structure and takes the 32-channel features used to predict the segmentation mask in the improved SAM visual model as input. Its output is used to calculate the affinity map of the image in the horizontal and vertical directions. The category prediction network includes a task attention module and a classification head for predicting categories. The second stage of training is performed without using any prompts to obtain an optimized multi-task medical segmentation model.

[0012] 5) Input the medical image to be segmented into the optimized multi-task medical segmentation model and output the final segmentation result.

[0013] In a second aspect, the present invention provides a multi-task medical image instance segmentation system for implementing the method of claim 1, comprising:

[0014] A data processing module, used for preprocessing the medical image dataset and the medical image dataset to be segmented;

[0015] Improved SAM visual large model construction and training module, used to build an improved SAM visual large model with task labels and perform the first phase of training;

[0016] The multi-task medical segmentation model construction and training module introduces the affinity graph prediction network and category prediction network based on the improved SAM visual model to obtain a multi-task medical segmentation model. The multi-task medical segmentation model is trained as a whole to obtain an optimized multi-task medical segmentation model.

[0017] The segmentation module is used to input the medical image to be segmented into the optimized multi-task medical segmentation model and output the segmentation result.

[0018] The present invention, through the above technical solution, produces the following beneficial effects:

[0019] (1) Powerful multi-task medical image segmentation capability: Based on the original SAM visual model, this paper leverages its powerful feature extraction and generalization capabilities to introduce task labels to restructure the image encoder and decoder, making them adaptable to a variety of medical image data. Furthermore, it introduces an affinity map prediction module and a category prediction module to achieve high-precision instance segmentation.

[0020] (2) Improve the defects of the SAM model: As a large model for visual segmentation, the SAM model has two obvious defects. First, a large number of prompts are required to segment all instances, which will seriously affect the efficiency of segmentation. By introducing the affinity graph prediction module and the graph partitioning operation, all instances in the image can be segmented at once. Second, SAM only segments the target but cannot recognize the target. By introducing the category prediction module, the segmented instances can be classified.

[0021] (3) Reduce development costs: Change the traditional model of training a separate model for each task, and only need to improve and optimize a single SAM model, which greatly reduces the workload of data collection, labeling and model design, reduces manpower and time costs, and makes the development process more efficient and economical;

[0022] (4) Simplified maintenance and management: Only one multi-task model needs to be maintained. Compared with multiple independent models, management tasks such as version updates and troubleshooting are simpler, and the system stability and reliability are higher, which facilitates management optimization in actual applications. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] Figure 1 A schematic diagram of a flow chart of a multi-task medical image instance segmentation method provided according to an embodiment of the present invention;

[0024] Figure 2 A schematic diagram of the model structure of an improved SAM model provided according to an embodiment of the present invention;

[0025] Figure 3 This is a schematic diagram of the DownConv module structure;

[0026] Figure 4 Schematic diagram of the TAM module structure;

[0027] Figure 5 Schematic diagram of the calculation process of the task labeling fusion module;

[0028] Figure 6 Schematic diagram of the overall structure of the multi-task medical image segmentation model;

[0029] Figure 7 This is the electron microscope neuron image segmentation result. DETAILED DESCRIPTION

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

[0031] In one embodiment of the present invention, a medical image segmentation method capable of processing multiple tasks simultaneously is described as follows: Figure 1 Shown, including:

[0032] 1) obtaining a medical image dataset in the same field as the medical image to be segmented, wherein the dataset includes image data and annotated segmentation masks;

[0033] 2) Preprocessing the dataset, including scaling, Gaussian noise, small-angle rotation, cropping, and normalization;

[0034] 3) Build an improved SAM visual large model that introduces task tokens (hereinafter referred to as the improved SAM model for ease of description); set the input image resolution to 512×512, replace the patch embedding module in the original SAM image encoder with a DownConv module, replace the mask tokens in the original SAM decoder with task tokens, and replace the self-attention module in the original SAM image encoder with a task attention module. Use the parameters of the original SAM as the initial parameters, use a combination of Dice loss and Focal loss as the loss function, and only use randomly selected pixels in the area indicated by the ground truth mask as cues for the first stage of training;

[0035] 4) Based on the improved SAM visual model, an affinity map prediction network and a category prediction network are introduced. The affinity map prediction network uses the UNETR model structure, takes the 32-channel features used to predict the segmentation mask in the improved SAM visual model as input, and its output is used to calculate the affinity map of the image in the horizontal and vertical directions; the category prediction network includes a task attention module and a classification head for predicting the category of the entity; each classification head includes a linear mapping layer and a softmax function; the second stage of training is performed without using any prompts to obtain an optimized multi-task medical segmentation model;

[0036] 5) Input the medical image to be segmented into the optimized multi-task medical segmentation model and output the final segmentation result.

[0037] Step 1) describes the data acquisition requirements of the present invention. Researchers need to obtain relevant medical image datasets from public medical image databases or in cooperation with medical institutions based on the target medical image segmentation field of interest, such as tumor detection and segmentation, image analysis of neurological diseases, etc. These datasets can not only contain data of various types (such as tumor images, cell images, etc.), but also image data of various modalities (such as CT, MRI, ultrasound, etc.). These data must also have corresponding accurately manually annotated segmentation masks to provide accurate reference standards for subsequent model training and evaluation.

[0038] It should be noted that in the present invention, "multi-task" refers to the medical segmentation model learning and performing multiple related or unrelated segmentation tasks simultaneously, rather than focusing on a single task. For example, the model processes multiple segmentation targets simultaneously during training (for example, simultaneously segmenting different structures such as tumors, organs, and blood vessels). There are some typical multi-task scenarios in medical images, such as the simultaneous segmentation of multiple organs such as the liver, kidneys, and spleen in abdominal CT; for example, segmentation tasks combining different imaging modalities (such as T1, T2, and DWI sequences of MRI).

[0039] Step 2) describes the data preprocessing and enhancement methods, which specifically include the following:

[0040] Data augmentation and preprocessing techniques were applied to the dataset, including scaling, adding Gaussian noise, small-angle rotation, random cropping, and normalization. Gaussian noise was added because noise interference is unavoidable in real-world medical imaging. To enhance the robustness of the model and enable accurate image segmentation in real-world settings, Gaussian noise was added to the data. Scaling, adding Gaussian noise, small-angle rotation, and random cropping are all data augmentation methods. Because obtaining a large number of medical images is difficult, data augmentation is used to increase the amount of data. Normalization, which involves uniformly mapping image pixel values ​​to the range [-1, 1], is crucial because it accelerates model convergence and improves training efficiency. Specifically, the mean and standard deviation of the pixel values ​​in each channel of all images are calculated. The mean is then subtracted from each channel and divided by the standard deviation. Furthermore, because images from different modalities have different properties, some are RGB images while others are single-channel. To ensure that all data can be processed by the same model, the single-channel images are replicated twice and merged into three channels, treating them as RGB images.

[0041] In order to verify the method of the present invention, the present invention divides the data set in step 1) into a training set and a test set. Among them, the training set data is used for training and optimization of the multi-task medical segmentation model, and the test set is used to test the performance of the optimized multi-task medical segmentation model, and to evaluate the performance of the method of the present invention and the corresponding model. Step 3) describes the first training phase after completing data collection and preprocessing. Its core is to train the improved SAM visual large model that introduces task tokens (task tags), with the aim of adapting the basic SAM model to different types of new data in advance. The structure of the improved visual large SAM model is as follows Figure 2 As shown in Figure 1, it is based on the original SAM model with the following modifications:

[0042] 3-1) The input image resolution is set to 512×512. The image is processed into an input image with a resolution of 512×512 in the data preprocessing stage;

[0043] 3-2) The input image resolution of the original SAM model is 1024×1024. To adapt to the new image resolution, the patch embedding module in the image encoder of the original SAM model is replaced with a DownConv module. To ensure that the feature map resolution entering the attention module in the image encoder is 64×64, consistent with the original SAM model;

[0044] 3-3) Replace the mask tokens in the decoder of the original SAM model with task tokens. Each task token corresponds to a segmentation task, and the task token is a learnable vector.

[0045] 3-4) The self-attention module (Self-Attention) in the image encoder of the original SAM model is replaced with a task-attention module (TAM), so that the image encoder can further adapt to multi-modal and multi-type medical images.

[0046] Improve the structure of the DownConv module of the SAM visual large model as follows Figure 3 As shown, the network consists of the first convolutional layer, the first BatchNorm normalization layer, the first GELU activation layer, the second convolutional layer, the second BatchNorm normalization layer, the second GELU activation layer, the third convolutional layer, the third BatchNorm normalization layer, the third GELU activation layer, and the fourth convolutional layer. The first convolutional layer has a kernel size of 3 and a stride of 1, which is used for image feature extraction. The kernel size of all other convolutional layers is set to 5 and the stride is set to 2, which means that the second, third, and fourth convolutional layers all have a 2x downsampling effect. These three convolutional layers take a 512×512 image as input and output a 64×64 image feature map. This feature map is input to the Task Attention Module (TAM) following the DownConv module. The BatchNorm normalization layer adjusts the distribution of the input data, making network training more efficient and stable. The GELU activation layer enhances the model's nonlinear modeling capabilities. Therefore, the DownConv module ensures that the feature map entering the attention module has a resolution of 64×64, consistent with the original SAM model. The feature map output by the DownConv module has a resolution of 64×64 and a channel count of d, which is consistent with the original SAM model. The image encoder of the original SAM model uses three Vision Transformer architectures as the backbone network: SAM-B, SAM-L, and SAM-H, with d of 768, 1024, and 1280, respectively.

[0047] The structure of the Task Attention Module (TAM) is as follows Figure 4 As shown in the figure, a task tag fusion module is added to the original self-attention module. The original self-attention operation is to perform self-attention on image features, while the task tag fusion module uses image features and task tags for cross attention. The output result is added to the output result of the original self-attention operation. This module can integrate the task tag information into the calculation of the image feature map. The calculation process of the task tag fusion module is as follows: Figure 5 As shown, first mark the task t s Input to the first linear mapping layer and the second linear mapping layer, output V and K respectively, and expand the 64×64×d feature map into a dimension of 642 The feature map f of ×d is input to the third linear mapping layer, and the output is Q, and then the calculation Then the attention matrix is ​​obtained through the softmax function Finally, perform matrix multiplication on the attention matrix and V to obtain the output of the task labeling fusion module.

[0048] The improved SAM visual model uses the pre-trained parameters of the original SAM model as initial parameters and the collected medical images as training data. The training process is as follows:

[0049] 3-5) In each training step, randomly select a task s from the S tasks and randomly select preprocessed image data X from the corresponding dataset. s , and in X s Randomly select an instance mask m from the annotation gt , and finally in the mask m gt A pixel in the indicated area is randomly selected as the cue point p0;

[0050] 3-6) Image X s The task label t corresponding to the selected segmentation task s s The two are input to the image encoder together, and the output is the image code F, which has a size of 64×64×d;

[0051] 3-7) Input the prompt point p0 into the prompt encoder and output the prompt point code pnt, where the image encoder is consistent with the original SAM model, and pnt is a 256-dimensional vector, which is consistent with the task label;

[0052] 3-8) Intersection-over-Union mark iou, prompt point code pnt and task mark t s They are all 256-dimensional vectors, and they are combined into a whole, called a label group;

[0053] 3-8) Input the image code F and the tag group together into the SAM decoder, and output the selected true mask m gt Prediction of m p , and output the true mask m gt and the predicted mask m p The intersection and union of iou gt Prediction of iou p ;

[0054] 3-9) The loss function is divided into two parts. One part is the loss function L related to the intersection-over-union score prediction. iou , the other part is the loss function L related to mask prediction m Among them, L iou =(iougt -iou p ) 2 Describes the predicted intersection-over-union score iou p with iou gt Squared error; mask loss L m It is the focal lossL of the two categories focal and dice lossL dice The combination of two categories of focal loss is:

[0055]

[0056] Among them, α is the balance weight, which is used to adjust the contribution of positive and negative samples to the loss. It usually takes a value between 0 and 1, and α is generally 0.25 or 0.75; γ is the focus parameter, which is used to control the degree of attention paid to difficult and easy samples. γ is usually 2. dice The formula is:

[0057] L dice =1-2m p ·m gt / (m p +m gt )

[0058] The mask loss function is:

[0059] L m =μL focal +L dice

[0060] Among them, μ is a parameter that balances the contribution of the two losses. The complete loss function is the sum of the two parts:

[0061] L total =L iou +L m

[0062] 3-10) In this embodiment, the number of iterative training is set to 100,000 times as the training termination condition. Therefore, training steps 3-5) to 3-9) are repeated 100,000 times to complete the first stage of training.

[0063] Figure 6The figure shows the model architecture of the second stage training in step 4. Based on the model trained in the first stage (the improved SAM model), the affinity map prediction network module and the category prediction network module are introduced and the modules in the decoder that do not participate in the calculation of image features are discarded to obtain a multi-task medical segmentation model. Among them, the affinity map prediction network module includes a 2x upsampling operation and a 2D UNETR model with only 4 transformer layers. This module takes a 32-channel feature map as input and a 16-channel feature map as output. The output feature map can be used to calculate the affinity map of the image in the horizontal and vertical directions; the category prediction network module consists of two parts. The first part is the concatenation of the TAM modules in the image encoder of the three improved SAM models, and the second part is S classification heads. Each classification head contains a Linear mapping layer and a softmax function, where S represents the number of segmentation tasks. The second stage training is to further train the entire multi-task medical segmentation model based on the results of the first stage training. The specific process includes:

[0064] 4-1) The pre-processed image to be segmented X s Input to the image encoder of the improved version of SAM, and output the image code F, which has a size of 64×64×256; among them, the image to be segmented X s From the sth segmentation task;

[0065] 4-2) The iou token (intersection-over-union token) in the improved SAM model and the task token t corresponding to the segmentation task s are s Combined into a whole, recorded as a tag group, the image code F and the tag group are input into the decoder of the improved version of SAM. The decoder here is as follows Figure 6 As shown, this is a simplified version of the original decoder, and no prompt information is needed for training SAM at this time. The output is a 32-channel image feature G for calculating the mask, with a size of 256×256×32;

[0066] 4-3) Input the 32-channel image feature G into the affinity map prediction module. First, it is upsampled by 2 times to obtain a new feature G′ of size 512×512×32. Then G′ is input into the UNETR model, and the output is a 512×512×16 feature map H.

[0067] 4-4) After obtaining the image feature map H, the cosine similarity of the features at two positions in H is used as the affinity value of the corresponding pixel. Specifically, the cosine similarity of the features of each pixel with the previous pixel is calculated in both the horizontal and vertical directions to obtain the affinity map in both directions, which is collectively referred to as affinity map A. pThe affinity map represents the relationship between image pixels. If two pixels belong to the same target, the affinity value between them is large, otherwise it is small. It is reasonable to use the feature map G as input to predict the affinity map, because Figure 7 (a) and Figure 7 As shown in (b), Figure 7 (a) is an electron microscope neuron image. After the improved SAM model after the first stage of training is processed to obtain the corresponding feature H, PCA dimensionality reduction and visualization are performed. The visualization results are as follows Figure 7 As shown in (b), we can see that the features within the same instance are similar, and there are clear boundaries between different instances. This feature is very suitable for calculating affinity graphs;

[0068] 4-5) Using the graph partitioning algorithm, the image can be divided into different regions according to the affinity graph, which is equivalent to completing the image segmentation and obtaining a series of segmentation masks Where N is the number of masks; the affinity graph A represents the relationship between pixels. The elements in A are considered as edges, and the pixels are considered as vertices. The edges and vertices can form a graph structure. The affinity value is considered as the weight of the edge. Then, the graph partitioning algorithm can be used to partition the graph into multiple subgraphs, each of which represents a segmentation instance.

[0069] 4-6) Encode the image F and the task tag t s The two are input together into the first part of the category prediction module, and the output is the feature F′;

[0070] 4-7) Calculate each predicted segmentation mask m i With the true mask m gt The intersection-over-union ratio of each instance in m i The instance with the largest intersection-and-union ratio. If this maximum intersection-and-union ratio is greater than the set threshold γ, then m i Matching is achieved with this instance; according to the predicted segmentation mask m i , extract the corresponding area in the feature F′, that is, extract the features of the corresponding area. These features are 256-dimensional vectors, average these features, and then input them into the corresponding classification head C s , get the prediction y for the matching instance category p ; Among them, the threshold γ is generally taken as 0.5;

[0071] 4-8) The complete multi-task medical image instance segmentation model can be seen as having two branches, one is the affinity map prediction branch, and the output is the affinity map A p , which is consistent with the real affinity graph A gt Calculate the mean square loss L2; ​​the second is the category prediction branch, the output is the category prediction y p, and the true category y gt Calculate the multi-classification cross entropy loss L CE ; The calculation formulas of the two loss functions are as follows:

[0072]

[0073] Among them, K s is the number of categories of the s-th segmentation task, n represents A p The total loss function is the sum of the two:

[0074] L=L2+L CE

[0075] After calculating the loss, the model can be trained using the Adam optimization algorithm.

[0076] 4-9) In this embodiment, training steps 4-1) to 4-8) are repeated 100,000 times to complete the second stage of training.

[0077] Figure 7 The comparison of the segmentation results of the method of the present invention and the existing methods on electron microscope neuron images is demonstrated; Figure 7 (c) in FIG. 1 is the segmentation result of the method of the present invention. Figure 7 (d) is the segmentation result of the PEA method. It can be seen that the results of the PEA method in the red rectangle have obvious errors. The method of the present invention has more stable segmentation performance. The PEA method is cited from the following paper: Wei Huang, Shiyu Deng, Chang Chen, Xueyang Fu, Zhiwei Xiong. Learning to Model Pixel-Embedded Affinity for Homogeneous Instance Segmentation. The Thirty-Sixth AAAI Conference on Artificial Intelligence (AAAI-22)

[0078] In order to evaluate the performance of the method described in this paper in the field of medical image instance segmentation, the aggregated Jaccard Index (AJI) is used as the evaluation index of the segmentation results. In order to cope with the situation of multi-class instance segmentation, the AJI is calculated for each class separately and the average is finally calculated. The AJI calculation process for a single class is as follows:

[0079] 5-1) Let the set of true masks be {G i}, the set of prediction masks is recorded as {S k};

[0080] 5-2) Initialize the number of intersection pixels C = 0 and the number of union pixels U = 0;

[0081] 5-3) For the real mask G i , find the mask S with the largest intersection ratio in the predicted mask set j , and update C and U as: C←C+|G i ∩S j |,U←U+|G i ∪S j |; and mark S j been used;

[0082] 5-4) For the real mask {G i All masks in} are operated in step 5-3);

[0083] 5-5) For the prediction mask S k , if S k If no mark has been used, update U: U←U+|S k |;

[0084] 5-6) For the prediction mask {S k All masks in} are operated in step 5-4);

[0085] 5-7) Calculate AJI = C / U.

[0086] for Figure 7 The two predicted segmentation masks (c) and (d) have AJIs of 0.8162 and 0.7828 respectively, indicating that the method of the present invention has a significant advantage over PEA in the AJI index.

[0087] Finally, it should be noted that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Those skilled in the art will readily appreciate that various modifications and variations of the present invention are possible. Any modifications, equivalent substitutions, or improvements made within the scope of the present invention are intended to be included within the scope of the claims.

Claims

1. A multi-task medical image segmentation method, characterized in that: The following steps are involved: 1) obtaining a medical image dataset in the same field as the medical image to be segmented, wherein the dataset includes image data and annotated segmentation masks; 2) preprocessing the data set; 3) Build an improved SAM visual model that introduces task tokens; The parameters of the original SAM model are used as the initial parameters, a combination of Dice loss and Focal loss is used as the loss function, and only randomly selected pixel points in the area indicated by the true mask are used as prompts for the first stage of training; 4) Based on the improved SAM visual model, an affinity map prediction network and a category prediction network are introduced. The affinity map prediction network uses the UNETR model structure and takes the 32-channel features used to predict the segmentation mask in the improved SAM visual model as input. Its output is used to calculate the affinity map of the image in the horizontal and vertical directions. The category prediction network includes a task attention module and a classification head for predicting categories. The second stage of training is performed without using any prompts to obtain an optimized multi-task medical segmentation model. 5) Input the medical image to be segmented into the optimized multi-task medical segmentation model and output the final segmentation result.

2. The multi-task medical image segmentation method according to claim 1, characterized in that: The medical image dataset includes medical image data of different pathological areas and medical image data of different modalities; these medical image data have accurately labeled segmentation masks.

3. The multi-task medical image segmentation method according to claim 1, characterized in that: The preprocessing includes one or more of scaling, Gaussian noise, small-angle rotation, cropping, and normalization. The resolution of the preprocessed image is set to 512×512.

4. The multi-task medical image segmentation method according to claim 3, characterized in that: In step 3), the improved SAM visual model is obtained by making the following modifications on the basis of the original SAM model: 3-1) The input image resolution is set to 512×512. The image is processed into an input image with a resolution of 512×512 in the data preprocessing stage; 3-2) Replace the patch embedding module in the original SAM model image encoder with the DownConv module; 3-3) Replace the mask tokens in the original SAM model decoder with task tokens. Each task token corresponds to a segmentation task, and the task token is a learnable vector. 3-4) The self-attention module in the original SAM model image encoder is replaced with a task attention module, so that the image encoder can further adapt to multimodal and multi-type medical images.

5. The multi-task medical image segmentation method according to claim 4, characterized in that: The DownConv module is used to ensure that the feature map resolution of the attention module entering the image encoder is 64×64, which is consistent with the original SAM model; the DownConv module includes a first convolutional layer, a first BatchNorm normalization layer, a first GELU activation layer, a second convolutional layer, a second BatchNorm normalization layer, a second GELU activation layer, a third convolutional layer, a third BatchNorm normalization layer, a third GELU activation layer, and a fourth convolutional layer, wherein the step size of the first convolutional layer is set to 1 for image feature extraction, and the step sizes of all other convolutional layers are set to 2 for image feature extraction and image 2x downsampling; the BatchNorm normalization layer makes the network training more efficient and stable by adjusting the distribution of input data, and the GELU activation layer can increase the nonlinear modeling capability of the model.

6. The multi-task medical image segmentation method according to claim 4, characterized in that: The task attention module is obtained by adding a task label fusion module to the original self-attention module. The task label fusion module uses image features and task labels for cross-attention, and its output result is added to the output result of the original self-attention module as the output of the task attention module.

7. The multi-task medical image segmentation method according to claim 6, characterized in that: The task mark fusion module combines the task mark t s Input to the first linear mapping layer and the second linear mapping layer, output V and K respectively, and expand the 64×64×d feature map into a dimension of 64 2 The feature map f of ×d is input to the third linear mapping layer, and the output is Q, and then the calculation Then the attention matrix is ​​obtained through the softmax function Finally, perform matrix multiplication on the attention matrix and V to obtain the output of the task labeling fusion module.

8. The multi-task medical image segmentation method according to claim 1, characterized in that: The affinity map is used to characterize whether adjacent pixels belong to the same target. The affinity map is segmented into different regions by a graph partitioning algorithm to obtain a series of segmentation masks.

9. The multi-task medical image segmentation method according to claim 1, characterized in that: Step 4) The second phase of training includes: 4-1) The preprocessed input image X s Input the image encoder of the improved SAM visual model to obtain the image encoding F, where the input image X s From the sth segmentation task; 4-2) Improve the intersection-and-union (iou) token in the SAM visual model and the task tag t corresponding to the segmentation task s s Merge them into a whole, recorded as a label group, and input the image code F and the label group together into the decoder of the improved SAM visual large model. At this time, no prompt information is required, and the 32-channel image feature G used to calculate the mask is output, with a size of; 4-3) Input the 32-channel image feature G into the affinity map prediction module; first perform a 2x upsampling to obtain a new feature G′; then input G′ into the affinity map prediction network, outputting a feature map H; 4-4) After obtaining the image feature map H, calculate the cosine similarity of each pixel with the previous pixel in both the horizontal and vertical directions to obtain the affinity map in both directions, which is collectively referred to as affinity map A. p ; 4-5) Use the graph partitioning algorithm to divide the image into different regions according to the affinity graph and obtain a series of segmentation masks Where N represents the number of masks; 4-6) Encode the image T and the task tag t s The two are input together into the first part of the category prediction module, and the output is the feature F′; 4-7) Calculate each predicted segmentation mask m i With the true mask m gt The intersection-over-union ratio of each instance in m i The instance with the largest intersection-and-union ratio. If this maximum intersection-and-union ratio is greater than the set threshold γ, then m i Matching is achieved with this instance; according to the predicted segmentation mask m i , extract the corresponding area in the feature F′, that is, extract the features of the corresponding area, average these features, and then input them into the corresponding classification head C s , get the prediction y for the matching instance category p ;; 4-8) Construct loss function: The complete multi-task medical image instance segmentation model has two branches, one is the affinity map prediction branch, and the output is the affinity map A p , which is consistent with the real affinity graph A gt Calculate the mean square loss L2; ​​the second is the category prediction branch, the output is the category prediction y p , and the true category y gt Calculate the multi-classification cross entropy loss L CE ; The calculation formulas of the two loss functions are as follows: Among them, K s is the number of categories of the s-th segmentation task, n represents A p The number of elements in ; the complete loss function is the sum of the two: L=L2+L CE After calculating the loss, the model is trained using the Adam optimization algorithm.

10. A multi-task medical image instance segmentation system implementing the method according to any one of claims 1 to 9, characterized in that: include: A data processing module, used for preprocessing the medical image dataset and the medical image dataset to be segmented; Improved SAM visual large model construction and training module, used to build an improved SAM visual large model with task labels and perform the first phase of training; The multi-task medical segmentation model construction and training module introduces the affinity graph prediction network and category prediction network based on the improved SAM visual model to obtain a multi-task medical segmentation model. The multi-task medical segmentation model is trained as a whole to obtain an optimized multi-task medical segmentation model. The segmentation module is used to input the medical image to be segmented into the optimized multi-task medical segmentation model and output the segmentation result.

Images