Lightweight general high-definition medical image segmentation network

Abstract

The invention provides a lightweight general high-definition medical image segmentation network, and the network employs an encoder-decoder architecture as the basis, and integrates a shallow ViT image encoder and a patch-level cross-multi-scale decoder, a channel interaction layer and a mark interaction layer are adopted in the shallow ViT image encoder to learn low-layer features of a given image, a long-distance dependency relationship is processed with linear complexity to obtain an embedded image, a patch-level cross-multi-scale decoder learns decoding information by creating a query embedded word, a self-attention mechanism is adopted to update mask query, and a patch-level cross-multi-scale decoder obtains a patch-level cross-multi-scale image. And the embedded image is updated through a bidirectional multi-scale cross attention layer to obtain a segmentation mask. The method has the beneficial effects that the expressive force of linear attention can be improved, and the memory cost is reduced.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image networks, and more particularly, to a lightweight general high-definition medical image segmentation network. Background Art

[0002] Medical imaging involves multiple modalities, such as two-dimensional microscopy, three-dimensional computed tomography, and video colonoscopy, which are regarded as important tools for identifying disease types and grades. Traditional manual image analysis methods are very time-consuming. In this challenging context, computer-aided diagnosis has received great attention from researchers and medical experts. In particular, pixel-level segmentation of disease or organ regions is a crucial step. Among them, UNet, as an encoder-decoder deep learning architecture, has shown excellent adaptability in various medical segmentation tasks. However, classical UNet variants use convolutional neural networks (CNNs) to learn feature representations. Due to the limitation of the receptive field, it is difficult to capture global context.

[0003] In the past few years, Vision Transformer (ViT) has become a promising alternative to CNN in various computer vision tasks. It uses self-attention mechanisms for long-sequence modeling. Although many ViT-based U-shape frameworks have shown excellent performance in different medical image segmentation tasks, the quadratic complexity of ViT rapidly increases the number of parameters in the model, resulting in high computational costs and high latency. In addition, mobile devices have shown outstanding applicability in achieving fast and more convenient medical image analysis and clinical diagnosis. These devices usually have a small memory capacity. Therefore, it is necessary to reduce the size of existing transformer models to meet the requirements of such devices. Linear attention is a potential solution that reduces the computational complexity from O(n^2) to O(n) by replacing the non-linear Softmax function in the self-attention mechanism with a linear normalization. Although this operation is very efficient, its insufficient feature representation greatly reduces the performance of ViT.

[0004] In addition, the generation of high-resolution segmentation masks can also provide great assistance for accurate disease diagnosis. The decoder of the classical method based on UNet combines high-level and low-level semantic information using a bottom-up structure, so it is applicable to mask generation of any scale. In addition, some methods in the prior art apply multi-scale feature fusion to each skip connection layer, thereby improving the accuracy of mask prediction. However, when faced with high-resolution segmentation mask prediction, these pyramid decoding operations based on large-size feature maps will result in high memory costs and higher model complexity. Given the high cost of GPU-equipped systems in many research institutions, reducing the model size and GPU memory usage during training and deployment is crucial for the widespread application of resource-efficient and accessible medical diagnostic technologies. Summary of the Invention

[0005] The technical problem to be solved by the present invention is to improve the expressiveness of linear attention and reduce memory costs. To overcome the defects of the above prior art (or related art), the present invention provides a lightweight general high-definition medical image segmentation network.

[0006] The present invention provides a lightweight general high-definition medical image segmentation network, which is based on an encoder-decoder architecture and integrates a shallow ViT image encoder and a patch-level cross multi-scale decoder. Among them, the shallow ViT image encoder uses a channel interaction layer and a token interaction layer to learn the low-level features of a given image, and processes long-range dependencies with linear complexity to obtain an embedded image. The patch-level cross multi-scale decoder learns decoding information by creating query embedding words, uses a self-attention mechanism to update the mask query, and updates the embedded image through a bidirectional multi-scale cross-attention layer to obtain a segmentation mask.

[0007] Compared with the prior art, a lightweight general high-definition medical image segmentation network of the present invention has the following advantages: In the present invention, an encoder-decoder architecture is adopted as the basis, integrating a shallow ViT image encoder and a patch-level cross multi-scale decoder, to capture long-distance global information in medical images with linear computational complexity, and combining a gating mechanism to overcome the problem of insufficient expressiveness of linear attention; and introducing a patch-level cross multi-scale decoder, using precise high-dimensional feature space decoding to achieve a direct mapping between low-resolution and high-resolution segmentation masks, eliminating the requirement of pyramid ratio recovery, and reducing the memory cost during the training process.

[0008] In a possible implementation manner, the shallow ViT image encoder uses two MBConvs as the channel interaction layer to learn the low-level features of the given image, and each MBConvs includes a 1 * 1 convolution connected in sequence , a 3 * 3 depth convolution and a 1 * 1 projection convolution perform patch embedding on the given image for linear complexity processing to obtain the embedded image.

[0009] In a possible implementation, the expression for the MBConvs to perform patch embedding on the given image is as follows: where represents the given image; H represents the height of the given image; W represents the width of the given image; S represents the pre-configured patch size; C represents the embedding channel; represents the activation function.

[0010] In a possible implementation, the shallow ViT image encoder adopts a dual-gated linear attention mechanism to form the token interaction layer for learning global features, calculates the similarity between all query-key pairs through the dual-gated linear attention mechanism to update the given image, and processes the given image using a multi-layer perceptron and layer normalization method to obtain the embedded image.

[0011] In a possible implementation, the expression of the dual-gated linear attention mechanism is as follows: where represents the th row of matrix V; H represents the height of the given image; W represents the width of the given image; S represents the pre-configured patch size; represents the kernel function; represents the given image; , , represents a learnable linear projection matrix; C represents the embedding channel; d represents the projection dimension; represents the SiLU function.

[0012] In a possible implementation, the calculation formula for obtaining the embedded image by processing the given image using the multi-layer perceptron and the layer normalization method is as follows: where, represents the embedded image; represents processing using the multi-layer perceptron; represents processing using the layer normalization method; represents the SiLU function; represents the given image; represents a learnable linear projection matrix; C represents the embedding channel; d represents the projection dimension.

[0013] In a possible implementation, the expression for the patch-level cross multi-scale decoder to update the mask query is as follows: where, represents the bidirectional multi-scale cross-attention layer; represents the self-attention mechanism; represents the query embedding word; represents the number of predicted classes; represents the decoding dimension.

[0014] In a possible implementation, the expression for the patch-level cross multi-scale decoder to update the embedded image is as follows: Among them, represents the embedded image; represents the position encoding; d represents the projection dimension; represents a multi-scale module for enhancing semantic information at different scales; represents the given image; represents a learnable linear projection matrix.

[0015] In a possible implementation manner, the expression of the patch-level cross multi-scale decoder for obtaining the segmentation mask is as follows: Among them, represents the segmentation mask; represents performing transposed convolution twice; represents a bilinear interpolation function for directly restoring the mask size; represents processing using the multi-layer perceptron; H represents the height of the given image; W represents the width of the given image; represents the number of predicted classes. Description of the Drawings

[0016] Figure 1 is a schematic diagram of the overall structure of the present invention; Figure 2 is a schematic diagram of the evaluation results of the two-dimensional medical image dataset of dermoscopy, X-ray, fundus, ultrasound, and microscope of the present invention; Figure 3 is a schematic diagram of the evaluation results of the three-dimensional medical image segmentation model on various organs of the present invention; Figure 4 is a schematic diagram of the medical video segmentation experiment on two different colonoscopy datasets of the present invention; Figure 5 is a schematic diagram of the evaluation results of the 2D medical image segmentation of the present invention; Figure 6 is a schematic diagram of the evaluation results of the 3D medical image segmentation of the present invention; Figure 7 is a schematic diagram of the evaluation results of the medical video segmentation of the present invention; Figure 8 Schematic diagram of the evaluation result of zero - point generalization of the present invention. Detailed implementation manners

[0017] First of all, those skilled in the art should understand that these implementation manners are only used to explain the technical principles of the embodiments of the present invention, and are not intended to limit the protection scope of the embodiments of the present invention. Those skilled in the art can adjust them according to needs to adapt to specific application scenarios.

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

[0019] See Figure 1 , an embodiment of the present invention discloses a lightweight general - purpose high - definition medical image segmentation network, defined as YumiNet. Specifically, first, SViT is introduced as a lightweight image encoder. The SViT of YumiNet is divided into two stages, with a total of only eight blocks. In the first stage, a channel interaction layer is used to efficiently learn low - level representations. In the second stage, a token interaction layer is constructed, combined with dual - gated linear attention to simulate long - range dependencies. The core operation uses a gating mechanism to improve the expressiveness of linear attention. Secondly, the present invention designs a PLCM decoder to generate high - resolution segmentation masks, which focuses on decoding high - dimensional but small - size feature spaces, eliminating the need for pyramid scale recovery. In addition, the SAM series has shown overwhelming performance in natural image and video segmentation tasks. In order to utilize the base model pre - trained with large - scale data, the present invention adopts a knowledge distillation method to pre - train the proposed model, and constructs a simple feature extraction strategy to transfer the knowledge of SAM to SViT. Experimental results show that YumiNet of the present invention performs better than its peers in various high - resolution medical image segmentation tasks, and requires less computational resources at the same time. When SViT is combined with the SAM series, they show excellent zero - point generalization ability, which can provide high - resolution segmentation with high memory efficiency and superior performance in various medical imaging modalities. The goal of the present invention is to train a model: , where represents the learning parameters, is the given image, is the predicted segmentation mask. Each pixel in the segmentation mask is assigned a hard label according to a predefined list of categories to achieve the highest similarity with the ground truth. YumiNet in the present invention is based on an encoder - decoder architecture, integrating a shallow - layer ViT image encoder and a patch - level cross - multi - scale decoder. In addition, feature extraction technology is also utilized during the pre - training process to fully exert the potential of YumiNet.

[0020] Continue to refer to Figure 1, Considering that current efficient ViTs mainly use depthwise convolution, group attention, or multi-scale learning modules to increase the network width, this method maintains sufficient feature representation based on large-scale natural image datasets, but does not consider the depth of the model. Due to the high cost of pixel-level annotations, medical datasets usually contain limited fully annotated labels. Therefore, deep models are more likely to cause overfitting in medical datasets. Thus, in the present invention, a Shallow Vision Transformer (SViT) is proposed as the image encoder of YumiNet. It adopts a simple but memory-efficient two-stage layout design, and the key component is a local gated linear attention module, which can improve the efficiency of the network in terms of memory, computation, and parameters.

[0021] Continue to refer to Figure 1 , To build a shallow ViT, the present invention proposes a two-stage layout. First, two MBConvs are used as channel interaction layers to efficiently learn low-level representations. Specifically, each MBConv contains a 1 * 1 convolution for channel expansion, a 3 * 3 depthwise convolution followed by a 1 * 1 projection convolution , for channel communication. Given a set of patch embeddings in , where H and W represent the height and width of the given image respectively, S is the predefined patch size, and C is the embedding channel, the calculation can be expressed as: where, represents an activation function (such as GELU). Based on this, the channel interaction layer in the present invention enhances the model's ability by introducing an inductive bias of local structural information.

[0022] Continue to refer to Figure 1 , Considering that the recently emerged Mamba architecture has become a competitive ViT alternative in long-sequence vision tasks (semantic segmentation) (such as semantic segmentation), inspired by the success of the gating mechanism in Mamba, in the next stage, the present invention designs Dual-Gated Linear Attention (DGLA) to form a token interaction layer for learning global representations. Specifically, the self-attention mechanism is defined as follows: where is a learnable linear projection matrix, where d represents the projection dimension. The self-attention mechanism updates the given image x by calculating the similarity between all query-key pairs, and this operation can be implemented through other implementation methods (such as gating mechanisms). Therefore, in the present invention, DGLA rewrites the formula as: The inter-gate network improves the efficiency of text value representation using non-linear transformation. In addition, in the present invention, an external gating network is also applied to DGLA to distinguish the utility of the given image x. The final embedded image embedding can be calculated as follows: In summary, the proposed SViT adopts a channel interaction layer and a token interaction layer to efficiently learn low-level representations and processes long-range visual dependencies with linear complexity.

[0023] Continue to refer to Figure 1 , considering that the traditional U-shaped architecture uses skip connections and pyramid decoding operations to restore the mask size. However, for high-resolution medical image segmentation, the computational cost of this structure is very high. To solve this problem, the present invention designs a patch-level cross multi-scale decoder (PLCM-Decoder), which adds multi-scale decoding information to the high-dimensional but small-sized image embedding, thus eliminating the need for pyramid decoding. First, a query embedding word q is created in to learn the decoding information, where T is the number of predicted classes and D is the decoding dimension. Then, self-attention is applied to update the masked query: , and then the image embedding is updated with the bidirectional multi-scale cross-attention layer of q'. The above operations can be expressed as: Among them, ψ is the positional encoding, which is used to enhance the dependence between the geometric position and the type; δ(·) represents the multi-scale module, including parallel pooling operators with three different kernel sizes, which is used to enhance the semantic information of different scales without additional parameters. Finally, the generation of the segmentation mask can be defined as follows: In summary, the proposed PLCM-Decoder does not require pyramid decoding, thus reducing the computational cost of high-resolution medical image segmentation tasks.

[0024] In addition, considering that pre-training is crucial for unleashing the full potential of the network in downstream tasks, recent research has shown that directly pre-training small models on large-scale data has limited benefits. To improve the quality of the feature representations generated by the image encoder (which is crucial for effective segmentation decoding), the present invention adopts a direct feature refinement method to pre-train the SViT component of YumiNet. Recently, SAM has demonstrated excellent performance in natural scenes as it is trained on massive data. To leverage this advantage, the image encoder of SAM-base is selected as the distillation target in the present invention. Specifically, the present invention connects a neck layer to the SViT to align the feature dimensions with those of the image encoder of SAM. The neck layer consists of a 1*1 convolution and a 3*3 convolution. The present invention uses the mean squared error (MSE) loss on 1% of the samples from the SA-1B dataset to optimize the distillation process.

[0025] In the original ViT, the expansion ratio of the MLP is generally set to 4, which makes the hidden dimension 4* wider than the embedding dimension C. Although this setting improves the model capacity, the wider dimension consumes a large portion of the memory resources. To alleviate this bottleneck, inspired by recent research, YumiNet sets the expansion ratio to 2 in both MBConvs and MLP. In addition, the width C of the embedding dimension and the depth L of the network are also the main factors affecting the memory cost. Previous research has shown that wider dimensions are beneficial for learning hard sample representations during training, while feature distillation can avoid the drawbacks of hard samples as it does not require the use of ground truth labels. Therefore, YumiNet of the present invention adopts a narrower embedding dimension and a shallower structure (e.g., C = 96, L = 6).

[0026] Participate Figures 2 to 4 , to comprehensively verify the effectiveness of YumiNet proposed in the present invention, it was first evaluated on two-dimensional medical image datasets in five different modalities, such as Figure 2 shown. These datasets include dermoscopy, X-ray, fundus, ultrasound, and microscopy. Secondly, as Figure 3 shown, the three-dimensional medical image segmentation model of the present invention was evaluated on various organs. Finally, as Figure 4As shown, medical video segmentation experiments were conducted using two different colonoscopy datasets. To quantify the performance of 2D and video segmentation, two metrics widely used in medical image segmentation, Dice and mIoU, were selected. In 3D medical image segmentation, evaluation metrics commonly used in existing research were adopted. YumiNet was compared with 13 models based on CNN, ViT, Mamb, and lightweight architectures. All models were re-implemented according to their official settings. The best and second-best performance values are marked in bold and underlined respectively. The evaluation results are as follows: 2D medical image segmentation: As Figure 5 shown, UNext and DCSAU-Net use fewer parameters but are not as good as nnUNet, MADGNet, or EMCAD. In contrast, the YumiNet framework of the present invention outperforms its peers on 2D medical segmentation datasets in five different modes. In particular, compared with UNext, the GPU memory usage of YumiNet is reduced by 92.31%, 62.07%, 95.71%, 82.71%, and 86.35% respectively; compared with UNext, DCSAU-Net, nnUNet, MADGNet, and EMCAD, YumiNet is reduced by 92.31%, 62.07%, 95.71%, 82.71%, and 86.35% respectively.

[0027] 3D medical image segmentation: As Figure 6 shown, the experimental results of the 3D multi-organ segmentation dataset are presented. The experimental results report the Dice scores of individual organ evaluations. The results show that in most organ segmentation tasks, the proposed YumiNet of the present invention outperforms the second-ranked EMCAD. The average Dice and mIoU are increased by 1.01% and 1.12% respectively. Among them, the Hausdorff Distance (HD) of YumiNet is the lowest, only 36.09, indicating stronger organ boundary localization ability.

[0028] Medical video segmentation: As Figure 7 shown, in two polyp segmentation datasets, YumiNet outperforms all state-of-the-art methods. Compared with nnUNet and EMCAD, the Dice scores of YumiNet in CVC-ColonDB and CVC-ClinicDB are significantly increased by 1.36% and 3.36% respectively. Therefore, the proposed YumiNet of the present invention has achieved new SOTA results in various medical segmentation tasks and shows a low computational cost.

[0029] Zero-shot generalization: As Figure 8 shown, the present invention further combines the pre-trained SViT with the prompt encoder and mask decoder of SAM and SAM2.Figure 8 Shows the average zero-point performance of all eight medical datasets in the box prompt mode. It can be seen that compared with lightweight SAM backbones (such as RepViT-SAM, EfficientSAM), the proposed SViT achieves lower parameter costs and complexities, and when used in combination with SAM2, it shows better generalization performance in medical 2D, 3D, and video segmentation tasks.

[0030] In summary, the present invention proposes YumiNet for general high-resolution medical image segmentation. Specifically, SViT is introduced to build a long-range dependence model with linear complexity. Then, a PLCM decoder is designed to eliminate the need for traditional pyramid scale recovery and reduce the memory cost during training. In addition, a pre-trained feature distillation technique is adopted to unleash the potential of SViT. A large number of experiments prove that YumiNet outperforms its peers in various medical segmentation tasks and has a lower computational cost. When SViT is combined with SAM, the combined architecture exhibits excellent zero-point generalization ability.

[0031] In the description of the present invention, the description referring to terms such as "one embodiment", "some embodiments", "in this embodiment", "specific examples", or "some examples" means that the specific features, mechanisms, materials, or characteristics described in connection with the embodiment or example are included in at least one embodiment or example of the present invention. In this specification, the schematic expressions of the above terms do not necessarily refer to the same embodiment or example. Moreover, the specific features, mechanisms, materials, or characteristics described can be combined in a suitable manner in any one or more embodiments or examples. In addition, without contradiction, those skilled in the art can combine and combine the different embodiments or examples described in this specification and the features of different embodiments or examples.

[0032] The above is only the specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any changes or substitutions that can be easily thought of by those skilled in the art within the technical scope disclosed by the present invention should be covered by the protection scope of the present invention. Therefore, the protection scope of the present invention should be subject to the protection scope of the claims.

Claims

1. A lightweight general-purpose high-definition medical image segmentation network, characterized in that: The lightweight general high-definition medical image segmentation network is based on an encoder-decoder architecture, integrating a shallow ViT image encoder and a patch-level cross-multi-scale decoder, wherein the shallow ViT image encoder uses a channel interaction layer and a label interaction layer to learn the low-level features of a given image, and processes long-distance dependencies with linear complexity to obtain an embedded image, and the patch-level cross-multi-scale decoder learns decoding information by creating query embeddings, uses a self-attention mechanism to update mask queries, and updates the embedded image through a bidirectional multi-scale cross-attention layer to obtain a segmentation mask.

2. The lightweight universal high-definition medical image segmentation network according to claim 1, characterized in that: The shallow ViT image encoder uses two MBConvs as the channel interaction layer to learn the low-level features of the given image, each of which includes a 1 * 1 convolution connected in sequence. , a 3 * 3 depth convolution And a 1*1 projection convolution Patch embedding is performed on the given image to perform linear complexity processing to obtain the embedded image.

3. The lightweight universal high-definition medical image segmentation network according to claim 2, characterized in that: The expression of the patch embedding of the given image by MBConvs is as follows: in, representing the given image; H represents the height of the given image; W represents the width of the given image; S represents the preconfigured patch size; C represents the embedded channel; Represents the activation function.

4. The lightweight universal high-definition medical image segmentation network according to claim 1, characterized in that: The shallow ViT image encoder adopts a dual-gated linear attention mechanism to form the labeled interaction layer for learning global features, calculates the similarities between all query-key pairs through the dual-gated linear attention mechanism to update the given image, and processes the given image using a multi-layer perceptron and layer normalization method to obtain the embedded image node.

5. The lightweight universal high-definition medical image segmentation network according to claim 4, characterized in that: The expression of the dual-gated linear attention mechanism is as follows: in, represents the matrix V OK; H represents the height of the given image; W represents the width of the given image; S represents the preconfigured patch size; represents the kernel function; representing the given image; , , represents a learnable linear projection matrix; C represents the embedding channel; d represents the projection dimension; Represents the SiLU function.

6. The lightweight universal high-definition medical image segmentation network according to claim 5, characterized in that: The calculation formula for obtaining the embedded image by processing the given image using the multilayer perceptron and the layer normalization method is as follows: in, representing the embedded image; Indicates that the multilayer perceptron is used for processing; Indicates that the layer normalization method is used for processing; represents the SiLU function; representing the given image; represents a learnable linear projection matrix; C represents the embedding channel; d represents the projection dimension.

7. The lightweight universal high-definition medical image segmentation network according to claim 1, characterized in that: The expression of the patch-level cross-multi-scale decoder update mask query is as follows: in, represents the bidirectional multi-scale cross attention layer; represents the self-attention mechanism; represents the query embedding word; Indicates the number of predicted categories; Indicates the decoding dimension.

8. The lightweight universal high-definition medical image segmentation network according to claim 7, characterized in that: The expression for updating the embedded image across the multi-scale decoder at the patch level is as follows: in, representing the embedded image; Indicates positional encoding; d represents the projection dimension; Represents a multi-scale module, which is used to enhance semantic information at different scales; representing the given image; represents a learnable linear projection matrix.

9. The lightweight universal high-definition medical image segmentation network according to claim 8, characterized in that: The expression of the segmentation mask obtained by processing the patch-level cross-multi-scale decoder is as follows: in, represents the segmentation mask; Indicates performing two transposed convolutions; Represents a bilinear interpolation function that directly restores the mask size; Indicates that the multilayer perceptron is used for processing; H represents the height of the given image; W represents the width of the given image; Indicates the number of predicted categories.

Images