A system and registration method for deformable medical image registration model based on Swin Transformer

Abstract

The present invention discloses a deformable medical image registration model based on Swin Transformer, comprising: an encoder, a decoder, and a spatial transformation network, connected in sequence. The encoder comprises: a first downsampling convolutional neural network, a second downsampling convolutional neural network, a convolutional neural network, a first WWA-based Swin Transformer module, a first feature resolution restoration module, a second WWA-based Swin Transformer module, a third downsampling convolutional neural network, a third WWA-based Swin Transformer module, a fourth downsampling convolutional neural network, and a fourth WWA-based Swin Transformer module, connected in sequence. The decoder comprises: first to fifth convolutional neural networks and the second convolutional neural network, connected in sequence. The present invention also provides a deformable medical image registration method based on Swin Transformer.

Description

Technical Field

[0001] The present invention relates to a system and a registration method for a deformable medical image registration model based on Swin Transformer, and belongs to the field of deep learning technology. Background Art

[0002] Medical registration is a fundamental task in medical image analysis. Its purpose is to spatially align anatomical structures in medical images from different patients and at different stages of their course. Registration plays a key role in medical atlas analysis and image-guided radiotherapy.

[0003] Over the past decade, deep learning techniques based on convolutional neural networks have developed rapidly and made significant progress in the field of medical image analysis. Deep learning techniques have been widely used in image-based disease diagnosis, segmentation, and registration.

[0004] With the recent development of Transformers in computer vision tasks, a growing number of studies have demonstrated their ability to model the correlation of spatial information across a wider range, thereby improving the performance of visual models. However, for 3D medical image registration, the sheer volume of information inherent in the data makes Transformers difficult to apply in practice. With the development of the Swin Transformer for vision tasks, Transformers, such as TransMorph, can now be applied to practical medical image registration tasks.

[0005] However, due to the overly fixed window division of Swin Transformer, the model cannot construct anatomical information at uncertain distances; at the same time, due to the large number of parameters and computational complexity of methods such as Transformer, in practical applications, Swin Transformer can only be used at a lower resolution stage, so that Swin Transformer can only model coarse-grained image features.

[0006] In order to accelerate the implementation of artificial intelligence technology in the field of medical image registration and the reliability of registration results, it is urgent to improve the contribution of Swin Transformer in the registration model in practice and build Swin Transformer to model anatomical information at uncertain distances in medical images. Summary of the Invention

[0007] The present invention designs and develops a system of deformable medical image registration models based on Swin Transformer. The Swin Transformer module based on WWA is used to model the global information of medical images. The feature resolution module is used to amplify the role of the Swin Transformer module based on WWA in the model. A 3D deformable medical image registration model is constructed to improve the registration accuracy and generalization ability of modeling anatomical information at uncertain distances in medical images.

[0008] The present invention also designs and develops a deformable medical image registration method based on Swin Transformer, which solves the problem of low registration accuracy in abdominal image registration tasks and effectively improves the accuracy of medical image registration.

[0009] The technical solution provided by the present invention is:

[0010] A system for deformable medical image registration model based on Swin Transformer,

[0011] Includes the following connected in sequence:

[0012] Encoder, decoder, and spatial transformer networks;

[0013] The encoder consists of the following connected in sequence:

[0014] First downsampling convolutional neural network, second downsampling convolutional neural network, first convolutional neural network, first WWA-based Swin Transformer module, first feature resolution restoration module, second WWA-based Swin Transformer module, third downsampling convolutional neural network, third WWA-based Swin Transformer module, fourth downsampling convolutional neural network, fourth WWA-based Swin Transformer module;

[0015] The decoder includes: first to fifth upsampling convolutional neural networks and a second convolutional neural network connected in sequence;

[0016] The output of the first feature resolution restoration module is connected to the input of the second WWA-based SwinTransformer module and the input of the second feature resolution restoration module. The output of the second feature resolution restoration module is connected to the input of the third feature resolution restoration module and the fifth upsampling convolutional neural network. The output of the third feature resolution restoration module is connected to the input of the second convolutional neural network. The output of the second convolutional neural network is connected to the spatial transformation network through the deformation field.

[0017] The output end of the fourth WWA-based Swin Transformer module is connected to the input end of the first upsampling convolutional neural network, the output end of the third WWA-based Swin Transformer module is connected to the input end of the second upsampling convolutional neural network, and the output end of the first WWA-based Swin Transformer module is connected to the first feature resolution restoration module and the third upsampling convolutional neural network.

[0018] Each of the first to fourth WWA-based Swin Transformer modules consists of two consecutive Swin Transformer submodules, and the input feature of the Swin Transformer module is set to z 1 , the calculation formula of two consecutive WWA-based Swin Transformer submodules is as follows:

[0019]

[0020] Where W-MSA is the standard window multi-head attention, SW-MSA is the shifted window multi-head attention, MLP is the multi-layer perception module, LN is the layer normalization operation, For WWA, LN and WMSA and Z 1 The calculated characteristic matrix, Z 2 is calculated by LN and MLP, and Z 1 The obtained feature matrix, For the SW-MSA, LN and WWA and Z 2 The calculated characteristic matrix, Z 3 After MLP and LN and The calculated feature matrix;

[0021] The weighted attention mechanism WWA of window features has two multi-layer perceptron module MLP layers; the first MLP layer calculates the global channel-wise feature interactions across channels; the second MLP layer calculates the feature interactions between windows across windows.

[0022] Preferably,

[0023] The first to fourth downsampling convolutional neural networks contain one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 2; the convolutional neural network contains two convolutional layers with a convolution kernel size of 3×3×3 and a stride of 1; the first to third upsampling convolutional neural networks contain one inverted convolutional layer with a convolution kernel size of 2×2×2 and a stride of 2, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1.

[0024] A deformable medical image registration method based on Swin Transformer, and a system using the deformable medical image registration model based on Swin Transformer, comprising:

[0025] S1: Concatenate the 3D fixed image and floating image of the given depth D, height H, width W, and the first channel into the image pair of the second channel, and input them into the first downsampling convolutional neural network;

[0026] S2: After two consecutive first and second downsampling convolutional neural networks, the first convolutional neural network is used for dimensionality increase and the output is sent to the first WWA-based Swin Transformer module;

[0027] S3: The first WWA-based Swin Transformer module extracts features and outputs them to the first feature resolution recovery module;

[0028] S4: The first to third feature resolution recovery modules all amplify the feature resolution;

[0029] S5: The feature map output by the first feature resolution recovery module is input into the second WWA-based SwinTransformer module, and features are extracted in sequence until the features are finally output to the decoder;

[0030] S6: The first to fifth upsampling convolutional neural networks in the decoder sequentially upsample the input features;

[0031] S7: The decoder connects the corresponding features in the encoder and decoder through jump connections, outputs the D×H×W×3 deformation field to the spatial transformation network, and the image to be deformed and the deformation field are input into the spatial transformation network, which outputs the deformed image.

[0032] Preferably, the model as a whole outputs a deformation field of the third channel, and uses a spatial transformation network to apply the output deformation field to the input image to output a deformed image.

[0033] Preferably,

[0034] Given image features of depth D, height H, width W, and channel C, the feature resolution module expands the features in the channel dimension. After expansion, the feature size is 8C×D×H×W. The expanded features are then reorganized into C×2D×2H×2W image features to restore the feature resolution.

[0035] The beneficial effects of the present invention are:

[0036] 1. The encoder of the present invention proposes a feature resolution recovery module, which is used to amplify the feature resolution of the Swin Transformer module based on WWA, solving the problem of limited performance of the Swin Transformer module in the registration model.

[0037] 2. In the WWA-based Swin Transformer module of the present invention, the weighted window attention mechanism can build associations between window features on a global scale, solving the problem of inflexible and limited information interaction of the Swin Transformer on a global scale, and realizing the ability to globally model image features.

[0038] 3. The Swin Transformer-based deformable medical image registration model RFR-WWANet of the present invention uses the WWA-based Swin Transformer module to model the global information of medical images, and uses the feature resolution module to amplify the role of the WWA-based Swin Transformer module in the model, constructing a 3D deformable medical image registration model similar to Unet. Through a large number of experiments, compared with the current excellent 3D deformable medical image registration methods, the registration model RFR-WWANet has superior registration accuracy and good generalization ability.

[0039] The encoder RFRNet restores the resolution of the feature maps output by the Swin Transformer in the registration model encoder, enabling the Swin Transformer module to provide fine-grained spatial features. Simultaneously, the proposed weighted window attention mechanism (WWA) module constructs an attention mechanism for information interaction of global window features, giving the WWA-based Swin Transformer module the internal force to conduct global information interaction within the feature space, thereby achieving global modeling capabilities. This method is used in the study of general medical image registration tasks, solving the problem of low registration accuracy in abdominal image registration tasks and effectively improving the accuracy of medical image registration. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] Figure 1 This is a flowchart of the deformable medical image registration method based on Swin Transformer described in the present invention.

[0041] Figure 2 Flowchart of an embodiment of a robust window attention mechanism in the present invention. DETAILED DESCRIPTION

[0042] The present invention will be described in further detail below in conjunction with the accompanying drawings so that those skilled in the art can implement the invention with reference to the description.

[0043] like Figure 1-2 As shown, the present invention provides a system for a deformable medical image registration model based on Swin Transformer, comprising:

[0044] The deformable medical image registration model based on Swin Transformer is used to stitch a 3D fixed image and a floating image of a given depth D, height H, width W, and the first channel into an image pair of the second channel and then input it into the model to predict the deformed image of the corresponding size D×H×W; it includes an encoder, a decoder and a spatial transformation network, and the encoder and decoder fuse shallow features with deep features through jump connections; the encoder is an RFRNet containing a feature resolution recovery module, and RFRNet includes a convolutional neural network and a Swin Transformer module based on WWA connected in sequence; the convolutional neural network is used to model image features at the same resolution stage; the downsampling convolutional neural network is used to reduce the input image features to the next level of resolution; the convolutional neural network is used to model image features at the same resolution stage; the Swin Transformer module based on WWA is used to extract image features in a global range; the feature resolution recovery module is used to restore the Swin based on WWA The features output by the Transformer module are restored to resolution; the decoder is composed of sequentially connected upsampling convolutional neural networks, which are connected to the encoder output features through jump connections; the image features output by the decoder are modeled using a convolutional neural network, and the deformation field of the third channel is output to the spatial transformer network. The spatial transformer network is used to apply the deformation field to the input floating image, and the final deformed image is output.

[0045] There are two neural networks for same-layer feature extraction, each containing two convolutional layers with a kernel size of 3×3×3 and a stride of 1; there are four downsampling convolutional neural networks, each containing a convolutional layer with a kernel size of 3×3×3 and a stride of 1, and one convolutional layer with a kernel size of 3×3×3 and a stride of 2; there are five upsampling convolutional neural networks, each containing one inverted convolutional layer with a kernel size of 2×2×2 and a stride of 2, and one convolutional layer with a kernel size of 3×3×3 and a stride of 1.

[0046] The encoder consists of the following connected in sequence:

[0047] First downsampling convolutional neural network, second downsampling convolutional neural network, first convolutional neural network, first WWA-based Swin Transformer module, first feature resolution restoration module, second WWA-based Swin Transformer module, third downsampling convolutional neural network, third WWA-based Swin Transformer module, fourth downsampling convolutional neural network, fourth WWA-based Swin Transformer module;

[0048] The decoder includes: first to fifth upsampling convolutional neural networks and a second convolutional neural network connected in sequence;

[0049] The output of the first feature resolution restoration module is simultaneously connected to the input of the second WWA-based SwinTransformer module and the input of the second feature resolution restoration module. The output of the second feature resolution restoration module is simultaneously connected to the input of the third feature resolution restoration module and the fifth upsampling convolutional neural network. The output of the third feature resolution restoration module is connected to the input of the second convolutional neural network. The output of the second convolutional neural network is connected to the spatial transformation network through the deformation field.

[0050] The output end of the fourth WWA-based Swin Transformer module is connected to the input end of the first upsampling convolutional neural network, the output end of the third WWA-based Swin Transformer module is connected to the input end of the second upsampling convolutional neural network, and the output end of the first WWA-based Swin Transformer module is simultaneously connected to the first feature resolution recovery module and the third upsampling convolutional neural network.

[0051] The first to fourth downsampling convolutional neural networks contain one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 2; the convolutional neural network contains two convolutional layers with a convolution kernel size of 3×3×3 and a stride of 1; the first to third upsampling convolutional neural networks contain one inverted convolutional layer with a convolution kernel size of 2×2×2 and a stride of 2, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1.

[0052] Given image features of depth D, height H, width W, and channel C, the feature resolution module expands the features in the channel dimension. After expansion, the feature size is 8C×D×H×W. The expanded features are then reorganized into image features of C×2D×2H×2W to restore the feature resolution.

[0053] Different from the general registration model using encoder and decoder, the present invention uses the proposed feature resolution recovery module and WWA-based Swin Transformer to construct the encoder.

[0054] This framework solves the problem of limited performance of Swin Transformer in the registration model, and uses WWA to enable Swin Transformer to have the ability to model global information, effectively improving the accuracy of deformable registration of 3D medical images.

[0055] The present invention also provides a deformable medical image registration method based on Swin Transformer. A system using the deformable medical image registration model based on Swin Transformer includes:

[0056] S1: Concatenate the 3D fixed image and floating image of the given depth D, height H, width W, and the first channel into the image pair of the second channel, and input them into the first downsampling convolutional neural network;

[0057] S2: After two consecutive first and second downsampling convolutional neural networks, the first convolutional neural network is used for dimensionality increase and output to the first WWA-based Swin Transformer module

[0058] S3: The first WWA-based Swin Transformer module extracts features and outputs them to the first feature resolution recovery module;

[0059] S4: The first to third feature resolution recovery modules all amplify the feature resolution;

[0060] S5: The feature map output by the first feature resolution recovery module is input into the second WWA-based SwinTransformer module, and features are extracted in sequence until the features are finally output to the decoder;

[0061] S6: The first to fifth upsampling convolutional neural networks in the decoder sequentially upsample the input features;

[0062] S7: The decoder connects the corresponding features in the encoder and decoder through jump connections, outputs the D×H×W×3 deformation field to the spatial transformation network, and the image to be deformed and the deformation field are input into the spatial transformation network, which outputs the deformed image.

[0063] like Figure 1 As shown, the system of the deformable 3D medical image registration model based on feature resolution restoration and weighted window attention mechanism includes the following steps:

[0064] A two-channel image pair consisting of a fixed image and a floating image is input into the model. Two consecutive downsampling convolutional neural networks are used to reduce the image resolution and extract the high-dimensional features of the image. The high-dimensional features of the image are further modeled using a convolutional neural network. The high-dimensional features of the image are modeled globally using a Swin Transformer module based on WWA.

[0065] like Figure 2 As shown in Figure 2, the calculation of the weighted window attention mechanism WWA includes the following steps: given the image features of depth D, height H, width W, and channel C, after the window feature division of the Swin Transformer, the image features become a sequence of window features of shape N×K×C, and the calculation formula of N is:

[0066]

[0067] Where d, h, and w are the window sizes in the Swin Transformer set according to requirements, and K = d × h × w;

[0068] WWA is divided into two stages. "Cross-Channel Attention" calculates the average value in the channel direction, uses MLP to map the mean vector, and uses Sigmoid to normalize the output vector to [0, 1]. After expansion, the vector is dot-multiplied with the input window feature sequence to obtain the channel-weighted window feature sequence.

[0069] "Cross-Window Attention" calculates the average value of the window feature sequence obtained in the previous step in the window direction, uses MLP to map the average direction, and uses Sigmoid to normalize the output vector to the range [0, 1]. After expansion, this vector is dot-multiplied with the originally input window feature sequence to obtain the window feature sequence after weighting in the window direction.

[0070] The calculation method of the WWA-based Swin Transformer module to model the input image features is:

[0071]

[0072] Where W-MSA is the standard window multi-head attention, SW-MSA is the shifted window multi-head attention, MLP is the multi-layer perception module, LN is the layer normalization operation, For WWA, LN and WMSA and Z 1 The calculated characteristic matrix, Z 2 is calculated by LN and MLP, and Z 1 The obtained feature matrix, For the SW-MSA, LN and WWA and Z 2 The calculated characteristic matrix, Z 3 After MLP and LN and The calculated feature matrix.

[0073] After modeling in the first WWA-based Swin Transformer module, the output image features are input into the feature resolution restoration module. The feature resolution restoration module then inputs the restored features into both the next feature resolution restoration module and the next WWA-based Swin Transformer module.

[0074] The image features pass through the downsampling convolutional neural network and the WWA-based Swin Transformer module in sequence until they reach the bottom of the model and are input into the decoder; the upsampling convolutional neural network in the decoder sequentially upsamples the image features and connects them with the features in the encoder.

[0075] The last convolutional neural network models the image features output by the decoder as a deformation field of size D×H×W×3, applies the deformation field to the floating image using a spatial variation network, and outputs the deformed image.

[0076] Although the embodiments of the present invention have been disclosed above, they are not limited to the applications listed in the description and implementation methods. They can be fully applied to various fields suitable for the present invention. For those familiar with the art, additional modifications can be easily implemented. Therefore, without departing from the general concept defined by the claims and the scope of equivalents, the present invention is not limited to the specific details and illustrations shown and described herein.

Claims

1. A system for deformable medical image registration model based on Swin Transformer, characterized in that: Includes the following connected in sequence: Encoder, decoder, and spatial transformer networks; The encoder consists of the following connected in sequence: First downsampling convolutional neural network, second downsampling convolutional neural network, first convolutional neural network, first WWA-based Swin Transformer module, first feature resolution restoration module, second WWA-based Swin Transformer module, third downsampling convolutional neural network, third WWA-based Swin Transformer module, fourth downsampling convolutional neural network, fourth WWA-based Swin Transformer module; The decoder includes: first to fifth upsampling convolutional neural networks and a second convolutional neural network connected in sequence; The output of the first feature resolution restoration module is connected to the input of the second WWA-based SwinTransformer module and the input of the second feature resolution restoration module. The output of the second feature resolution restoration module is connected to the input of the third feature resolution restoration module and the fifth upsampling convolutional neural network. The output of the third feature resolution restoration module is connected to the input of the second convolutional neural network. The output of the second convolutional neural network is connected to the spatial transformation network through the deformation field. The output end of the fourth WWA-based Swin Transformer module is connected to the input end of the first upsampling convolutional neural network, the output end of the third WWA-based Swin Transformer module is connected to the input end of the second upsampling convolutional neural network, and the output end of the first WWA-based Swin Transformer module is connected to the first feature resolution restoration module and the third upsampling convolutional neural network. Each of the first to fourth WWA-based Swin Transformer modules consists of two consecutive Swin Transformer submodules, and the input feature of the Swin Transformer module is set to z 1 , the calculation formula of two consecutive WWA-based Swin Transformer submodules is as follows: Where W-MSA is the standard window multi-head attention, SW-MSA is the shifted window multi-head attention, MLP is the multi-layer perception module, LN is the layer normalization operation, For WWA, LN and WMSA and Z 1 The calculated characteristic matrix, Z 2 is calculated by LN and MLP, and Z 1 The obtained feature matrix, For the SW-MSA, LN and WWA and Z 2 The calculated characteristic matrix, Z 3 After MLP and LN and The calculated feature matrix; The weighted attention mechanism WWA of window features has two multi-layer perceptron module MLP layers; the first MLP layer calculates the global channel-wise feature interactions across channels; the second MLP layer calculates the feature interactions between windows across windows.

2. The system of deformable medical image registration model based on Swin Transformer according to claim 1, characterized in that: The first to fourth downsampling convolutional neural networks contain one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 2; the convolutional neural network contains two convolutional layers with a convolution kernel size of 3×3×3 and a stride of 1; the first to third upsampling convolutional neural networks contain one inverted convolutional layer with a convolution kernel size of 2×2×2 and a stride of 2, and one convolutional layer with a convolution kernel size of 3×3×3 and a stride of 1.

3. A deformable medical image registration method based on Swin Transformer, characterized in that: A system using the Swin Transformer-based deformable medical image registration model according to any one of claims 1 to 2, comprising: S1: Concatenate the 3D fixed image and floating image of the given depth D, height H, width W, and the first channel into the image pair of the second channel, and input them into the first downsampling convolutional neural network; S2: After two consecutive first and second downsampling convolutional neural networks, the first convolutional neural network is used for dimensionality increase and the output is sent to the first WWA-based Swin Transformer module; S3: The first WWA-based Swin Transformer module extracts features and outputs them to the first feature resolution recovery module; S4: The first to third feature resolution recovery modules all amplify the feature resolution; S5: The feature map output by the first feature resolution recovery module is input into the second WWA-based Swin Transformer module, and features are extracted in sequence until the features are finally output to the decoder; S6: The first to fifth upsampling convolutional neural networks in the decoder sequentially upsample the input features; S7: The decoder connects the corresponding features in the encoder and decoder through jump connections, outputs the D×H×W×3 deformation field to the spatial transformation network, and the image to be deformed and the deformation field are input into the spatial transformation network, which outputs the deformed image.

4. The deformable medical image registration method based on Swin Transformer according to claim 3, characterized in that: The model as a whole outputs the deformation field of the third channel, and uses the spatial transformation network to apply the output deformation field to the input image to output the deformed image.

5. The deformable medical image registration method based on Swin Transformer according to claim 4, characterized in that: Given image features of depth D, height H, width W, and channel C, the feature resolution module expands the features in the channel dimension. After expansion, the feature size is 8C×D×H×W. The expanded features are then reorganized into C×2D×2H×2W image features to restore the feature resolution.

Images