Abdominal ct image segmentation method based on branch growing neural network architecture search

Abstract

The application discloses a method for abdominal CT image segmentation based on branch growth neural network architecture search, comprising the following steps: constructing a search space network based on local width branch and global depth branch; alternately connecting the local width branch and the global depth branch to the growable down-sampling layer of an encoder based on the search space network to construct a grown candidate network; training and testing the grown candidate network to obtain an image segmentation model; inputting the abdominal CT image data to be tested into the image segmentation model to complete abdominal CT image segmentation. The application improves the accuracy of abdominal CT image segmentation.

Description

Technical Field

[0001] The present invention belongs to the technical field of medical image processing, and in particular relates to an abdominal CT image segmentation method based on branch-growth neural network architecture search. Background Art

[0002] Medical image segmentation is a key step in medical image analysis. Convolutional neural networks (CNNs) have been increasingly widely used in medical image segmentation, with the encoder-decoder U-Net being a prominent example. Numerous improved models based on U-Net have subsequently been proposed. For example, Xiao et al. combined residual blocks with U-Net to propose Res-UNet. By adding residual connections between convolutional blocks, they achieved promising results in retinal vessel segmentation. To enhance contextual information extraction, Huang et al. proposed FECC-Net, which added an ASPP module and an enhanced encoder branch to U-Net, improving the accuracy of brain MRI image segmentation. Oktay et al. combined the attention mechanism with U-Net to propose Attention U-Net, achieving good performance in abdominal pancreas segmentation. While CNNs have made significant progress in medical segmentation tasks, they still have inherent limitations in capturing global features. Their limited receptive field makes them incapable of capturing global information, resulting in insufficient global modeling capabilities. To address this issue, researchers have proposed various improvements, such as dense connections, image pyramids, and self-attention modules. At the same time, inspired by the success of the Transformer in natural language processing, a visual attention model (ViT) was proposed based on the Transformer structure. This model uses a self-attention mechanism to capture the dependencies between positions in the input sequence, successfully solving the problem of long-distance dependencies. Learning global dependencies across all image patches through the self-attention mechanism has also shown excellent results in the field of medical image segmentation. For example, Chen et al. first utilized the global context modeling capabilities of the Transformer and the high-resolution localization capabilities of the U-Net in medical image segmentation, achieving good performance in multi-organ segmentation tasks. Zhang et al. proposed a model with two encoders, a Transformer branch and a CNN branch, in parallel, achieving good segmentation results in tasks such as colon polyps, skin, and prostate segmentation. Although these methods have achieved good results in the field of medical segmentation, these networks are all hand-designed and have poor adaptability. Summary of the Invention

[0003] To solve the above technical problems, the present invention proposes an abdominal CT image segmentation method based on branch-growing neural network architecture search to improve the accuracy of abdominal CT image segmentation.

[0004] To achieve the above object, the application provides an abdominal CT image segmentation method based on branch growth neural network architecture search, comprising:

[0005] Based on local width branch and global depth branch, a search space network is constructed;

[0006] Based on the search space network, the local width branch and the global depth branch are alternately connected to the growable down-sampling layer of the encoder to construct a grown candidate network;

[0007] The grown candidate network is trained and tested to obtain an image segmentation model;

[0008] Abdominal CT image data to be tested is obtained, which is input into the image segmentation model to complete abdominal CT image segmentation.

[0009] Optionally, the local width branch Conv and Res2net are composed, and a Conv+Res2net×N structure is adopted, wherein Res2net×N represents N Res2net modules.

[0010] Optionally, the global depth branch is composed of a Transformer layer, and the Transformer layer captures global features of long-distance dependence in a global range through a self-attention mechanism.

[0011] Optionally, the growable down-sampling layer is composed of a pooling layer, a convolution layer and a branch connection layer;

[0012] The pooling layer is used to reduce the scale of a feature map to facilitate the extraction of different scale features.

[0013] The convolution layer is used to extract local features.

[0014] The branch connection layer is used to connect a new branch to enable the network to continuously grow.

[0015] Optionally, the branch connection layer is provided with a visited attribute for distinguishing whether it is connected by a branch, including a case without branch connection and a case with branch connection.

[0016] In the case without branch connection, the branch connection layer is directly skipped, and the result processed by the convolution layer is taken as the final output.

[0017] In the case with branch connection, the output processed by the convolution layer is fused with the connected branch through a concatenation layer, and then dimension normalization is performed through convolution.

[0018] Optionally, the implementation process of the grown candidate network comprises:

[0019] The first generation growth initial encoder accesses the growable down-sampling layer searched by the encoder and accesses the first local width branch, and the first generation growth of the encoder is obtained after growth; the first generation growth of the encoder grows from the depth direction, accesses the first global depth branch after the last growable down-sampling layer, and obtains the second generation growth of the encoder; the second generation growth of the encoder continues to search the growable down-sampling layer accessed and accesses the new local width branch, and obtains the third generation growth of the encoder; the third generation growth of the encoder accesses the new global depth branch after the last global depth branch, and the above is repeated.

[0020] The technical effect of the present application: the present application discloses an abdominal CT image segmentation method based on branch growth neural network architecture search, in order to automatically construct a network model, the present application proposes a local-global branch iterative growth algorithm, which makes the network grow on the last generation network model, realizes the dynamic expansion of the network architecture, and the training of each generation is initialized based on the network parameters of the last generation, thereby reducing the training time. In order to make the network effectively obtain multi-scale features and global context information, the present application designs growth branches based on Res2net modules and Transformer modules, grows the network in two directions of depth and width, and then obtains more rich semantic information. The generated network is verified on the Synapse multi-organ CT image dataset, and the experimental results show that the present application has higher evaluation accuracy compared with the artificially designed image segmentation method. BRIEF DESCRIPTION OF DRAWINGS

[0021] The accompanying drawings, which form a part of this application, are included to provide a further understanding of the application and are incorporated in and constitute a part of this application. The embodiments of this application and its description are used to explain the application without imposing any undue limitation. In the drawings:

[0022] Figure 1 The flowchart of the abdominal CT image segmentation method based on branch growth neural network architecture search of the embodiments of the present application is shown in the figure;

[0023] Figure 2 The local width branch and the global depth branch of the embodiments of the present application are shown in the figure, wherein (a) is the local width branch, and (b) is the global depth branch;

[0024] Figure 3 The schematic diagram of a growable down-sampling layer of the embodiments of the present application is shown in the figure;

[0025] Figure 4 The size change schematic diagram of the embodiments of the present application is shown in the figure, wherein (a) is the branchless access, and (b) is the branch access;

[0026] Figure 5 The local-global branch iterative growth strategy schematic diagram of the embodiments of the present application is shown in the figure;

[0027] Figure 6 This is a schematic diagram of the network architecture of LG-BG-NAS according to an embodiment of the present invention;

[0028] Figure 7 Schematic diagram of the local-global branch iterative growth process according to an embodiment of the present invention;

[0029] Figure 8 This is a schematic diagram for visually comparing the segmentation results of different methods on the Synapse dataset according to an embodiment of the present invention. DETAILED DESCRIPTION

[0030] It should be noted that, in the absence of conflict, the embodiments and features of the embodiments in this application can be combined with each other. The present application will be described in detail below with reference to the accompanying drawings and in combination with the embodiments.

[0031] It should be noted that the steps shown in the flowcharts of the accompanying drawings can be executed in a computer system such as a set of computer-executable instructions, and that, although a logical order is shown in the flowcharts, in some cases, the steps shown or described can be executed in an order different from that shown here.

[0032] like Figure 1 As shown, this embodiment provides an abdominal CT image segmentation method based on branch-growing neural network architecture search, including:

[0033] Construct a search space network based on local width branches and global depth branches;

[0034] Based on the search space network, the encoder's growable downsampling layer is alternately connected to the local width branch and the global depth branch to construct a candidate network after growth;

[0035] Train and test the grown candidate network to obtain an image segmentation model;

[0036] Obtain the abdominal CT image data to be tested, input the image segmentation model, and complete the abdominal CT image segmentation.

[0037] Specifically, this embodiment proposes a local-global branch iterative growth algorithm that allows the network to grow on the previous generation network model, achieving dynamic expansion of the network architecture. Each generation of training is initialized based on the parameters of the previous generation network, thereby reducing training time. In order to enable the network to effectively obtain multi-scale features and global context information, this embodiment designs growth branches based on the Res2net module and the Transformer module, growing the network in both depth and width, thereby obtaining richer semantic information.

[0038] Further, such asFigure 2 The local width branch (LB) shown in (a) consists of a ConvNet and a Res2Net, using a Conv+Res2Net×N structure, where Res2Net×N represents N Res2Net modules. Its primary purpose is to extract local features of the input data. To capture richer local features using multi-branch convolutional layers, the LB is typically grown across the width of the encoder.

[0039] Further, such as Figure 2 The global depth branch (GB) shown in (b) consists of a Transformer layer, which consists of layer normalization, multi-head attention, layer normalization, and a multi-layer perceptron. The Transformer layer captures long-distance dependent global features globally through the self-attention mechanism.

[0040] Furthermore, in order to dynamically insert branches into the encoder to achieve branch growth, a growable downsampling layer (Grownable pooling, GPool) is designed. Figure 3 As shown in the figure, GPool consists of a pooling layer, a convolutional layer, and a branch access layer (Accesslayer), which allows branches to access the network through the Access layer, allowing the network to continue to grow. A visited attribute is set in the Access layer to distinguish whether it is accessed by a branch. When there is no branch access (visited = false), the Access layer will be skipped directly, and the result processed by the Conv layer will be used as the final output. When there is a branch access (visited = true), the output processed by the Conv layer will be concatenated (Concat) with the access branch, and the dimension will be normalized through 1×1 convolution. Assume that the input size is (3, H, W), and the size change of the feature map is as follows Figure 4 As shown. Figure 4 In (a), there is no branch access, the input passes through the convolution layer, the pooling layer in GPool and the convolution layer, and the final output size is 2C, H / 2, W / 2; when there is a branch access, as shown in Figure 4 As shown in (b), the input passes through the convolutional layer and the Res2net layer, so that the final output is also 2C, H / 2, W / 2, so as to facilitate concat with the original branch.

[0041] In order to improve the local and global feature extraction capabilities of the candidate neural network encoder during the search, a local-global branch iterative growth strategy is adopted. First, the LB is used to extract richer fine-grained local features at similar scales, and then the GB is grown to better capture the global feature relationship. The above branch growth process is repeated until the algorithm converges. The local-global branch iterative growth strategy is as follows: Figure 5 As shown in Figure 1, LB growth is performed first, followed by GB growth, and LB is connected to the network from shallow to deep. m Indicates the encoder's mth generation growth, with the initial encoder being G0. First, the first generation of growth, G0, searches for accessible GPools from the encoder and accesses the first local width branch (LB1), obtaining G1 after growth. Then, G1 grows in the depth direction, i.e., after the last GPool, it accesses the first global depth branch (GB1), obtaining G2. G2 continues to search for accessible GPool layers and accesses new LBs. m / 2+1 , get G3; then G3 accesses the new one after the last GB In order to ensure that the two inputs of the Access layer (from LB and GPool itself) have the same shape, LB i There should be a Res2net module in it.

[0042] The training and testing of the candidate network after growth specifically includes: using the encoding-decoding structure U-Net as the initial network, after multiple generations of LG-BG-NAS growth, the trained encoding-decoding architecture neural network is directly used for testing, the structure is as follows Figure 6 As shown in Figure 2, the decoder part is consistent with U-Net, consisting of L upsampling and a 1×1 convolution. The network does not require retraining or parameter adjustment and can directly segment images in the test dataset.

[0043] A specific experimental example of this embodiment is as follows:

[0044] The experiment is performed on a computer configured with an Intel i5 processor, DDR4 16G x 2 memory and an RTX 3060 GPU using Pytorch. The average similarity coefficient (Dice) and the average Hausdorff distance (HD) are used to compare with the SOTA methods such as Att-UNet, TransUnet, SwinUnet, etc. to verify the effectiveness of the method. The network structure of U-Net is the same as the initial encoder in the method proposed in the embodiment; TransUnet is a hybrid model combining CNN and Transformer, thus having the advantages of both CNN and Transformers; SwinUnet is a pure Transformer model. Therefore, these methods are selected for comparison with the proposed method. The experiment uses the Synapse multi-organ dataset for training and testing. The Synapse dataset includes 30 cases of 3779 abdominal axial clinical CT images, 18 samples are divided into a training set, and 12 samples are divided into a test set. All 3D bodies are segmented in a piece-by-piece manner to evaluate 9 abdominal organs (aorta, gallbladder, spleen, left kidney, right kidney, liver, pancreas, spleen, and stomach).

[0045] The experiment uses 224x224 pictures as input, the batch size is 3, the SGD optimization algorithm is used during training, the initial learning rate is 0.01, the momentum is set to 0.9, the weight decay is 0.0001, and the cosine annealing method is used to update the learning rate. The epochs of G0 are 200, and the epochs of G1-G5 are 60. The threshold T is set to 0.1, and L is set to 4.

[0046] The average Dice coefficient and the average HD are used to evaluate the segmentation accuracy of the model. The Dice coefficient is a set similarity index used to calculate the similarity of two samples, and its calculation formula is as follows:

[0047]

[0048] Where |X| and |Y| represent the true label and the prediction result respectively, and |X∩Y| represents the intersection thereof.

[0049] HD is also a set similarity index, which defines the similarity between two sets of points by calculating the distance between two points, and the smaller the value, the higher the similarity between the two sets. The calculation formula of HD is as follows:

[0050] HD(A, B) = max(h(A, B), h(B, A)) (2)

[0051]

[0052] Where ‖·‖ is the distance form between point set A and point set B, h(A, B) and h(B, A) are the one-way Hausdorff distances from set A to set B and from set B to set A, respectively.

[0053] First, we search for candidate networks on the Synapse dataset. Since this dataset does not have an independent validation set, we directly perform validation on its test set. The network growth process is as follows: Figure 7 As shown in the figure, width-wise growth is first performed, starting with GPool1 in the encoder to search for an accessible GPool layer. Initially, if GPool1 has no branches connected, a corresponding LB is generated and connected to GPool1, causing G0 to grow into G1 and then train. The network parameters of G1 are initialized based on the trained G0, so there is no need to retrain the entire network, only the newly connected branches. Then, the trained G1 is grown in the depth direction, that is, a GB is grown after the last GPool layer, causing G1 to grow into G2. Similarly, G2 is initialized and trained based on the network parameters of the trained G1. G2 then continues to grow in the width direction to generate G3 and train it. Depth-wise growth is then performed on G3 to obtain G4. This growth process is iterated continuously until the average Dice improvement is less than the threshold T.

[0054] Table 1 shows the average Dice, training time, number of parameters, and computational effort for G0-G5. First, because G0 trains for 200 epochs, training time is relatively long. From G1 to G5, the training time of each generation of networks increases with the addition of branches. The addition of LB significantly increases training time, while the addition of GB has little effect on training time. Both the number of parameters and computational effort show a gradual increase from G0 to G5. The increase in parameters is more pronounced after the addition of GB, which is due to the Transformer's multi-head self-attention mechanism. Furthermore, the computational effort increases even more significantly after the addition of LB, due to the parallel convolution operations of the Res2Net module. Overall, each generation of networks shows significant improvements in parameter and computational effort, while training time fluctuates significantly after the addition of LB, reflecting the increase in computing resources and time consumption while improving network performance. From the average Dice in the table, we can see that with the increase of growth generations, the average Dice growth rate decreases. After 4 generations of growth, the average Dice is basically stable. When it grows to the G5 model, the average Dice value is less than the threshold T compared with G4, and the HD value is not as good as G4. Therefore, G4 is selected as the best model.

[0055] Table 1

[0056]

[0057] Table 2 compares the test results of LG-BG-NAS with traditional CNN segmentation models such as V-Net, DARR, R50 U-Net, and Att-UNet on the Synapse dataset. LG-BG-NAS performs best, with G4 achieving an average Dice value and HD value of 83.07% and 18.87 mm on the test set, respectively. Att-UNet, which ranks second among these methods, incorporates an attention mechanism into skip connections, achieving an average Dice value and HD value of 77.77% and 36.02 mm, respectively.

[0058] Table 2 also compares the test results of CNN+Transformer hybrid models on the Synapse dataset, including R50 ViT, TransUnet, Swin UNet, and TransClaw U-Net. Again, LG-BG-NAS performed best. Swin Unet, ranked second, uses a pure Transformer encoder-decoder architecture, achieving average Dice and HD values ​​of 79.13% and 21.55 mm, respectively. TransClaw U-Net, ranked third, combines the Transformer and Claw U-Net architectures, achieving average Dice and HD values ​​of 78.09% and 26.38 mm, respectively. TransUnet has a similar architecture to LG-BG-NAS, but with fixed CNN and Transformer modules. Its average Dice and HD values ​​are 77.48% and 31.69 mm, respectively. We also tested other networks during the growth process. When G0 grows to G4, the average Dice increases rapidly; when G4 grows to G5, the average Dice increases only 0.05%, reaching the algorithm's convergence condition and no further growth is achieved.

[0059] Table 2

[0060]

[0061] Figure 8 Shows a visual comparison of the segmentation results of different methods on the Synapse dataset. Figure 8 As can be seen, the Att-Unet method is most prone to over-segmentation or under-segmentation, such as the over-segmentation of the gallbladder (green) in the first sample (row 1) and the under-segmentation of the stomach (light blue) in the second sample (row 2). TransUnet and Swin Unet also suffer from under-segmentation, such as the pancreas (cyan) in the third sample (row 3). Compared to other Transformer-incorporated models, LG-BG-NAS uses LB to obtain multi-scale feature information and GB to better capture global feature relationships, achieving better segmentation results than other methods.

[0062] The application discloses a method for abdominal CT image segmentation based on branch growing neural network architecture search, in order to automatically construct a network model, the application proposes a local-global branch iterative growing algorithm, which enables the network to grow on the last generation network model, realizes dynamic expansion of the network architecture, and the training of each generation is initialized based on the network parameters of the last generation, thereby reducing the training time. In order to enable the network to effectively obtain multi-scale features and global context information, the application designs a growing branch based on a Res2net module and a Transformer module, grows the network in two directions of depth and width, and then obtains more rich semantic information. The generated network is verified on a Synapse multi-organ CT image dataset, and experimental results show that the application has higher evaluation precision compared with an artificially designed image segmentation method.

[0063] The above is only a preferred specific embodiment of the application, but the protection scope of the application is not limited thereto, any person skilled in the art can easily think of changes or replacements within the technical range disclosed by the application, which should be covered in the protection scope of the application. Therefore, the protection scope of the application should be subject to the protection scope of the claims.

Claims

1. An abdominal CT image segmentation method based on branch-growing neural network architecture search, characterized in that: include: Construct a search space network based on local width branches and global depth branches; Alternately accessing the local width branch and the global depth branch to the growable downsampling layer of the encoder based on the search space network to construct a grown candidate network; Training and testing the grown candidate network to obtain an image segmentation model; Obtaining abdominal CT image data to be tested, inputting the image segmentation model, and completing abdominal CT image segmentation; The local width branch consists of Conv and Res2net, and adopts a Conv+Res2net×N structure, where Res2net×N represents N Res2net modules; The global depth branch consists of a Transformer layer that captures long-range dependent global features globally through a self-attention mechanism; The growable downsampling layer consists of a pooling layer, a convolutional layer, and a branch access layer; The pooling layer is used to reduce the scale of the feature map to facilitate the extraction of features of different scales; The convolutional layer is used to extract local features; The branch access layer is used to access new branches to enable the network to continue to grow; The branch access layer sets a visited attribute to distinguish whether it is accessed by a branch, including a case where there is no branch access and a case where there is branch access; In the case of no branch access, the branch access layer will be skipped directly, and the result processed by the convolution layer will be used as the final output; In the case of branch access, the output of the convolution layer is fused with the access branch through the splicing layer, and the dimension is normalized through convolution; The implementation process of the grown candidate network includes: The first generation growth initial encoder searches for the connected growable downsampling layer from the encoder and connects to the first local width branch, and after growth, the first generation growth of the encoder is obtained; the first generation growth of the encoder grows from the depth direction, and connects to the first global depth branch after the last growable downsampling layer to obtain the second generation growth of the encoder; the second generation growth of the encoder continues to search for the connected growable downsampling layer and connects to the new local width branch to obtain the third generation growth of the encoder; the third generation growth of the encoder connects to a new global depth branch after the last global depth branch, and so on.

Images