Multi-task registration and segmentation method based on multi-scale attention and accumulation learning

Abstract

The application discloses a multi-task registration and segmentation method based on multi-scale attention and accumulation learning, and realizes mutual promotion and enhancement between registration and segmentation tasks. The application realizes efficient feature coding between tasks through a double-flow shared encoder, a multi-scale attention module guides the model to focus on challenging areas, and further improves the accuracy of registration and segmentation through an accumulation learning strategy, efficiently realizes multi-task registration and segmentation of medical images, and obtains registration and segmentation accuracy superior to current mainstream methods. The application has important clinical application significance and reference value for disease diagnosis, lesion monitoring and surgical navigation assisted by medical images.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of medical image processing, and particularly to a multi-task registration and segmentation method based on multi-scale attention and accumulation learning. BACKGROUND

[0002] Lung cancer is one of the most serious diseases that threaten human life, and the initial clinical manifestation of lung cancer is a solitary spherical nodule, i.e., a lung nodule. In the efficient diagnosis of lung nodules by medical assistants, an accurate registration and segmentation method plays a crucial role.

[0003] Medical image registration and segmentation are two key technologies in the field of medical image processing, and they have wide applications in clinics. Among them, the medical image registration technology models a deformation field to warp the floating image and align it with the reference image, which is a key technology used in disease diagnosis, radiotherapy and surgical navigation. The medical image segmentation technology is used to segment specific organs or lesions, and is widely used in intraoperative auxiliary positioning and lesion growth monitoring.

[0004] Although the current medical image registration and segmentation methods have made remarkable development, they still face some challenges. For example, registration is usually modeled in an unsupervised form, which is essentially an ill-posed problem; and in the absence of prior information, the feasibility of registration will be greatly reduced. The segmentation of specific organs or lesions requires the supervision of segmentation labels, and the acquisition of segmentation labels requires a certain amount of time, effort and professional knowledge; it is a costly process, especially for 3D medical images.

[0005] Registration and segmentation also have high relevance and complementarity: the deformation field generated by registration can provide deformation augmentation for the segmentation task, and the anatomical prior information obtained by segmentation can provide auxiliary supervision for registration. In addition, both tasks are usually modeled as the same structure, both containing an encoder and a decoder. Therefore, applying multi-task learning to registration and segmentation can achieve mutual promotion and enhancement between the two tasks. However, registration usually requires splicing input images for feature encoding, while segmentation does not. The different encoding forms of registration and segmentation prevent the establishment of an efficient multi-task feature sharing mechanism between the two tasks. SUMMARY

[0006] The present application aims to overcome the shortcomings and deficiencies of the prior art, and proposes a multi-task registration and segmentation method based on multi-scale attention and accumulation learning, which realizes the mutual promotion and enhancement between the registration and segmentation tasks, and obtains a performance superior to the most advanced multi-task registration and segmentation method, which has important clinical application significance and reference value for medical image assisted disease diagnosis, lesion monitoring and surgical navigation.

[0007] To achieve the above object, the technical scheme provided by the present application is: a multi-task registration and segmentation method based on multi-scale attention and accumulation learning, comprising the following steps:

[0008] 1) input the reference image and the floating image into a double-stream shared encoder for feature extraction, to obtain the encoded feature maps of the reference image and the floating image respectively;

[0009] 2) input the encoded feature maps of the reference image and the floating image into a multi-scale attention module to extract multi-scale anatomical attention;

[0010] 3) share the multi-scale anatomical attention to a registration decoder and a segmentation decoder, output the deformation field from the floating image to the reference image by the registration decoder, and output the segmentation results of the reference image and the floating image by the segmentation decoder;

[0011] 4) the registration process is supervised by the similarity loss of the reference image and the deformation image the regularization loss of the deformation field together, the segmentation process is supervised by the segmentation loss of the segmentation results of the reference image and the floating image and the corresponding segmentation labels ; wherein the consistency loss is used to ensure the consistency of the registration and the segmentation on the anatomical structure.

[0012] Further, in step 1), the double-stream shared encoder is constructed with Swin Transformer as the backbone network; the double-stream shared encoder is defined as DSSE(f,m), and the encoded feature maps of the reference image and the floating image extracted therefrom are represented as:

[0013] F1,F2,…,F 10 =DSSE(f,m)

[0014] In the formula, F1,F2,…,F5 represent the encoded feature maps of different resolution levels extracted from the reference image, and F6,F7,…,F 10 represent the encoded feature maps of different resolution levels extracted from the floating image; f is the reference image, and m is the floating image.

[0015] Further, in step 2), the multi-scale attention module is defined as MSAM(f,m), and the multi-scale anatomical attention extracted therefrom is represented as:

[0016]

[0017] In the formula, is the multi-scale anatomical attention map for segmentation decoding, is the multi-scale anatomical attention map for registration decoding; denotes the segmentation decoded feature map from the lower resolution level, F i denotes the encoded feature map from the current resolution level.

[0018] Further, in step 3), the registration decoder and the segmentation decoder both apply the accumulated learning strategy, and are constructed with convolution as the backbone network; the decoding process of the registration decoder applying the accumulated learning strategy is represented as follows:

[0019]

[0020] where CONV(·) represents convolution operation, RH(·) is a registration head, and UPS(·) represents up-sampling operation; r i is the registration decoded feature map of the current resolution level, is the registration decoded feature map from the lower resolution level; φ i and are the deformation fields of the current resolution level before and after up-sampling, respectively, is the deformation field from the lower resolution level; is the multi-scale anatomical attention map of the floating image, represents deformation operation;

[0021] The decoding process of the segmentation decoder applying the accumulated learning strategy is represented as follows:

[0022]

[0023] where SH(·) is a segmentation head, s i is the segmentation decoded feature map of the current resolution level, is the segmentation decoded feature map from the lower resolution level; and are the deformation fields of the current resolution level before and after up-sampling, respectively, is the deformation field from the lower resolution level.

[0024] Further, in step 4), the similarity loss of the reference image and the deformation image is specifically represented as:

[0025]

[0026] where MSE(,) is a mean square error operator; is the deformation image, which is generated by deforming the floating image m by the deformation field φ; Ω is a three-dimensional space in which the reference image and the floating image are located, p is a pixel in Ω, and p∈Ω;

[0027] The regularization loss of the deformation field Specifically expressed as:

[0028]

[0029] In the formula, The gradient operator is ||·||, and the L2 norm operator is 2 The deformation field is defined as the sum of the identity transformation I and the displacement field u, that is, φ=I+u; therefore, the regularization of the deformation field φ is achieved by calculating the L2 norm of the pixel displacement gradient In the displacement field u;

[0030] The segmentation loss Specifically expressed as:

[0031]

[0032] In the formula, f l And m l Are the segmentation labels of the reference image and the floating image respectively, f s And m s Are the segmentation results of the reference image and the floating image respectively; The Dice similarity loss function is represented by The binary cross-entropy loss function is represented by

[0033] The consistency loss Specifically expressed as:

[0034]

[0035] In the formula, The deformed floating image segmentation result is generated by deforming the floating image segmentation result m s Through the deformation field φ.

[0036] Compared with the prior art, the present application has the following advantages and beneficial effects:

[0037] 1. The present application realizes the joint learning of registration and segmentation in the form of multi-task, fully utilizes the correlation and complementarity between registration and segmentation tasks, and realizes the mutual promotion and enhancement between the two tasks.

[0038] 2. The present application realizes efficient feature sharing between registration and segmentation tasks, breaks through the obstacle of different coding forms of registration and segmentation through the proposed double-flow shared encoder, and realizes efficient feature sharing between the two tasks by using a unified feature coding form.

[0039] 3. This invention performs well in challenging registration and segmentation regions. It extracts multi-scale anatomical attention through the proposed multi-scale attention module and shares it with the registration decoder and segmentation decoder, guiding the model to focus on challenging registration and segmentation regions and achieving higher registration and segmentation accuracy.

[0040] 4. This invention applies an accumulation learning strategy to the registration and segmentation decoder, optimizes the registration process by achieving coarse-to-fine registration, and optimizes the segmentation process by fusing segmentation maps of multiple scales, thereby further improving the accuracy of registration and segmentation. Attached Figure Description

[0041] Figure 1 This is a flowchart illustrating the method of the present invention.

[0042] Figure 2 This is a 3D registration result diagram showing the comparison between the method of the present invention and the comparative method.

[0043] Figure 3 This is a visualization of the 3D segmentation results of the method of the present invention and the comparative method. Detailed Implementation

[0044] The present invention will be further described in detail below with reference to the embodiments and accompanying drawings, but the embodiments of the present invention are not limited thereto.

[0045] like Figure 1 As shown, this embodiment discloses a multi-task registration and segmentation method based on multi-scale attention and cumulative learning. Through multi-task learning, it achieves mutual promotion and enhancement between tasks. The specific implementation of this method includes the following steps:

[0046] 1) Input the reference image and the floating image into the dual-stream shared encoder for feature extraction to obtain the encoded feature maps of the reference image and the floating image respectively;

[0047] The dual-stream shared encoder is constructed using a Swing Transformer as its backbone network; defining the dual-stream shared encoder as DSSE(,), the encoded feature maps of the reference image and the floating image extracted by it are represented as follows:

[0048] F1,F2,…,F 10 =DSSE(f,m)

[0049] In the formula, F1, F2, ..., F5 represent coded feature maps at different resolution levels extracted from the reference image, and F6, F7, ..., F 10 This represents encoded feature maps at different resolution levels extracted from the floating image; f is the reference image, and m is the floating image.

[0050] 2) input the encoding feature maps of the reference image and the floating image into the multi-scale attention module to extract multi-scale anatomical attention;

[0051] The multi-scale attention module is defined as MSAM(,,), and the multi-scale anatomical attention extracted therefrom is represented as:

[0052]

[0053] In the formula, is a multi-scale anatomical attention map for segmentation decoding, is a multi-scale anatomical attention map for registration decoding; represents a segmentation decoding feature map from a lower resolution level, F i represents an encoding feature map from a current resolution level.

[0054] 3) share the multi-scale anatomical attention to the registration decoder and the segmentation decoder, output a deformation field from the floating image to the reference image by the registration decoder, and output a segmentation result of the reference image and the floating image by the segmentation decoder;

[0055] Both the registration decoder and the segmentation decoder apply an accumulated learning strategy and are constructed with a convolution as a backbone network; the decoding process of the registration decoder applying the accumulated learning strategy is represented as follows:

[0056]

[0057] In the formula, CONV(·) represents a convolution operation, RH(·) is a registration head, and UPS(·) represents an up-sampling operation; r i is a registration decoding feature map of a current resolution level, is a registration decoding feature map from a lower resolution level; φ i and are a deformation field of the current resolution level before and after up-sampling, respectively, is a deformation field from a lower resolution level; is a multi-scale anatomical attention map of the floating image, represents a deformation operation;

[0058] The decoding process of the segmentation decoder applying the accumulated learning strategy is represented as follows:

[0059]

[0060] In the formula, SH(·) is a segmentation head, and s i is a segmentation decoding feature map of a current resolution level, is a segmentation decoding feature map from a lower resolution level; and deformation field before and after upsampling respectively, deformation field from the lower resolution level.

[0061] 4) The registration process is supervised by the similarity loss of the reference image and the morphed image The regularization loss of the deformation field The segmentation process is supervised by the segmentation loss of the segmentation results of the reference image and the floating image and the corresponding segmentation labels The consistency loss is used to supervise the registration and segmentation processes to ensure the consistency of the registration and segmentation on the anatomical structure.

[0062] The similarity loss of the reference image and the morphed image Specifically represented as:

[0063]

[0064] In the formula, MSE(,) is the mean square error operator; is the morphed image, which is generated by morphing the floating image m through the deformation field φ; Ω is the three-dimensional space where the reference image and the floating image are located, p is a pixel in Ω, p∈Ω;

[0065] The regularization loss of the deformation field Specifically represented as:

[0066]

[0067] In the formula, is the gradient operator, and ||·|| is the L2 norm operator. 2 The deformation field φ is defined as the sum of the identity transformation I and the displacement field u, that is, φ=I+u; therefore, the regularization of the deformation field φ is achieved by calculating the L2 norm of the pixel displacement gradient in the displacement field u.

[0068] The segmentation loss Specifically represented as:

[0069]

[0070] In the formula, f l and m l are the segmentation labels of the reference image and the floating image respectively, f s and m s are the segmentation results of the reference image and the floating image respectively. represents the Dice similarity loss function, represents the binary cross-entropy loss function.

[0071] The consistency loss Specifically represented as:

[0072]

[0073] In the formula, is the deformed floating image segmentation result, which is generated by deforming the floating image segmentation result m s After deformation by the deformation field φ.

[0074] The effect of the multi-task registration and segmentation method based on multi-scale attention and accumulation learning described above in this embodiment will be further illustrated through the following experiments:

[0075] 1. Experimental conditions:

[0076] Experimental environment: The computing power platform of this experiment is NVIDIA Tesla V100 (32GB) GPU and Intel Xeon Processor (Skylake, IBRS) CPU, and PyTorch is used to build a deep learning framework and train a neural network.

[0077] Experimental data: 420 cases of chest follow-up CT from 210 high-risk patients with lung cancer are used as experimental data in this embodiment, which includes not only CT images, but also corresponding lung parenchyma segmentation labels and registration key point data.

[0078] Evaluation index: The evaluation index of the method registration performance adopts target registration error TRE (mm) and lung parenchyma DSC, and the evaluation index of the segmentation performance adopts IoU and ASD (mm). Among them, the target registration error TRE (mm) calculates the average Euclidean distance between key points, evaluates the accuracy of global image registration, and the smaller the value, the higher the registration accuracy; the lung parenchyma DSC measures the coincidence degree of the lung parenchyma, and the larger the value, the higher the coincidence degree of the lung parenchyma; the IoU calculates the ratio of the intersection and union between the segmentation result and the segmentation label, and the larger the value, the higher the coincidence degree; the ASD (mm) quantifies the distance between the boundaries of the segmentation result and the segmentation label, and the smaller the value, the higher the boundary coincidence degree.

[0079] Comparative method: In order to verify the effectiveness of the registration and segmentation of the method, a plurality of currently most advanced methods are selected for comparison experiments, including three registration methods (VoxelMorph-diff, TransMorph-diff, LapIRN), three segmentation methods (3D U-net, UNet++, nnU-Net), and two multi-task registration and segmentation methods (Deep Atlas, Cross-stitch).

[0080] 2. Experimental Results:

[0081] The registration and segmentation index quantification results of the method of this invention and the comparison method are shown in Table 1 below:

[0082] Table 1: Quantitative results (average) of registration and segmentation indices for different methods

[0083]

[0084]

[0085] As can be seen from Table 1, the method of the present invention has the smallest TRE (mm) value and the largest Lung DSC value, which proves that the method of the present invention has the highest registration accuracy compared with the comparison method. Moreover, the method of the present invention has the largest IoU value, and the ASD (mm) value is only larger than that of nnU-Net, but smaller than other methods, which proves that the method of the present invention has excellent segmentation accuracy.

[0086] The visualization 3D registration results of the method of this invention and the comparison method are as follows: Figure 2 As shown in the figure, the specific representation is a 3D overlay of the deformed image and the lung parenchyma label of the reference image generated by each method. The lung parenchyma labels of the deformed image in each sub-figure are from: (a) floating image segmentation label, (b) VoxelMorph-diff, (c) TransMorph-diff, (d) LapIRN, (e) Deep Atlas, (f) Cross-stitch, and (g) the method of this invention. Figure 2 It can be observed that before registration, there are significant differences in the lung parenchyma labels between the floating image and the reference image, especially in the lower right corner near the diaphragm where there is considerable deformation. Compared to the comparison method, the method of this invention achieves the highest registration accuracy, particularly for areas with large deformations.

[0087] The visualized 3D segmentation results of the method of this invention and the comparison method are as follows: Figure 3 As shown in the figure, the specific display is a 3D overlay of the segmentation results and segmentation labels of each method. The segmentation results in each sub-figure are from: (a) 3D U-net, (b) UNet++, (c) nnU-Ne, (d) Deep Atlas, (e) Cross-stitch, and (f) the method of this invention. Figure 3 It can be observed that, compared with the comparative method, the method of the present invention achieves the highest precision segmentation results, especially in the upper left region of the lung parenchyma edge.

[0088] The above experimental results show that the medical image multi-task registration and segmentation can be effectively realized by the application, and the highest registration and segmentation accuracy is obtained. The application realizes efficient feature coding between tasks through a double-flow shared encoder, guides the model to focus on challenging areas through a multi-scale attention module, and further improves the accuracy of registration and segmentation through an accumulation learning strategy, realizes mutual promotion and enhancement between registration and segmentation tasks, has important reference value for clinical practical application, and is worth promoting.

[0089] The above-described embodiments are only the preferred embodiments of the present application, and are not intended to limit the scope of the application. Any changes made in the shape and principle of the application should be covered within the scope of protection of the application.

Claims

1. A multi-task registration and segmentation method based on multi-scale attention and cumulative learning, characterized in that, The method comprises the following steps: 1) inputting a reference image and a floating image into a double-stream shared encoder for feature extraction, to obtain encoded feature maps of the reference image and the floating image respectively; The double-stream shared encoder is constructed with Swin Transformer as a backbone network; the encoded feature maps of the reference image and the floating image extracted by the double-stream shared encoder are represented as: F1, F2,..., F 10 = DSSE(f, m) where F1, F2,..., F5 represent the encoded feature maps of different resolution levels extracted from the reference image, F6, F7,..., F 10 represent the encoded feature maps of different resolution levels extracted from the floating image; f is the reference image and m is the floating image. 2) inputting the encoded feature maps of the reference image and the floating image into a multi-scale attention module, to extract multi-scale dissection attention; The multi-scale attention module is defined as MSAM(,), and the multi-scale dissection attention extracted by the multi-scale attention module is represented as: wherein is a multi-scale anatomical attention map for segmentation decoding, is a multi-scale anatomical attention map for registration decoding; denotes a segmentation decoding feature map from a lower resolution level, F i denotes an encoding feature map from the current resolution level; 3) sharing the multi-scale dissection attention to a registration decoder and a segmentation decoder, outputting a deformation field from the floating image to the reference image by the registration decoder, and outputting segmentation results of the reference image and the floating image by the segmentation decoder; 4) The registration process is supervised by a similarity loss between the reference image and the warped image The regularization loss of the warping field The segmentation process is supervised by a segmentation loss between the segmentation results of the reference image and the floating image and the corresponding segmentation labels The registration process is supervised by a consistency loss between the reference image and the warped image The consistency between registration and segmentation on anatomical structure is guaranteed.

2. The multi-task registration and segmentation method based on multi-scale attention and cumulative learning according to claim 1, characterized in that, In step 3), the registration decoder and the segmentation decoder both apply an accumulation learning strategy, and are constructed with convolution as a backbone network; the decoding process of the registration decoder applying the accumulation learning strategy is represented as follows: where CONV(·) represents a convolution operation, RH(·) is a registration head, and UPS(·) represents an up-sampling operation; r i is a registration-decoded feature map at the current resolution level, is a registration-decoded feature map from a lower resolution level; φ i and are deformation fields at the current resolution level before and after up-sampling, respectively, is a deformation field from a lower resolution level; is a multi-scale anatomical attention map of the floating image, represents a deformation operation; The decoding process of the segmentation decoder applying the accumulation learning strategy is represented as follows: where SH(·) is the segmentation head, s i is the segmentation-decoded feature map at the current resolution level, is the segmentation-decoded feature map from the lower resolution level; and are the current resolution level’s morphing fields before and after upsampling, respectively, is the morphing field from the lower resolution level.

3. The multi-task registration and segmentation method based on multi-scale attention and cumulative learning according to claim 2, characterized in that, In step 4), the similarity loss of the reference image and the morphed image Specifically represented as: where MSE(·) is the mean square error operator; is the deformed image, which is generated by deforming the floating image m by the deformation field φ; Ω is the three-dimensional space where the reference image and the floating image are located, p is a pixel in Ω, p ∈ Ω; The regularization loss of the deformation field Specifically represented as: wherein is the gradient operator, || · || 2 is the L2 norm operator, the deformation field φ is defined as the sum of the identity transformation I and the displacement field u, i.e. φ = I + u; thus, regularization of the deformation field φ is achieved by computing the L2 norm of the pixel displacement gradients in the displacement field u. The segmentation loss Specifically represented as: In the formula, f l and m l are the segmentation labels of the reference image and the floating image respectively, f s and m s are the segmentation results of the reference image and the floating image respectively. represents the Dice similarity loss function, represents the binary cross-entropy loss function. the consistency loss is specifically represented as: In the formula, is the segmentation result of the deformed floating image, which is generated by deforming the segmentation result m s generated after deforming by the deformation field φ.

Images