SSRepVM-UNet-based lightweight medical image segmentation method

Abstract

The invention provides a lightweight medical image segmentation method based on SSRepVM-UNet. The method comprises three steps of image processing, model training and image segmentation. The method comprises the following steps: firstly, dividing a downloaded data set, and then performing size reduction to obtain a uniform size (256 * 256) and generating a corresponding. Npy file; after data enhancement is carried out on the training set, the training set is sent to an SSRepVM-UNet model for training, and after 250 rounds of training, the model parameter with the minimum loss of the verification set is stored as best.pt; and finally, testing the test set by utilizing best.pt to obtain each index result, and storing a reasoning result, so that the reasoning result can be conveniently compared with a mask, and the segmentation effect can be intuitively felt.

Description

Technical Field

[0001] The present invention relates to the field of medical image technology, and in particular to a lightweight medical image segmentation method based on SSRepVM-UNet. Background Art

[0002] With the rapid development of deep learning, automatic medical image segmentation has begun to be widely used in the medical field. Since the efficiency and accuracy of its segmentation results can help doctors make better and faster pathological judgments and promote the improvement of medical efficiency, automatic medical image segmentation technology has strong practical significance and has also attracted widespread attention from scientific researchers. At present, automatic medical image segmentation is mainly based on deep learning networks represented by CNN and Transformer.

[0003] UNet is one of the main representatives based on CNN models. Because of its strong scalability and simple structure, it can extract multi-scale information through different encoders and decoders, and has gained a prominent position in the field of medical image segmentation. The convolutional neural network model has excellent local feature extraction capabilities, but it is limited by the local receptive field and it is difficult to establish the correlation of remote information. Even if a large convolution kernel is used to obtain remote information, the segmentation results are still not satisfactory. Therefore, the U-shaped architecture of Transformer appeared, using the self-attention mechanism to solve the problem of remote information extraction. This modification and innovation has achieved good segmentation results. However, the computational quadratic complexity of the self-attention mechanism grows quadratically with the size of the input image. Therefore, in the segmentation of high-resolution medical images that require dense predictions, Transformer will introduce a lot of calculations, occupy a huge amount of memory, and is not suitable for deployment on resource-constrained medical equipment.

[0004] After Vmamba was proposed in January 2024, VM-UNet first proposed a U-shaped architecture based on Vmamba and applied it to medical image segmentation. H-vmunet proposed the H-SS2D high-order visual state space based on VM-UNet, and directly added the shared weight SAB and CAB modules to the jump connection between the encoder and decoder to strengthen the fusion of features of different sizes. Compared with VM-Unet, it reduces the number of parameters by 67.28%, and has been significantly improved on the ISIC2017, ISIC2018 and PH2 datasets. Its parameter volume and computational complexity are 8.97M and 0.74GFLOPs respectively. However, the segmentation effect, model size and computational complexity of H-vmunet on the ISIC2018 dataset still have room for optimization and improvement. UltraLight VM-UNet is extremely lightweight based on H-vmunet, with a model size and computational complexity of only 0.05M and 0.06GFLOPs. UltraLight VM-UNet uses only one 3×3 convolution for feature extraction in the first three modules of the encoder and decoder, and uses global maximum pooling and bilinear interpolation for downsampling and upsampling. The weight-sharing SAB and CAB modules are also used in the jump connection between the encoder and decoder. In the 6 PVMLayer, the input of Mamba is split into 4 parts for separate calculations, which greatly reduces the model size and the number of parameters. Good results have also been achieved on the above three datasets, but for the sake of extreme lightweight, the segmentation effect on ISIC2018 is not much improved compared to the previous model. In summary, H-vmunet and UltraLightVM-UNet can be appropriately improved in terms of balancing segmentation effect and model size.

[0005] In the automated diagnosis process for dermatoscopes, image segmentation is of great significance. Accurate lesion boundaries can provide an important basis for the subsequent analysis of lesion types. However, since the lesion area often has blurred boundaries, complex and diverse colors, and severe hair occlusion, automatic image segmentation of dermatoscopes has become a very challenging task. The same problem occurs in the automated diagnosis process of gastrointestinal polyps. Such as the variability of polyp morphology, the fuzziness of boundaries, the similarity of color to surrounding tissues, the noise problem of endoscopic images, and the reflection interference of intestinal fluid and light. These factors together increase the difficulty of the model to accurately segment polyps in actual clinical applications. Moreover, in actual application, the model needs to be deployed on medical testing equipment with limited computing resources, which puts higher requirements on model size and computational complexity. The image segmentation model based on CNN has difficulty in extracting remote information, the image segmentation model based on Transformer has high complexity and occupies a lot of memory, and the H-vmunet based on vmamba and the UltraLight VM-UNet based on mamba are not balanced in model size, computational complexity and segmentation accuracy. In summary, the existing models have obvious advantages and disadvantages, and lack a well-balanced model in all aspects. This patent proposes SSRepVM-UNet based on UltraLight VM-UNet, which improves the segmentation effect of the model on three dermoscopic datasets while maintaining lightweight and linear complexity. In particular, for the ISIC2018 dataset, the DSC of UNet, H-vmunet and UltraLight VM-UNet does not exceed 0.895. However, polyp image segmentation models, such as VM-UNetV2, have the problem of large number of parameters and complex calculation. Summary of the invention

[0006] In view of this, the purpose of the present invention is to provide a lightweight medical image segmentation method based on SSRepVM-UNet, which is still very advantageous in terms of lightweight. In addition, its training time is short and no pre-training parameters need to be imported, which is very beneficial to the subsequent promotion of various medical images.

[0007] To achieve the above objectives, the present invention adopts the following technical scheme: a lightweight medical image segmentation method based on SSRepVM-UNet, including three steps: image processing, model training and image segmentation; first, the downloaded data set is divided, then resized to a uniform size and the corresponding .npy file is generated; then the training set is data enhanced and sent to the SSRepVM-UNet model for training, and after 250 rounds of training, the model parameters with the smallest loss of the validation set are saved as best.pt; finally, the test set is tested using best.pt to obtain the results of each indicator, and the inference results are saved and compared with the mask to feel the segmentation effect.

[0008] In a preferred embodiment, the SSRepVM-UNet model uses 3×3 convolution as the stem of the model, increases the 3-channel image to 8 channels, and the encoder uses maximum pooling for downsampling; in modules 2 and 3 of the encoder and decoder, depth-separable structured re-parameterized convolution is used to extract features from the image; in modules 4, 5 and 6 of the encoder and decoder, parallel Vision-Mamba structured re-parameterized convolution modules are used; in the jump connection of the encoder and decoder, the weight-sharing SAB and CAB proposed in MALUNet are used; bilinear interpolation is used for upsampling in the decoder; the last upsampling uses deconvolution; the channel parameters of each module in the whole are set to [8, 16, 24, 32, 48, 64]; the overall parameter amount and computational amount of SSRepVM-Unet are 0.17M and 0.18GFLOPs.

[0009] In a preferred embodiment, the structural re-parameters consist of a 3×3 convolution, a 1×1 convolution and an Identity; the 1×1 convolution and the Identity are converted into a 3×3 convolution, and similarly, the three BN layers are combined together; during inference, the multi-branch structure is converted into a single-path structure, and the convolution layer and the BN layer are fused;

[0010] The formulas for the convolutional layer with bias and the BN layer are:

[0011]

[0012] Substituting the result of the convolutional layer into the BN formula, we get:

[0013]

[0014] make Get the final fusion expression:

[0015] BN(Conv(x))=W fused +B fused

[0016] The 3×3 convolution is optimized in the computing library CuDNN, Inter MKL and hardware. During inference, the three branches in the structural parameterization are equivalent to one 3×3 convolution. The fusion of the convolution layer and the BN layer ultimately enables the module to have a high-performance inference rate.

[0017] In a preferred embodiment, the RepConv module in the SSRepVM-UNet model first changes the number of channels using a 3×3 convolution, and adds a RepViT block after the convolution and before the downsampling to improve the block structure of MoblieNetV3, advances the deep convolution to separate the token mixer and the channel mixer, and introduces structural re-parameters during training to introduce a multi-branch topology structure for the deep filter; a 3×3 and a 1×1 depth-separable convolution are used in the token mixer; the channel mixer consists of a 1×1 extended convolution layer and a 1×1 projection layer, uses GELU as the activation function, and uses a residual structure to help information flow in the deep network; and an SE attention module is added between the token mixer and the channel mixer.

[0018] In a preferred embodiment, the SAB in the SSRepVM-UNet model operates on multi-stage features through shared two-dimensional convolution to generate a feature map on the spatial axis; SAB performs average pooling and maximum pooling operations on the channel dimension for the feature map of each stage, and concatenates them to obtain a feature map with two channels of unchanged size; uses a shared dilated convolution operation, and then uses a Sigmoid activation function to generate a spatial attention map for each stage; finally, the generated spatial attention map is element-wise multiplied by the original feature map, and a residual connection is added to obtain the final output.

[0019] In a preferred embodiment, the CAB in the SSRepVM-UNet model consists of a one-dimensional convolution and a fully connected layer, which globally and locally fuses multi-stage features to generate a feature map on the channel axis; the acquisition of multi-stage and multi-scale information plays a vital role in the segmentation of targets of different scales; CAB connects the features of different stages on the channel axis in series to generate a channel attention map, thereby better integrating information.

[0020] In a preferred embodiment, the SS-RepVM module first splits the channels of the input feature map in half, one half is used for convolution operation, and the other half is used for Mamba to capture long-range dependencies; for the convolution operation, the separated half of the channels is first regularized, and then the RepViT module is selected as the convolution feature extraction layer; for the other half of the channels, the 2D selective linear scanning SS2D proposed by Vmamba is used for feature extraction.

[0021] In a preferred embodiment, SS2D consists of a scan expansion operation, an S6 block, and a scan merge operation; the scan expansion operation converts the input image into four sequences from four different directions: upper left to lower right, from lower right to upper left, from upper right to lower left, and from lower left to upper right, and then uses the S6 block to extract features from these four sequences to ensure that feature information in all directions can be scanned, thereby obtaining different features; finally, the scan merge operation sums and merges the sequences from the four directions to keep the input and output image sizes consistent; S6 adjusts the parameters of the SSM according to the input features.

[0022] In a preferred embodiment, X2 is first reshaped to make the feature dimension consistent with the SS2D input dimension; after changing the dimension, the feature map is normalized along the channel direction using LN; then the feature map is separated again and input into SS2D respectively, and then the output is residually connected to improve the model's ability to obtain long-range spatial information; then the outputs of the two SS2Ds are concatenated, normalized again on the channel, and reshaped again to restore the dimension to (B, C, H, W); after X1 and X2 are concatenated in the channel dimension, the channel is shuffled using the Channel Shuffle operation; finally, a common 3×3 convolution is used to perform channel mapping on the feature map outputted at the end; the specific formula is:

[0023]

[0024] Out = Conv[shuffle(X out )].

[0025] Compared with the prior art, the present invention has the following beneficial effects: SSRepVM-UNet is very friendly to training conditions and only requires 672MB of video memory during training. The training speed is also very fast, with an average of 33 seconds to complete a round of training, and the inference time is short, with an average FPS of 3.46. SSRepVM-UNet does not need to import pre-trained weights, and the model size used to verify inference after training is only 868KB. In terms of model accuracy, SSRepVM-UNet has achieved a significant lead. The DSC of 0.9124 in ISIC2018 data far exceeds H-vmunet's 0.8966 and UltraLight VM-UNet 0.8940, and SSRepVM-UNet's SE of 0.9001, SP of 0.9804 and ACC of 0.9633 are also the best among all current lightweight models. Similarly, SSRepVM-UNet also achieved excellent performance on the PH2 and ISIC07 datasets. Although SSRepVM-Net's 0.17M model parameters and 0.18GFLOPs model computation are not as good as the current lightest UltraLightVM-UNet, it is also the second lightest model known so far, and it is better than UltraLight VM-UNet in a total of 12 indicators on the three data of ISIC2018, PH2 and ISIC2017. The performance on the Kvasir and ClinicDB datasets is also excellent. Compared with VM-UNetV2, the mIoU index has been improved by 0.0136 and 0.0048, respectively, and the DSC index has been improved by 0.085 and 0.027, respectively. And it has achieved the lightest model in the polyp image segmentation model. Therefore, in terms of model parameters, model computation and image segmentation accuracy, SSRepVM-UNet is more balanced and more suitable for deployment on edge devices, promoting the development of automatic medical image segmentation technology. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 This is a schematic diagram of the SSRepVM-UNet network structure of a preferred embodiment of the present invention;

[0027] Figure 2 It is a structural parameter principle diagram of a preferred embodiment of the present invention;

[0028] Figure 3 A schematic diagram of a RepConv module in a preferred embodiment of the present invention;

[0029] Figure 4 A schematic diagram of the SAB and CAB modules of a preferred embodiment of the present invention;

[0030] Figure 5A schematic diagram of a SS-RepVM module in a preferred embodiment of the present invention;

[0031] Figure 6 The data set, samples and inference results of the preferred embodiment of the present invention; Figure 6 (a) is a sample randomly selected from the dermatoscopy datasets ISIC2017, ISIC2018, PH2 test set and the polyp datasets ClinicDB and Kvasir test set. Figure 6 (b) makes the corresponding Ground Truth, Figure 6 (c) is the inference result of our model.

[0032] Figure 7 The folder structure for storing training sets and test sets in a preferred embodiment of the present invention;

[0033] Figure 8 It is the log trained by ISIC2018 of the preferred embodiment of the present invention. DETAILED DESCRIPTION

[0034] The present invention will be further described below in conjunction with the accompanying drawings and embodiments.

[0035] It should be noted that the following detailed descriptions are illustrative and are intended to provide further explanation of the present application. Unless otherwise specified, all technical and scientific terms used herein have the same meanings as those commonly understood by those skilled in the art to which the present application belongs.

[0036] It should be noted that the terms used herein are only for describing specific embodiments and are not intended to limit the exemplary embodiments according to the present application; as used herein, unless the context clearly indicates otherwise, the singular form is also intended to include the plural form. In addition, it should be understood that when the terms "comprise" and / or "include" are used in this specification, they indicate the presence of features, steps, operations, devices, components and / or their combinations.

[0037] A lightweight medical image segmentation method based on SSRepVM-UNet, including three steps: image processing, model training, and image segmentation. First, the downloaded data set is divided, then resized to a uniform size (256×256) and the corresponding .npy file is generated. Then the training set is data enhanced and sent to the SSRepVM-UNet model for training. After 250 rounds of training, the model parameters with the smallest loss of the validation set are saved as best.pt. Finally, the test set is tested using best.pt to obtain the results of various indicators, and the inference results are saved for easy comparison with the mask to intuitively feel the segmentation effect.

[0038] The SSRepVM-UNet model is described as follows:

[0039] The overall network architecture of SSRepVM-UNet is as follows Figure 1 As shown in the figure, the overall model still adopts a U-shaped structure, which is basically symmetrical from top to bottom. First, a common 3×3 convolution is used as the stem of the model to increase the 3-channel image to 8 channels for subsequent calculations. In the entire encoder, maximum pooling is used for downsampling, and the minimum number of parameters and calculations is sought without losing too many image features. In modules 2 and 3 of the encoder and decoder, a deep separable structure with heavy parameters is used for image feature extraction. The RepConv module structure is shown in the figure.

[0040] 3. Modules 4, 5, and 6 of the encoder and decoder all use parallel Vision-Mamba structured re-parameterized convolution modules, and their structures are as follows Figure 5 As shown in the figure. In the jump connection between the encoder and decoder, the weight-sharing SAB and CAB proposed in MALUNet are used. Similarly, the upsampling of the decoder uses a relatively simple and effective bilinear interpolation for upsampling. Unlike other models, the last upsampling uses deconvolution, which increases the amount of calculation and parameters compared to bilinear interpolation, but the learnable deconvolution brings better segmentation and generalization effects to the model. The channel parameters of each module in the overall are set to [8,16,

[0041] 24,32,48,64]. The total number of parameters and computational complexity of SSRepVM-Unet is 0.17M and 0.18GFLOPs, which is very lightweight compared to models such as UNet, VM-UNetv2, and H-vmunet. In addition, its training time is short and it does not require the import of pre-trained parameters, which is very beneficial for the subsequent promotion of various medical images.

[0042] 2. Structural re-parameter principle

[0043] Inception first introduced a multi-branch structure and achieved high performance, but the multi-way structure needs to save intermediate results, which will significantly increase the memory usage. Most importantly, it will also slow down the reasoning process, which is very unfavorable when computing resources are limited. Therefore, researchers proposed a structural re-parameterization, the principle of which is as follows: Figure 3 As shown in Figure 1. It consists of a 3×3 convolution, a 1×1 convolution and an Identity. Figure 2As shown in the figure, 1×1 convolution and Identity can be converted into 3×3 convolution, because the inputs of the three convolutions are the same, and the three convolutions can be combined together according to the associative law. Similarly, the three BN layers can also be combined together. In this way, the multi-branch structure is transformed into a single-path structure during reasoning, which speeds up the reasoning speed. In the code implementation, the convolution layer and the BN layer are also fused.

[0044] The formulas for the convolutional layer with bias and the BN layer are: Conv(x) = W(x) + b,

[0045] Substituting the result of the convolutional layer into the BN formula, we can get:

[0046]

[0047] make The final fusion expression can be obtained:

[0048] BN(Conv(x))=E fused +B fused

[0049] The 3×3 convolution is optimized in the computing library CuDNN, Inter MKL and hardware. During inference, the three branches in the structural parameterization can be equivalent to one 3×3 convolution. Combined with the fusion of the convolution layer and the BN layer, the module finally has a high-performance inference rate.

[0050] 3. RepConv module

[0051] The RepConv module structure is as follows Figure 3 As shown in the figure, the number of channels is changed by using 3×3 convolution. In order to increase the network depth and reduce the information loss caused by downsampling, a RepViT block is added after the convolution and before the downsampling. According to the research content of MetaFormer, the specific form of the token mixer is not very important. Under this concept, the RepViT block improves the block structure of MoblieNetV3, advances the deep convolution, and separates the token mixer and channel mixer. In order to improve performance, the structural re-parameter is introduced during training to introduce the deep filter Figure 3The multi-branch topology of the network is used to reduce the number of parameters and the amount of computation to achieve lightweight model. A 3×3 and a 1×1 depth-separable convolution are used in the token mixer, which directly improves the inference speed of the model. The channel mixer consists of a 1×1 extended convolution layer and a 1×1 projection layer. GELU is used as the activation function. The residual structure is used to help information flow in the deep network, alleviate the gradient vanishing problem, and improve network performance. An SE attention module is added between the token mixer and the channel mixer to enhance the model's ability to capture long-distance dependencies and local information. Four RepConv modules are added, which increases GFLOPs by 0.023. DSC can be improved by 0.1% on ISIC2017, which is still a very significant effect. SE attention can be optimized in the future to further reduce the amount of computation of the RepConv module.

[0052] 4.SAB and CAB modules

[0053] Depend on Figure 4 As shown in the figure, SAB operates on multi-stage features through shared two-dimensional convolution to generate feature maps on the spatial axis. SAB performs average pooling and maximum pooling operations on the channel dimension for the feature maps of each stage, and concatenates them to obtain a feature map with two channels of unchanged size. The shared dilated convolution operation (dilation rate is 3) is used, and then the Sigmoid activation function is used to generate the spatial attention map of each stage. Finally, the generated spatial attention map is element-wise multiplied by the original feature map, and a residual connection is added to obtain the final output. CAB consists of a one-dimensional convolution and a fully connected layer, which globally and locally fuses multi-stage features to generate a feature map on the channel axis. The acquisition of multi-stage and multi-scale information plays a vital role in the segmentation of objects of different scales. CAB concatenates the features of different stages on the channel axis to generate a channel attention map, thereby better integrating information.

[0054] 4.SS-RepVM module

[0055] UltraLightVM-UNet research and analysis shows that the number of input channels has an explosive impact on Mamba's parameters, so this method still uses the method of separated channel parallel computing to reduce the model size. Mamba, as a state-space model, can effectively simulate long-range interactions while maintaining linear complexity. The convolution layer has excellent local feature extraction capabilities. In order to combine the advantages of both, we first split the channels of the input feature map in half, one half for convolution operations and the other half for Mamba to capture long-range dependencies. For the convolution operation, first regularize the separated half of the channel, and then use Figure 3 The RepViT module shown is used as a convolutional feature extraction layer. In order to reduce parameters, Figure 3 The RepViT module is slightly different in that Figure 5 The RepVIT module in is between the token mixer and the channel mixer, and SE attention is not used. For the other half of the channels, we use the 2D Selective Linear Sweep (SS2D) proposed by Vmamba for feature extraction.

[0056] SS2D consists of three parts: scan expansion operation, S6 block and scan merging operation. The scan expansion operation converts the input image into four sequences from four different directions: upper left to lower right, from lower right to upper left, from upper right to lower left, and from lower left to upper right. Then the S6 block is used to extract features from these four sequences to ensure that the feature information in each direction can be scanned, thereby obtaining different features. Finally, the scan merging operation sums and merges the sequences from the four directions to keep the input and output image sizes consistent. S6 adjusts the parameters of SSM according to the input features, which enables the model to filter out irrelevant information and retain useful information. Since the input dimension of SS2D is (B, W, H, C), but the separated X2 dimension is (B, C, W, H), the Reshape operation is first performed on X2 to make the feature dimension consistent with the SS2D input dimension. After changing the dimension, the feature map is normalized along the channel direction using LN. The feature map is then separated again and input into SS2D respectively, and then the output is residually connected to improve the model's ability to obtain long-range spatial information. Then, the outputs of the two SS2Ds are concatenated and normalized again on the channel. In order to concatenate with the X1 feature map after convolution feature extraction, the reshape operation is performed again to restore the dimension to (B, C, H, W). After X1 and X2 are concatenated in the channel dimension, in order to improve the interaction between the two parts of the features on the channel, the Channel Shuffle operation is used to shuffle the channels. Finally, a normal 3×3 convolution is used to perform channel mapping on the final output feature map. The specific formula is:

[0057]

[0058] Out = Conc[shuffle(X out )]

[0059] Compared with the PVM module of UltraLight VM-UNet, the SS-RepVM module adds a total of 0.03GFLOPs, but by using a hybrid architecture of convolutional layers and SSM, it fully utilizes the advantages of both and greatly improves the modeling of fine-grained feature images. In the ISIC2018 test set, DSC increased by 0.87%, which is a significant effect. There are still many studies on the scanning extension direction of SS2D to enhance information extraction in various directions, and then SS2D can be directly replaced by other scanning modules. This also reflects the scalability of the SS-RepVM module and the rationality of the overall architecture.

[0060] The computer experimental environment conditions used in this patent are as follows: Ubuntu: 20.04, CPU: Inter(R) Xeon(R) Gold 6326CPU@2.90HGz, GPU: GeForce GTX 4090, Memory: 128G, Python: 3.8.10, Pytorch: 1.10.0

[0061] 1) Dataset creation

[0062] ISIC2017 and ISIC2018 are large-scale dermoscopic image datasets released by the International Skin Imaging Collaboration (ISIC). ISIC2017 contains 2000 training images, and all training images have corresponding ground truths. This method divides 1250 images into training sets, 150 images into validation sets, and 600 images into test sets, and the corresponding ground truths are also divided in this ratio. ISIC2018 contains 2594 training images, and similarly divides 1815 images into training sets, 259 images into validation sets, and 520 images into test sets, and the corresponding ground truths are also divided in this ratio. After the dataset is divided, the uniform size is resized to 256×256, and the corresponding .npy files are generated. The PH2 dataset is a dermoscopic image dataset obtained by Pedro Hispano in Matosinhos, Portugal. This dataset contains 200 images and corresponding ground truths. This method resizes them to 256×256 and then divides them all into test sets for external validation. Kvasir and ClinicDB are two public gastrointestinal polyp datasets, both of which contain segmentation masks manually annotated by clinicians. The Kvasir dataset contains 1,000 high-resolution color polyp images, and the ClinicDB dataset contains 612 images. We randomly selected 900 images from Kvasir and 550 images from ClinicDB as training sets, and the remaining images as test sets. All images are uniformly resized to 256×256 after division, and the corresponding .npy files are generated. Reference Figure 6 , Figure 6 (a) are randomly selected samples from the dermatoscopy datasets ISIC2017, ISIC2018, PH2 test set and the polyp datasets ClinicDB and Kvasir test set. Figure 6 (b) Make the corresponding Ground Truth. Figure 6 (c) is the inference result of our model.

[0063] ISIC dataset download address: https: / / challenge.isic-archive.com / data /

[0064] PH2 dataset download address:

[0065] https: / / www.dropbox.com / s / k88qukc20ljnbuo / PH2Dataset.rar /

[0066] Kvasir dataset download address: https: / / opendatalab.com / OpenDataLab / Kvasir

[0067] ClinicDB dataset download address: https: / / polyp.grand-challenge.org / CVCClinicDB /

[0068] In this method, all training images are uniformly regularized before being sent to model training, and data enhancements such as random cropping and random flipping are performed and corresponding labels are generated to improve the robustness of the model. The folder structure for storing the training set and test set is as follows: Figure 7 As shown:

[0069] 2) Definition of training hyperparameters

[0070] In this patent, in order to compare fairly with other segmentation methods, the settings of hyperparameters are basically consistent, and the main hyperparameters are as follows: Initial learning rate: 0.001, BatchSize: 8, epoch: 250, threshold: 0.5, loss function: BceDiceloss, optimizer: Adamw, image size: 256

[0071] 3) Model training

[0072] Before training, set the address of the training data in configuration-setting.py, and then directly run train.py to start training. This method does not load pre-trained weights and starts training from scratch. During training, the model parameters are saved according to the validation loss of each round. Only 672MB of video memory is required for training, and Figure 8 From the ISIC2018 training log, we can see that it only takes 33 seconds to train one round and calculate the validation loss, and the inference time for the last 520 test images is 150 seconds, so the FPS of SSRepVM-UNet is: 3.46. Figure 8 It can be calculated that 250 rounds of training only take 2.3 hours to complete. After the training, the best.pt saved is only 868KB, which is a great advantage for subsequent edge deployment.

[0073] 4). Model testing

[0074] This method uses the commonly used measurement indicators DSC, specificity, sensitivity and accuracy in dermoscopic image segmentation to evaluate the model. The inference test is performed on the three test sets divided in advance, and the values ​​of the above indicators are classified and output. When testing, first set the test set address and the address of best.pt saved after training in config-setting.py, and then directly run test.py to test the model. After the test is completed, the indicator results can be output and the inference results of all test set images will be saved. The test results of the three test sets are compared with the current most advanced models and the model parameters and calculation amount are shown in Tables 1 and 2:

[0075] Table 1

[0076]

[0077]

[0078] Table 2

[0079]

[0080] Comparison of parameters and computational complexity between SSRepVM-UNet and the current major lightweight medical image segmentation models. From the comparison results, we can see that SSRepVM-UNet's 0.171M parameters and 0.18GFLOPs are still very advantageous among all the models. And SSRepVM-UNet also has excellent segmentation effects in ISIC2018, PH2 and ISIC2017.

[0081] This application adds the mIoU metric to evaluate the performance of this method on the polyp dataset. We try to increase the number of channels of the model, expanding the original number of channels [8, 16, 24, 32, 48, 64] to [32, 64, 1288, 256, 512, 1024]. The number of model parameters becomes 16.47M, and the computational complexity is 3.54GFLOPs. The specific test results are shown in Table 3.

[0082] Table 3

[0083]

[0084]

[0085] For the polyp dataset, we tried to increase the number of channels of the model to achieve better performance, which shows that the SSRepVM-UNet model has strong generalization performance. And because the background and foreground of the polyp dataset are similar, our model can still achieve good performance, which shows that the module we designed has better feature extraction capabilities than other models and can adapt to more complex environments.

[0086] As mentioned above, SSRepVM-UNet is not only friendly to training conditions and fast in training speed, but also has both lightweight and accuracy. It is very suitable for deployment on edge devices and application in medical testing, helping doctors to better diagnose diseases and promote the advancement of medical technology.

Claims

1. A lightweight medical image segmentation method based on SSRepVM-UNet, characterized in that: It includes three steps: image processing, model training, and image segmentation. First, the downloaded data set is divided, then resized to a uniform size and the corresponding .npy file is generated. Then, the training set is enhanced and sent to the SSRepVM-UNet model for training. After 250 rounds of training, the model parameters with the smallest loss on the validation set are saved as best.pt. Finally, the test set is tested using best.pt to obtain the results of various indicators, and the inference results are saved and compared with the mask to feel the segmentation effect.

2. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 1, characterized in that: The SSRepVM-UNet model uses 3×3 convolution as the stem of the model, increases the 3-channel image to 8 channels, and the encoder uses maximum pooling for downsampling; in modules 2 and 3 of the encoder and decoder, depth-separable structured re-parameterized convolution is used to extract features from the image; in modules 4, 5, and 6 of the encoder and decoder, parallel Vision-Mamba structured re-parameterized convolution modules are used; in the jump connection of the encoder and decoder, the weight-sharing SAB and CAB proposed in MALUNet are used; bilinear interpolation is used for upsampling in the decoder; the last upsampling uses deconvolution; the channel parameters of each module in the whole are set to [8, 16, 24, 32, 48, 64]; the overall parameter amount and computational amount of SSRepVM-Unet are 0.17M and 0.18GFLOPs.

3. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 2, characterized in that: The structure parameterization consists of a 3×3 convolution, a 1×1 convolution and an Identity. The 1×1 convolution and Identity are converted into 3×3 convolutions. Similarly, three BN layers are combined together. During inference, the multi-branch structure is converted into a single-path structure, and the convolution layer and BN layer are fused. The formulas for the convolutional layer with bias and the BN layer are: Substituting the result of the convolutional layer into the BN formula, we get: make Get the final fusion expression: BN(Conv(x))=W fused +B fused The 3×3 convolution is optimized in the computing library CuDNN, Inter MKL and hardware. During inference, the three branches in the structural parameterization are equivalent to one 3×3 convolution. The fusion of the convolution layer and the BN layer ultimately enables the module to have a high-performance inference rate.

4. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 2, characterized in that: The RepConv module in the SSRepVM-UNet model first uses a 3×3 convolution to change the number of channels, and adds a RepViT block after the convolution and before the downsampling to improve the block structure of MoblieNetV3, advances the deep convolution to separate the token mixer and the channel mixer, and introduces a structural re-parameter during training to introduce a multi-branch topology structure for the deep filter; a 3×3 and a 1×1 depth-separable convolution are used in the token mixer; the channel mixer consists of a 1×1 extended convolution layer and a 1×1 projection layer, uses GELU as the activation function, and uses a residual structure to help information flow in the deep network; an SE attention module is added between the token mixer and the channel mixer.

5. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 2, characterized in that: The SAB in the SSRepVM-UNet model operates on multi-stage features through shared two-dimensional convolution to generate feature maps on the spatial axis; SAB performs average pooling and maximum pooling operations on the channel dimension for the feature maps of each stage, and concatenates them to obtain a feature map with two channels of unchanged size; uses shared dilated convolution operations and then uses the Sigmoid activation function to generate the spatial attention map of each stage; finally, the generated spatial attention map is element-wise multiplied by the original feature map, and a residual connection is added to obtain the final output.

6. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 2, characterized in that: The CAB in the SSRepVM-UNet model consists of a one-dimensional convolution and a fully connected layer, which globally and locally fuses multi-stage features to generate a feature map on the channel axis; the acquisition of multi-stage and multi-scale information plays a vital role in the segmentation of objects of different scales; CAB connects the features of different stages on the channel axis in series to generate a channel attention map, thereby better integrating information.

7. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 2, characterized in that: The SS-RepVM module first splits the channels of the input feature map in half, one half is used for convolution operation, and the other half is used for Mamba to capture long-range dependencies; for the convolution operation, the separated half of the channels is first regularized, and then the RepViT module is selected as the convolution feature extraction layer; for the other half of the channels, the 2D selective linear scan SS2D proposed by Vmamba is used for feature extraction.

8. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 7, characterized in that: SS2D consists of a scan extension operation, an S6 block, and a scan merge operation; the scan extension operation converts the input image into four sequences from four different directions: upper left to lower right, from lower right to upper left, from upper right to lower left, and from lower left to upper right, and then uses the S6 block to extract features from these four sequences to ensure that feature information in all directions can be scanned, thereby obtaining different features; finally, the scan merge operation sums and merges the sequences from the four directions to keep the input and output image sizes consistent; S6 adjusts the parameters of SSM according to the input features.

9. A lightweight medical image segmentation method based on SSRepVM-UNet according to claim 8, characterized in that: First, reshape X2 to make the feature dimension consistent with the SS2D input dimension; after changing the dimension, use LN to normalize the feature map along the channel direction; then separate the feature map again and input it into SS2D respectively, and then perform residual connection on the output to improve the model's ability to obtain long-range spatial information; then concatenate the outputs of the two SS2Ds, normalize them again on the channel, and reshape them again to restore the dimension to (B, C, H, W); after X1 and X2 are concatenated in the channel dimension, use the Channel Shuffle operation to shuffle the channels; Finally, a common 3×3 convolution is used to perform channel mapping on the final output feature map; the specific formula is: Out=Conv[shuffle(X out )]。

Images