Medical Image Segmentation Method

Abstract

The present invention provides a medical image segmentation method, including: optimizing a basic convolutional module and a model structure based on the nnU-Net framework to perform parameter expansion, so as to obtain an extensible medical image segmentation model; using a medical image data set to pre-train the medical image segmentation model to improve the generalization ability, performance and migration ability of the medical image segmentation model, so as to perform medical image segmentation tasks of different types and sources.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence technology, and particularly relates to a medical image segmentation method. Background Art

[0002] Medical image segmentation is a key task in the fields of computer vision and medical imaging, which includes separating the region of interest in the image from the background to assist in diagnosing, treating, and researching various diseases. With the development of deep learning technology, deep learning models based on the U-shaped network (U-Net), especially nnU-Net (no new-Net), have achieved remarkable results in the field of medical image segmentation. nnU-Net is an adaptive and general neural network structure for various medical image segmentation tasks. Based on the U-Net architecture, through automated network design and hyperparameter optimization, it realizes adaptive adjustment for different tasks. nnU-Net has achieved excellent results in multiple medical image segmentation competitions and has become a benchmark method in the field of medical image segmentation.

[0003] However, existing medical image segmentation models still have problems of scale limitations and insufficient generalization ability, which limit the performance and application scope of the models in practical applications. Specifically, there are certain limitations in the basic convolution module and model structure of nnU-Net, which are not suitable for large-scale expansion. In addition, the parameter scale of existing nnU-Net models is usually only in the millions. Compared with large-scale models in the fields of natural language processing and computer vision, their performance and generalization ability need to be further improved. Therefore, in view of these deficiencies, the present invention proposes an improved large-scale medical image segmentation model with scalability and transfer ability. Summary of the Invention

[0004] The purpose of the present invention is to provide a medical image segmentation method to solve the problems of scale limitations and insufficient generalization ability of existing nnU-Net models.

[0005] To solve the above technical problems, the present invention provides a medical image segmentation method, including:

[0006] Optimizing the basic convolution module and model structure based on the nnU-Net framework for parameter expansion to obtain a scalable medical image segmentation model;

[0007] Pre-training the medical image segmentation model using a medical image dataset to improve the generalization ability, performance, and transfer ability of the medical image segmentation model for medical image segmentation tasks of different types and sources.

[0008] Optionally, in the above-mentioned medical image segmentation method, it further includes:

[0009] Based on the nnU-Net framework, an extensible and transferable medical image segmentation network structure STU-Net is obtained, and the network depth and width are adjusted according to task requirements and computing resources;

[0010] Use a large-scale supervised pre-training strategy, including: pre-training on a public medical image segmentation dataset to improve the initial performance and transfer performance of the model.

[0011] Optionally, in the medical image segmentation method described above, it further includes:

[0012] Through the extensible and transferable medical image segmentation network structure STU-Net, perform automated configuration, parameter expansion, and cross-task transfer of medical image segmentation tasks, so as to optimize the basic convolution module, design an extensible network structure, and a large-scale supervised pre-training strategy based on the nnU-Net framework.

[0013] Optionally, in the medical image segmentation method described above, it further includes:

[0014] Optimizing the basic convolution module and model structure based on the nnU-Net framework includes: modifying the basic module of the nnU-Net framework, including introducing residual connections, adjusting the downsampling module, and changing the upsampling operation;

[0015] Introduce a downsampling module at the beginning of each encoder stage except the first resolution stage, and add several residual blocks after it;

[0016] Use interpolation operation to replace the original transposed convolution operation to avoid weight mismatch problems during cross-task transfer.

[0017] Optionally, in the medical image segmentation method described above, it further includes: performing the design of an extensible network structure, including:

[0018] Perform depth expansion and width expansion on the basis of STU-Net to achieve network structures of different scales;

[0019] Maintain the symmetry of the model, and simultaneously expand the depth and width of the encoder and decoder in each resolution stage.

[0020] Optionally, in the medical image segmentation method described above, the large-scale supervised pre-training strategy includes:

[0021] Pre-train the STU-Net model on the TotalSegmentator dataset to improve the generalization ability and cross-task transfer performance of the model;

[0022] During the pre-training process, adjust the training strategy, including increasing the number of training epochs to 4000 and using mirror data augmentation;

[0023] The pre-trained model is directly applied to downstream tasks with the same target segmentation classes or transferred through fine-tuning for learning.

[0024] The present invention also provides a computer-readable storage medium storing one or more programs, which can be executed by one or more processors to implement the steps of the above-mentioned medical image segmentation method.

[0025] The inventors of the present invention have found through research and summary that the prior art has the following problems:

[0026] (1) Scale limitation: Existing medical image segmentation models (such as nnU-Net) usually have a scale of millions of parameters. Compared with large-scale models in the fields of natural language processing and computer vision, their performance and generalization ability are limited.

[0027] (2) Structural limitation: The basic convolution module and model structure of existing models such as nnU-Net are not flexible enough and are not suitable for large-scale expansion.

[0028] (3) Insufficient migration ability: The migration ability of existing medical image segmentation models on different tasks and datasets is limited, resulting in a limited scope of application of the models in actual applications.

[0029] Based on the above insights, the medical image segmentation method provided by the present invention designs a large-scale medical image segmentation model with scalability and migration ability to solve the problems of scale limitation and insufficient generalization ability of medical image segmentation models in the prior art. For this purpose, the present invention improves the existing nnU-Net framework, optimizes the basic convolution module and model structure, thereby supporting expansion of different scales; the parameter scale can range from 10 million to 1.4 billion, including the largest 1.4 billion-parameter medical image segmentation model to date. To improve the model migration ability, the present invention uses a large-scale publicly available medical image dataset for pre-training. By implementing the present invention, the defects of existing medical image segmentation models in terms of scalability and migration ability can be effectively solved, providing strong support for model application between different tasks. Brief Description of the Drawings

[0030] Figure 1 is a schematic diagram of the segmentation performance of different models on the TotalSegmentator dataset;

[0031] Figure 2 is a schematic diagram of the architecture of the medical image segmentation model according to an embodiment of the present invention;

[0032] Figure 3 is a schematic diagram of the visualization results of STU-Net and nnU-Net of different scales;

[0033] Figure 4 It is a schematic diagram comparing the average DSC performance of STU-Net models of different scales trained with different proportions of training data on the TotalSegmentator validation set;

[0034] Figure 5 It is a schematic diagram comparing five specialized expert STU-Net models and a single general STU-Net model on the TotalSegmentator dataset. Specific embodiments

[0035] The present invention will be further described below in conjunction with specific embodiments with reference to the accompanying drawings.

[0036] It should be noted that the components in the respective drawings may be exaggerated for illustrative purposes and not necessarily drawn to scale. In the respective drawings, the same or functionally identical components are provided with the same reference numerals.

[0037] In the present invention, unless otherwise specified, "arranged on...", "arranged above...", and "arranged over..." do not exclude the presence of intermediate elements therebetween. In addition, "arranged on or above..." only represents the relative positional relationship between two components, and in certain cases, such as when the product direction is reversed, it can also be converted to "arranged under or below...", and vice versa.

[0038] In the present invention, each embodiment is only intended to illustrate the solution of the present invention and should not be construed as restrictive.

[0039] In the present invention, unless otherwise specified, the quantifiers "a" and "one" do not exclude the scenario of multiple elements.

[0040] It should also be noted here that in the embodiments of the present invention, for clarity and simplicity, only a part of the components or assemblies may be shown, but those of ordinary skill in the art can understand that, under the teaching of the present invention, the required components or assemblies can be added according to the specific scenario requirements. Additionally, unless otherwise stated, the features in different embodiments of the present invention can be combined with each other. For example, a certain feature in the second embodiment can be used to replace the corresponding or functionally identical or similar feature in the first embodiment, and the resulting embodiment also falls within the scope of disclosure or the scope of recording of the present application.

[0041] It should also be noted here that within the scope of the present invention, terms such as "identical", "equal", "equivalent" do not mean that the two values are absolutely equal, but allow for a certain reasonable error. That is to say, these terms also cover "substantially identical", "substantially equal", "substantially equivalent". By analogy, in the present invention, terms indicating direction such as "perpendicular to", "parallel to", etc. also cover the meanings of "substantially perpendicular to" and "substantially parallel to".

[0042] In addition, the numbering of the steps of each method of the present invention does not limit the execution order of the method steps. Unless otherwise specified, the method steps can be executed in different orders.

[0043] The following further elaborates on the medical image segmentation method proposed by the present invention in conjunction with the accompanying drawings and specific embodiments. According to the following description, the advantages and features of the present invention will be clearer. It should be noted that the accompanying drawings are in a very simplified form and use non-precise scales, solely for the purpose of facilitating and clearly assisting in the illustration of the embodiments of the present invention.

[0044] The object of the present invention is to provide a medical image segmentation method to solve the problems of the limited scale and insufficient generalization ability of the existing nnU-Net model.

[0045] To achieve the above object, the present invention provides a medical image segmentation method. Aiming at the limitations of existing medical image segmentation models in terms of scale, structure, and migration ability, an improved scalable medical image segmentation model is designed. By optimizing the basic convolution module and the model structure, it supports large-scale parameter expansion, and uses a large-scale medical image dataset for pre-training, thereby improving the generalization ability, performance, and migration ability of the model, and expanding its applicable range in practical applications.

[0046] Figure 1 Shows the segmentation performance of different models on the TotalSegmentator dataset. The area of each bubble is proportional to the FLOPs (floating-point operations per second) of the corresponding model at different scales. Different colors represent different models, and multiple bubbles of the same color represent the same model with different scales. The FLOPs are calculated based on an input size of 128×128×128.

[0047] Large-scale pre-trained models have had a profound impact on the development of deep learning. However, the current state-of-the-art models for medical image segmentation are still relatively small, with only tens of millions of parameters. Further expanding them to higher orders of magnitude has rarely been explored. An overall goal of exploring large-scale models is to train them on large-scale medical segmentation datasets to obtain better transfer capabilities. In this invention, a series of scalable and transferable U-Net models (STU-Net) are designed, with parameter scales ranging from 14 million to 1.4 billion. In particular, the STU-Net with 1.4 billion parameters is currently the largest medical image segmentation model. The STU-Net of this invention is based on the nnU-Net framework due to its popularity and excellent performance. This invention first optimizes the default convolutional blocks of nnU-Net to make them scalable. Then, different combinations of network depth and width expansion are empirically evaluated, and it is found that simultaneously expanding the depth and width of the model is optimal. This invention trains the scalable STU-Net model of this invention on the large-scale TotalSegmentator dataset and finds that increasing the model size brings stronger performance improvement. This observation indicates that large models have good prospects in medical image segmentation. In addition, this invention evaluates the transfer capabilities of the model for direct inference on 14 downstream datasets and further fine-tuning on 3 datasets, covering different modalities of images and different types of segmentation targets. The final simulation experiments observe that the pre-trained model of this invention performs well in direct inference and fine-tuning.

[0048] One of the main concerns of this invention through pre-training on large-scale datasets is how to design large-scale models so that they can be transferred to various medical segmentation tasks. However, large-scale models usually require more computing resources. This situation becomes more serious when used to train three-dimensional high-resolution medical images. Therefore, this invention also hopes that large-scale models can be scalable to different sizes to adapt to different computing budgets.

[0049] In summary, the objective of the present invention is to design a large-scale medical segmentation model with scalability and transferability. To this end, the present invention proposes a series of scalable and transferable U-Nets, called STU-Net, with parameter ranges from 10M to 1.4B. Notably, the 1.4B model is the largest model in the field of medical image segmentation to date. In addition to different model sizes for scalability, the present invention pre-trains on a large-scale dataset in a supervised manner to ensure the model has transferability. Specifically, the present invention constructs the model based on the nnU-Net framework because it has state-of-the-art baseline performance and is widely used. However, there are two obstacles to developing large-scale models using this framework. 1) For the scalability of the model, due to the gradient diffusion problem, the basic convolutional blocks in nnU-Net may not be suitable for expansion, and it is still unclear how to increase its parameters. 2) To evaluate the transferability of the model, fine-tuning on other downstream datasets is usually required. However, the model architectures in nnU-Net cannot be easily used for fine-tuning because their model structures are regarded as hyperparameters and are task-specific, and it is not feasible to change the model architecture during fine-tuning.

[0050] To overcome the first obstacle, the present invention improves the basic convolutional blocks of nnU-Net, for example, incorporating residual connections into its basic blocks to expand the model depth. Then, the present invention empirically evaluates different combinations of network depth and width. The inventors of the present invention found through research that simultaneously expanding the depth and width of the model is optimal. Therefore, the present invention can obtain the STU-Net-L model with 450 million parameters (2 times the depth, 2 times the width) and the huge STU-Net-H model with 1.4 billion parameters (3 times the depth, 3 times the width). Due to the unaffordable computational cost and GPU memory consumption on a single A100 GPU, the present invention does not perform further expansion.

[0051] To solve the second obstacle, the present invention replaces the transposed convolution in nnU-Net with nearest neighbor interpolation, followed by 1×1×1 convolution for upsampling, which avoids using the task-specific kernel / stride options in transposed convolution and makes the upsampling block transferable. The present invention further fixes the hyperparameters related to the model weights, for example, fixing the number of stages to 6 and using isotropic convolutional kernels, so that the model architecture remains the same during pre-training and fine-tuning, thus facilitating migration to other tasks.

[0052] To make the STU-Net of the present invention migratable, the present invention further adopts the large-scale TotalSegmentator dataset for supervised pre-training. TotalSegmentator is one of the largest datasets at present, containing 1204 whole-body CT images, among which 104 structures are accurately labeled. The STU-Net pre-trained on the large-scale dataset performs excellently in various downstream datasets during direct inference and fine-tuning, covering different modalities and segmentation targets. Benefiting from the large-scale supervised pre-training, the large-scale STU-Net of the present invention has excellent migration ability. The excellent performance shows that the main contributions of the present invention can be summarized as follows:

[0053] The present invention proposes an extensible STU-Net model, which can be extended according to different parameter sizes. The 1.4B STU-Net is the largest medical image segmentation model so far. Based on these STU-Net models, when trained on a large-scale dataset, the performance of the model increases significantly with the increase of the model size.

[0054] The present invention pre-trains the large-scale STU-Net model based on the large-scale TotalSegmentation dataset, which has good migration ability. It can be well generalized to various other datasets without additional adjustment or adaptation. It also performs excellently when fine-tuning downstream datasets.

[0055] The STU-Net of the present invention is based on the nnU-Net framework and benefits from its task-adaptive design. Therefore, the model of the present invention can adapt to different tasks and has good performance guarantee. At the same time, the present invention modifies the model architecture in nnU-Net to ensure that the model weights trained on the TotalSegmentator dataset can be easily transferred to downstream tasks.

[0056] Medical image segmentation is mainly dominated by two types of deep learning models: convolutional neural network (CNN)-based and Transformer-based models. U-Net is the earliest CNN model proposed for medical image segmentation. On this basis, residual connections, attention modules, and different feature aggregation strategies are applied to various tasks. Recently, the vision Transformer that has achieved success in natural image processing with the self-attention mechanism has been introduced into the medical image segmentation task. These medical image segmentation models only have a few million parameters, so they are not large enough. In addition, these models cannot be scaled and migrated under different computing budgets and simultaneously handle various medical image segmentation tasks.

[0057] In deep learning, improving model performance by scaling up the model size is a widely used method. Currently, the largest models are 540 billion and 22 billion parameters in the language and vision domains respectively. However, only a few works have scaled medical image segmentation models to large scales. In contrast, the present invention has successfully scaled the model size to an order of magnitude larger than previous works and evaluated their transfer capabilities on large-scale datasets. The model of the present invention is constructed based on the nnU-Net framework, which can automatically configure hyperparameters for different medical image segmentation tasks and achieve state-of-the-art performance on various medical image segmentation tasks.

[0058] The inventors found through research that nnU-Net adopts a symmetric encoder-decoder architecture based on skip connections. This architecture contains multiple resolution stages. Each stage includes two convolutional layers followed by instance normalization and leaky ReLU (denoted as Conv-IN-LeakyReLU). It does not contain residual connections, so simply stacking more layers at each stage may lead to gradient diffusion, making it difficult to optimize the entire model. This may limit the depth of nnU-Net and further limit its scalability. On the other hand, nnU-Net determines the input patch size and input spacing according to the dataset attributes. Then, the dataset-sensitive patch size and spacing are used to set hyperparameters related to the network architecture, such as the number of resolution stages, convolutional kernels, and downsampling / upsampling ratios. Therefore, these architecture-related hyperparameters vary between tasks, resulting in different network architectures for different tasks. As a result, a model trained on one task cannot be directly transferred to other tasks, which limits the evaluation of the model's transfer capabilities.

[0059] Therefore, the present invention makes modifications based on nnU-Net. The task-specific hyperparameters in nnU-Net can be divided into those related to model weights (e.g., convolution kernel size, number of resolution stages) and those unrelated to weights (such as pooling kernel, input block size, and input spacing). First, the present invention fixes the hyperparameters related to weights to keep the model architecture transferable to other tasks. Specifically, the present invention keeps the number of resolution stages for all tasks as 6 and uses isotropic convolution kernels (3, 3, 3) for all convolutional layers. For the hyperparameters unrelated to weights, the present invention adopts the default settings in nnU-Net to ensure its state-of-the-art performance. Comparing the settings of STU-Net with nnU-Net and 3D U-Net, Table 1 shows the comparison of different hyperparameters in the methods of nnU-Net, 3D U-Net, and STU-Net. The up (down) sampling ratios along different axes are (x, y, z). nnU-Net automatically configures task-specific hyperparameters, usually bringing state-of-the-art performance. However, this may reduce the transferability of the model. 3D-UNet uses a fixed network architecture and can easily transfer weights between different segmentation tasks, but its baseline performance is usually not as good as that of nnU-Net. The STU-Net of the present invention not only maintains the good performance of nnU-Net but also can be easily transferred.

[0060] Table 1

[0061]

[0062] Figure 2 Fig. 10 shows the schematic diagram of the STU-Net architecture of the present invention. Among them, it is modified on the basis of the nnU-Net architecture to improve its scalability and transferability. (a) is an overview of the STU-Net architecture. The blue arrows represent downsampling, while the brown arrows represent upsampling. (b) is the residual block for implementing large-scale models. (c) is the downsampling in the first residual block of each encoder stage. (d) and (e) are the starting module and segmentation head for input and output channel conversion. (f) is the weightless interpolation for upsampling, effectively solving the problem of weight mismatch between different tasks.

[0063] In the encoder and decoder of nnU-Net, each stage consists of a basic block. Each basic block includes two Conv-IN-LeakyReLU layers. When increasing the number of basic blocks in each stage, optimization problems may occur due to gradient diffusion. To solve these problems, the present invention introduces residual connections in the basic blocks. Except for the first resolution stage, each stage in the encoder starts with Figure 2 (b) the downsampling block in, followed by several Figure 2(c) The residual block. Different from nnU-Net which uses a separate convolution for downsampling, the present invention integrates downsampling into the first residual block of each stage. Such a downsampling block has two branches, namely Figure 2 the left branch and the right branch in (b). The left branch has two 3×3×3 convolutions with different strides, i.e., the stride of the former is 1 and the stride of the latter is 2. The right branch uses a 1×1×1 convolution kernel size with a stride of 2 to match the output shape of the left branch. This downsampling block has a residual architecture similar to Figure 2 the conventional residual block in (c), making the entire architecture concise and consistent.

[0064] The adjustments to the upsampling module include: By default, nnU-Net uses transposed convolutions with a stride of for upsampling. However, for different tasks, even at the same resolution stage, the convolution kernel and stride may vary between (2,2,2) and (2,2,1), which results in different weight shapes in the transposed convolutions for different tasks. When trying to transfer weights from one task to another, this further leads to weight mismatches. To solve this problem, the present invention uses nearest neighbor interpolation together with a 1×1×1 convolution with a stride of 1 to replace the transposed convolution because weightless interpolation has no weight shape problem. According to the experimental results (see Table 6), the present invention uses nearest neighbor interpolation for upsampling, and the experimental results show that nearest neighbor interpolation not only provides a faster processing speed but also performs comparably to cubic linear interpolation.

[0065] Based on the above improved modules, the present invention can obtain a new network architecture called STU-Net. The present invention further expands the depth and width of each stage of the model to achieve model expansion. Deeper networks usually have a larger receptive field and better representation ability. Wider networks tend to extract richer multi-scale features at each layer. Depth scaling and width scaling are not independent. Scaling the depth and width of the network simultaneously can achieve better accuracy and computational efficiency. To simplify the scaling problem, the present invention maintains the symmetry of the model, which means that the present invention scales the encoder and decoder simultaneously and scales the depth and width at the same ratio at each resolution stage. The present invention shows different scales of STU-Net in Table 2, "Base (basic type), Large (large type), Huge (huge type)". The depth is the number of residual blocks in each resolution stage, and the width is the number of channels in each resolution stage. The parameter calculation is based on a single-channel input and a 104-channel output (for 104 categories in TotalSegmentator).

[0066] Table 2

[0067]

[0068] The present invention pre-trains the STU-Net on the TotalSegmentator dataset. In the STU-Net, the last 1×1×1 convolutional layer for segmentation output has 104 channels, corresponding to the total number of annotated categories in TotalSegmentator. To make the pre-trained model more general and transferable, the present invention does not strictly follow the standard training process in nnU-Net but makes some modifications. Compared with the default 1000 training epochs in nnU-Net, the present invention pre-trains the model for 4000 epochs. In addition, the present invention finds that pre-training using mirror data augmentation can improve the transfer performance of downstream tasks.

[0069] The pre-trained model of the present invention can directly perform inference on downstream datasets containing CT images with the target segmentation categories among the upstream 104 classes without further tuning. If the downstream task has new labels or different modalities, the present invention uses the trained model as initialization and randomly initializes the segmentation output layer to match the number of target output classes. For fine-tuning, the segmentation head is randomly initialized, while the weights of the remaining layers are loaded from the pre-trained model. The fine-tuning learning rate of these weights (0.1×) is smaller than that of the segmentation head, thus obtaining better results.

[0070] The present invention trains STU-Nets of different scales on the TotalSegmentator dataset, which contains 1204 images with 104 anatomical structures (including 27 organs, 59 bones, 10 muscles, and 8 blood vessels). It covers most of the clinical segmentation targets of normal structures throughout the body. All images are resampled to an isotropic resolution of 1.5×1.5×1.5 mm. The present invention follows the original data partition, using 1081 cases for training, 57 cases for validation, and 65 cases for final testing. It should be noted that due to data privacy reasons, the face has been blurred. The present invention evaluates the direct inference of STU-Net on 14 public datasets and further fine-tuning on 3 public datasets to test the transferability of the trained model. The detailed attributes of these downstream datasets are shown in the appendix.

[0071] Evaluation metric: The Dice similarity coefficient (DSC) is used as the evaluation metric. A higher DSC score indicates better segmentation performance. For fair comparison, the training results of the model in the last round are reported instead of the best one.

[0072] Implementation details: All experiments were run in an environment of Python 3.8, CentOS 7, Pytorch 1.10, and nnU-Net 1.7.0. Roughly follow the default data preprocessing, data augmentation, and training procedures in nnU-Net. Use the SGD optimizer with Nestrov momentum 0.99 and weight decay 1e-3. The batch size was fixed at 2, and each epoch consisted of 250 iterations. The learning rate started at 0.01 and decayed following the polynomial learning rate strategy: (1 - epoch / 1000)^0.9. Data augmentation methods of brightness, gamma, rotation, scaling, mirroring, and elastic deformation were adopted during training. The pre-trained patch size on TotalSegmentator was 128×128×128. The fine-tuning patch size on downstream tasks was automatically configured by nnU-Net. The model was trained on an NVIDIA Tesla A100 GPU with 80GB VRAM.

[0073] To verify the effectiveness of STU-Net of different scales, the model was compared with other state-of-the-art methods on the TotalSegmentator validation set. For fair comparison, all models were trained for 1000 epochs in the nnU-Net framework. To improve the performance of other methods, the optimizers, learning rates, and learning rate decay strategies reported in their papers were used. Following the method in the original SwinUNETR paper, where the feature size of SwinUNETR-B is twice that of SwinUNETR-S, the feature size of SwinUNETR-B was further doubled to 96 to obtain SwinUNETR-L. This enabled observing the performance of the SwinUNETR model when scaled to a larger size.

[0074] As Figure 1 shown in Table 3, the STU-Net-B model exceeded the best CNN-based model nnU-Net and the best Transformer-based model SwinUNETR-B in terms of the average DSC for all classes by 0.36% and 4.48% respectively. Further expanding the base model of the present invention to large and giant sizes increased the average DSC scores by 1.59% and 2.94% respectively. It was also observed that expanding the SwinUNETR model led to a significant performance improvement, but it was still worse than STU-Net-B. The STU-Net-H of the present invention achieved the highest average DSC among all classes and five sub-class groups of the TotalSegmentator dataset.

[0075] The experimental results demonstrate the effectiveness of the model structure proposed in the present invention and the effectiveness of the model extension method. As shown in Table 3, the segmentation results of different methods on the TotalSegmentator validation set are evaluated by the average DSC (%↑). Due to space limitations, all 104 classes are divided into 5 groups, and the average DSC of these 5 groups of classes and all classes are reported separately. *: The maximum number of features in nnU-Net is changed from 320 to 512 to match the parameters of the STU-Net-B model. ·□: SwinUNETR-L is obtained by changing the feature size in SwinUNETR-B from 48 to 96.

[0076] Table 3

[0077]

[0078] As shown in Table 4, the transferability of the models trained on TotalSegmentator is evaluated. The average DSC (%) of direct inference on various downstream datasets is reported. AMOS-CT represents the results only in the CT modality.

[0079] Table 4

[0080]

[0081]

[0082] The present invention tests the transferability of the trained model in the following two ways: 1) directly inferring on a downstream CT dataset containing the same 104 classes as TotalSegmentator; 2) fine-tuning the trained model on 3 downstream datasets containing classes not in TotalSegmentator (e.g., lesions) and modalities other than CT (e.g., MR, PET).

[0083] The present invention uses the pre-trained STU-Net to perform direct inference on 14 downstream datasets, and the annotated classes of these datasets are all included in the TotalSegmentator dataset. These 14 datasets contain a total of 2494 cases, which is a strong evaluation platform for assessing the transferability of models pre-trained on large-scale datasets.

[0084] The TotalSegmentator dataset provides detailed annotations for normal organs, including separate labels for left and right organs (without lesion annotations). However, downstream datasets may contain lesion annotations within the organ regions or combined annotations of left and right organs. These differences in annotations may introduce additional labels that never appeared in TotalSegmentator, leading to label inconsistencies. To address this issue, the present invention merges the lesion annotations with the corresponding organ labels and, if necessary, merges the left and right organ annotations to create modified labels for evaluation purposes. Table 4 shows that, when pre-trained on TotalSegmentator, larger-scale models generally have higher average DSC scores in these 14 datasets. This conclusion generally also applies to each specific dataset. The better performance indicates that the carefully designed large-scale STU-Net of the present invention can have better transferability when pre-trained on large-scale datasets.

[0085] The present invention fine-tunes pre-trained STU-Nets of different sizes and applies them to three downstream datasets that contain new anatomical structures, modalities, and domains: AutoPET22 (with lesions and PET modality), AMOS22 (with MR modality), and FLARE22 (with multiple domains). Due to the architectural improvements made by the present invention to nnU-Net, all model weights except for the weights of the segmentation head can be easily transferred to these downstream tasks. It should be noted that these downstream datasets, such as AutoPET, may contain multimodal inputs, and the number of channels may be more (e.g., one channel) than the input data of the pre-training dataset (e.g., TotalSegmentator). More input channels require the first convolutional layer of the model to have correspondingly more channels, further resulting in a mismatch in the convolutional weight shapes between pre-training and downstream fine-tuning tasks. In this case, the present invention copies the channels of the weights of the first layer of the pre-trained model from one channel to multiple channels to adapt to the downstream tasks. In the AutoPET and AMOS datasets, the presence of multiple modalities means that fine-tuning and evaluation need to be performed separately for each modality (denoted by the suffix of the modality name in Table 5), as well as fine-tuning and evaluation on the mixed modalities. As shown in Table 5, fine-tuning the STU-Net model of the present invention (pre-trained on TotalSegmentator) has better segmentation performance on downstream datasets than models trained from scratch. In particular, the large model of the present invention (STU-Net-H-ft) outperforms all other models, achieving the highest average DSC of 80.69% on these downstream datasets. This observation emphasizes the importance of pre-training and large models in improving segmentation performance. The quality comparison of the segmentation results of the model of the present invention and nnU-Net is as Figure 3 shown.

[0086] As shown in Table 5, the results of fine-tuning on 3 downstream datasets are presented. The average DSC (%↑) is evaluated. AMOS-CT (or -MR) represents the results in only the CT (or MR) modality; otherwise, the results are for the mixed modality. The same holds for AutoPET. The suffix "ft" indicates fine-tuning.

[0087] Table 5

[0088]

[0089] Figure 3 Visualization results showing STU-Net and nnU-Net of different scales on various medical image datasets are presented. Each row shows a representative case from a different dataset, including the first row for the FLARE22 dataset, the second row for the AMOS dataset with CT images, the third row for the AMOS dataset with MR images, the fourth row for the AutoPET dataset with CT images, and the fifth row for the AutoPET dataset with PET images. The seven columns from left to right correspond to the original image, the ground truth label (gt), the nnU-Net result, and the STU-Net-B-ft, STU-Net-L-ft, and STU-Net-H-ft results of the present invention.

[0090] This visualization further highlights the advantages of the STU-Net model of the present invention during fine-tuning on downstream datasets and demonstrates the excellence of the method of the present invention in various medical imaging scenarios. Notably, the improvement on the AutoPET dataset is the most significant, which may be because pre-training on TotalSegmentator provides information about normal anatomy throughout the body, which is complementary to the only annotation information in AutoPET - the lesion information (and can serve as prior information). Therefore, fine-tuning the pre-trained model can effectively enhance the model's ability to segment lesions on AutoPET.

[0091] When training from scratch, the giant model is slightly inferior to the large model, probably because the limited training samples on these datasets cannot further improve the larger model. After pre-training with sufficient training samples from TotalSegmentator, the giant model can benefit from fine-tuning and generally significantly outperforms the large model.

[0092] Notably, even when TotalSegmentator focuses on CT scans, the pre-trained model shows significant performance improvement when fine-tuned on non-CT modalities, such as the AMOS-MR and AutoPET-PET datasets. This indicates that pre-training helps learn cross-modal basic features and structures, beyond modality-specific features.

[0093] The STU-Net-B model proposed by the present invention is an improved version of the default nnU-Net architecture, including several improvements. Table 6 compares the segmentation performance of different STU-Net-B architecture variants on the TotalSegmentator validation set. The STU-Net-B model of the present invention uses nearest neighbor interpolation for upsampling and introduces downsampling in the first residual block. To evaluate its performance, the present invention conducted an analysis comparing it with alternative upsampling methods using separate convolutions for downsampling and transposed convolutions or trilinear interpolation.

[0094] The present invention first increased the maximum number of features of the standard nnU-Net from 320 to 512 to match the parameters of the STU-Net-B of the present invention and called it nnU-Net*. The performance of nnU-Net* is better than that of the standard nnU-Net, but slightly lower than that of the STU-Net-B of the present invention. This comparison demonstrates the effectiveness of the improvements in the STU-Net-B.

[0095] Then, the present invention explored the downsampling design in STU-Net by introducing a variant that uses convolutional downsampling instead of integrating the downsampling process into the first residual block of each stage. This modification reduces performance and increases computational cost. The present invention further studied the improvement of upsampling and designed two variants of STU-Net-B. The first uses transposed convolutions to replace the default interpolation and convolution-based upsampling. This results in a 0.16% reduction in performance and makes the weights unable to be transferred for downstream fine-tuning. The second uses trilinear (or cubic-linear) interpolation to replace the nearest neighbor interpolation. This change reduces performance and slows down the running speed. Generally speaking, the default upsampling design in STU-Net has better performance, faster running speed, and better transferability.

[0096] As shown in Table 6, on the TotalSegmentator validation dataset, the segmentation performance of different architecture models is evaluated using the average DSC (%↑). *: The present invention changed the maximum number of features of the standard nnU-Net from 320 to 512 to match the parameter quantity of the STU-Net-B of the present invention. Abbreviations: DS (downsampling), US (upsampling)

[0097] Table 6

[0098]

[0099] Therefore, the improvements proposed by the present invention not only improve the effectiveness and efficiency of nnU-Net, but also endow it with important weight transferability and scalability properties. These properties are crucial for further expanding the model and promoting transfer learning.

[0100] The present invention applies three different scaling strategies to nnU-Net and STU-Net-base, that is, scaling is performed using different depth coefficients d ∈ [1.0, 2.0, 3.0, 4.0], width coefficients w ∈ [1.0, 2.0, 3.0, 4.0], and simultaneous scaling of depth and width coefficients [1.0, 2.0, 3.0]. The coefficient d (or w) represents that the depth (width) of each stage is enlarged by d (or w) times. Table 7 shows the results of different scaling strategies on TotalSegmentator. First, generally, wider nnU-Net* and STU-Net can achieve better performance, while deeper models are not necessarily better. Therefore, compared with depth scaling, width scaling is more effective in improving model performance on large-scale datasets, but it will significantly increase the computational consumption. Second, compared with the nnU-Net model, since the STU-Net model adopts a residual design, expanding the depth to a certain extent can better improve its performance, such as 87.12 → 87.72 → 87.99 → 87.58 vs. 86.94 → 85.65 → 83.70 → 81.45 of nnU-Net. Third, compound scaling, that is, simultaneously increasing the depth and width coefficients, is more effective and efficient in improving the performance of the STU-Net of the present invention. Finally, even using the same scaling strategy and similar parameters and FLOPs, the STU-Net of the present invention has better performance than nnU-Net* in all settings, which again verifies the effectiveness of the improvement of the present invention.

[0101] As shown in Table 7, segmentation is performed using models of different sizes on the TotalSegmentator dataset, and the average DSC (%↑) is evaluated. *: The present invention changes the maximum number of features of the standard nnU-Net from 320 to 512 to match the parameter quantity of the STU-Net-B of the present invention.

[0102] Table 7

[0103]

[0104] As shown in Table 8, comparison of the segmentation performance of STU-Net-L with different pre-training settings. The average DSC (%↑) is evaluated. The first and second rows compare whether to use mirror data augmentation for training the model from scratch (without pre-training). The third to sixth rows compare the performance of different mirror data augmentations with pre-training and fine-tuning. The last four rows study the influence of different numbers of pre-training epochs on the fine-tuning results.

[0105] Table 8

[0106]

[0107]

[0108] The present invention conducts an empirical study on pre-training, mirror data augmentation, and the number of training epochs during pre-training. First, the present invention studies the effectiveness of large-scale pre-training. The present invention uses the STU-Net-L model pre-trained on the TotalSegmentator dataset and fine-tunes it on the FLARE and AutoPET datasets to conduct this study. From the comparison between the first row (without pre-training) and the third row (with pre-training) in Table 8, the present invention finds that, without pre-training, the performance of the model is better than that with pre-training, being 0.84% and 3.67% respectively. Such better performance proves the effectiveness of large-scale pre-training.

[0109] Secondly, the present invention studies the mirror data augmentation used in the pre-training and fine-tuning stages. In pre-training, mirroring is useful for improving the DSC of downstream tasks because mirror data augmentation (rows 3 - 5) significantly improves the DSC of row 6 (without mirror data augmentation). It is worth noting that the mirror data augmentation used simultaneously in both the pre-training and fine-tuning stages obtains the best results.

[0110] Finally, when mirror data augmentation is used in both pre-training and fine-tuning, the inventors of the present invention find that different numbers of pre-training epochs h also affect the performance (see the last four rows). The best result is obtained at 4000 epochs, which can ensure that the model is sufficiently trained and converges.

[0111] Figure 4 Shows the comparison of the average DSC (%↑) performance of different scales of STU-Net models trained with different proportions of training data on the TotalSegmentator validation set.

[0112] The present invention studies the influence of the dataset size on the performance of the model when training the model on the TotalSegmentator dataset. It should be noted that different proportions of training cases are obtained through a stratified random selection process, ensuring that higher proportions of training cases also include lower proportions of data.

[0113] As Figure 4As shown, increasing the model size improves the segmentation performance on the TotalSegmentator subset, regardless of the number of training cases. For example, even when trained with only 5% of the data, the performance of STU-Net-H is better than that of STU-Net-S. Similarly, even when trained with only 20% of the data, STU-Net-H is better than STU-Net-B. These results indicate that large-scale models are more data-efficient for medical image segmentation compared to smaller models. Additionally, the performance of different models continuously improves as the number of cases increases, and this trend has not reached saturation. These observations suggest that increasing the number of training cases based on the TotalSegmentator dataset can further improve performance.

[0114] Figure 5 Figure 4 shows a comparison of five specialized expert STU-Net models and a single general STU-Net model on the TotalSegmentator dataset. Each expert model targets the segmentation of one of the five subcategories (i.e., organs, vertebrae, heart, muscles, and ribs), while the general model is trained on all 104 categories. The performance is measured using the average DSC across various anatomical categories: the five subcategories and the overall performance metric TotalSeg. The STU-Net architectures (S, B, L, and H) represent the expert and general models respectively. Lighter colors represent the expert models, and darker colors represent the general model.

[0115] The present invention evaluates a general STU-Net model trained on all 104 classes of the TotalSegmentator dataset, compared with five expert STU-Net models, each targeting one of the five subcategories (the same as in Table 3). Additionally, the present invention also studies the impact of model size on the performance of expert and general models.

[0116] As Figure 5 shown, as the model size increases, the performance of both expert and general models generally improves. The expert models perform excellently in the organ, vertebra, and heart subcategories, while the general model performs better in the muscle and rib subcategories. For the largest model (STU-Net-H), the general model outperforms the expert models, achieving the highest average DSC score of 90.06% across all categories of the TotalSeg dataset, while the highest average DSC score of the expert models is 89.07%.

[0117] The results show that while the expert model may perform better than the general model in specific sub - categories, the general model consistently provides strong performance across different anatomical structures. As the model size increases, the performance gap between the expert model and the general model varies across different sub - categories: for organs and vertebrae, the gap narrows; for the heart, the gap reverses; for muscles and ribs, the gap widens. These findings suggest that as the model size increases, the general model can segment multiple categories simultaneously and demonstrates promising performance improvements.

[0118] The present invention introduces a series of scalable and transferable medical image segmentation models based on the nnU - Net framework, called STU - Net. The largest model of the STU - Net model of the present invention consists of 1.4 billion parameters and is the largest medical image segmentation model to date. By training the STU - Net model on the large - scale TotalSegmentator dataset, the present invention demonstrates that when the model size is scaled up, transferring it to various downstream tasks can significantly improve performance. This highlights the potential of large - scale models in the field of medical image segmentation. In addition, the STU - Net - H model trained on the TotalSegmentator dataset shows strong transferability in multiple downstream datasets under direct inference and further fine - tuning scenarios. This observation emphasizes the practical value of leveraging large - scale pre - trained models for medical image segmentation tasks. In summary, the development of the scalable and transferable STU - Net model is expected to advance the state - of - the - art in medical image segmentation, opening up new avenues for research and innovation in the field of medical image segmentation.

[0119] In summary, the above - mentioned embodiments have elaborated on different configurations of the medical image segmentation method. Of course, the present invention includes but is not limited to the configurations listed in the above embodiments. Any content obtained by transforming on the basis of the configurations provided in the above embodiments falls within the scope protected by the present invention. Those skilled in the art can draw inferences by analogy based on the content of the above embodiments.

[0120] Each embodiment in this specification is described in a progressive manner. The key point of each embodiment is to illustrate the differences from other embodiments. The same or similar parts among the embodiments can be referred to each other. For the system disclosed in the embodiments, since it corresponds to the method disclosed in the embodiments, the description is relatively simple, and the relevant parts can be referred to the description of the method part.

[0121] The above description is only a description of the preferred embodiments of the present invention and does not limit the scope of the present invention in any way. Any changes and modifications made by those of ordinary skill in the art of the present invention based on the above disclosure fall within the scope of protection of the claims.

Claims

1. A medical image segmentation method, characterized in that, Including: Optimizing the basic convolutional module and model structure based on the nnU-Net framework for parameter expansion to obtain an extensible medical image segmentation model; And Using a medical image dataset to pre-train the medical image segmentation model to improve the generalization ability, performance, and transfer ability of the medical image segmentation model for medical image segmentation tasks of different types and sources; Among them, optimizing the basic convolutional module and model structure based on the nnU-Net framework includes: modifying the basic module of the nnU-Net framework, including introducing residual connections, adjusting the downsampling module, and changing the upsampling operation; Introducing a downsampling module at the beginning of each encoder stage except the first resolution stage and adding several residual blocks after it; and Using interpolation operations to replace the original transposed convolution operations to avoid weight mismatch problems during cross-task transfer, where nearest neighbor interpolation and 1×1×1 convolution with a stride of 1 are used together to replace the transposed convolution; Among them, the method further includes: performing extensible network structure design, including: Performing depth expansion and width expansion based on STU-Net to achieve network structures of different scales; and Maintaining the symmetry of the model and simultaneously expanding the depth and width of the encoder and decoder in each resolution stage.

2. The medical image segmentation method according to claim 1, wherein Also including: Obtaining an extensible and transferable medical image segmentation network structure STU-Net based on the nnU-Net framework and adjusting the network depth and width according to task requirements and computing resources; And Using a large-scale supervised pre-training strategy, including: pre-training on a public medical image segmentation dataset to improve the initial performance and transfer performance of the model.

3. The medical image segmentation method according to claim 1, wherein Also including: Through the extensible and transferable medical image segmentation network structure STU-Net, performing automated configuration, parameter expansion, and cross-task transfer of medical image segmentation tasks to implement optimizing the basic convolutional module based on the nnU-Net framework, designing an extensible network structure, and a large-scale supervised pre-training strategy.

4. The medical image segmentation method according to claim 1, wherein, The large-scale supervised pre-training strategy includes: Pre-training the STU-Net model on the TotalSegmentator dataset to improve the generalization ability and cross-task transfer performance of the model; During the pre-training process, adjusting the training strategy, including increasing the number of training rounds to 4000 rounds and using mirror data augmentation; and The pre-trained model is directly applied to downstream tasks with the same target segmentation categories or transferred through fine-tuning for learning.

5. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores one or more programs that can be executed by one or more processors to implement the steps of the medical image segmentation method as described in any one of claims 1-4.

Images