Exploiting full Resolution Feature Context for Liver Tumor and Vessel Segmentation via Integrate Framework: Application to Liver Tumor and Vessel 3D Reconstruction under embedded microprocessor

We introduce an encoder that can learn the global feature representation, which consists of a feature embedding module based on a feature extraction backbone and a feature extraction module that senses the semantically related information representation of the image based on the transformer [24]. This module adopts a brand-new feature extraction idea, by semantically representing the features of the picture and learning the global representation of semantic features.

The input image $i\in\mathbb{R}^{H\times W\times C}$ is first fed into the feature extraction backbone network. The network can extract the spatial information features of CT images and output the feature map $x\in\mathbb{R}^{H^{{}^{\prime}}\times W^{{}^{\prime}}\times C^{{}^{\prime}}}$. We divide the feature map x learned by the backbone into a series of patches $x_{p}^{i}\in\mathbb{R}^{P^{2}\times C},i=1,…,N$, where the size of each patch is $P\times P$,and the number of patches denote by $N=\frac{H^{{}^{\prime}}\times W^{{}^{\prime}}}{P^{2}}$. For each patch, we use a convolution operation with a kernel size of $P\times P$ to obtain the information $E_{info}^{i}$ of i-th patch to form an information matrix $\{E_{info}^{1},E_{info}^{2},...,E_{info}^{N}\}$. In order to better learn location information using Transformer, Dosovitskiy et al. [33] perform a learnable location embedding for each patch to obtain the location matrix $\{E_{pos}^{1},E_{pos}^{1},...,E_{pos}^{N}\}$ of the N patches. The feature of the i-th patch can be formulated by the following equation:

We adopt this position encoding method, so that the feature extraction module can effectively learn the position information of the features. We next input the above obtained feature matrix $E=\{E^{1},E^{2},…,E^{N}\}$ of x into multi Transformer layers to learn semantically representation of the feature map. In comparison to traditional convolution operation, transformer adopts a multi-head self-attention mechanism, and its core formulation is shown in equation 2:

where $h$ and $w$ denote by width and height of the feature matrix $E$ after feature extraction and location embedding. And $n$ is the number of self-attention mechanism heads. $Q_{ijk},K_{ijk},V_{ijk}$ denote the query, key and value obtained by three linear transformations of the input $E_{ij}$ in each self-attended head, respectively. $y\in\mathbb{R}^{h\texttimes w\texttimes C}$ denotes the output after one multi-headed self-attention. We stacked 12 transformer layers, and the output of the last layer can theoretically learn to incorporate a rich context feature representation of the CT image under a wider range of perceptual fields. We then feed the output of the Transformer layers into a three-layer convolution operation. The final output feature map consist global high-level abstract information, effectively solving the problem of missing information caused by perceptual field defects in traditional deep CNN networks. We call it semantic feature map.

Transformer-based extraction module is a very powerful for semantically information feature, because the Transformer feature extraction module has advantages in learning semantically related features. In many ways, however, Transformer is not an effective replacement for traditional convolutional operations. For extraction of more subtle feature in some images such as edge feature of interest regions and tiny vessel feature, CNN is nothing but the perfect solution. We designed a local residual network encoder based on multi-layer SEBottleNet stacking, as shown in Figure 2. The encoder consists of a six feature extraction module. A max pooling operation is performed to extract the high-level feature representation after feed feature map to each feature extraction block. The input CT image $x\in\mathbb{R}^{H\times W}$ is first fed forward to a CNN module for high-level feature extraction, and the feature map $u\in R^{H\times W\times C}$ is obtained. Then, the feature map $u$ is feed into a deep residual feature extractor stacked by five layers of SEBottleNet, each of which is used for learning the context features under the local perception field. BottleNet residual network[34] retains all the advantages of residual network and significantly reduces computation interval and computational burden. We introduced the Squeeze and Excitation(SE)[35] in the BottleNet to enhance the interdependence between feature map channels. The structure of the SEBottleNet is shown as in Figure 2. The mean value $e_{c}\in\mathbb{R}^{1\times 1\times C}$ of the feature embedding for each channel in the feature map $U\in\mathbb{R}^{H\times W\times C}$ can be obtained from the Squeeze section, as shown in the following equation:

Where the $u_{c}(i,j)\in\mathbb{R}^{1\times 1\times C}$ is the pixel in feature map $U$. The Excitation section can learn the feature weights $e_{c}$ for each channel by $s_{c}$:

Finally, the vector product $\tilde{O}$ of s and u is obtained by the Scale operation, and this is the final output of the SE module:

where $\tilde{O}_{c}$ is the feature map of a feature channel.

The SEBottleNet residual network splits the traditional convolutional operation into multiple modules to ensure that each module has a different feature extraction task. We introduced the Squeeze and Excitation module in the middle of the module to better learn the importance of the feature map channel dimensions, so that SEBottleNet has a stronger learning focus in the feature extraction process. Through the continuous stacking of SEBottleNet and maxpool, the encoder can continuously extract the local feature representation of the input CT image HIGH-level. Meanwhile, since each SEBottleNet is set with residual connections, it enables the encoder to effectively mitigate the degradation problem caused by network deepening.

Since the hepatic arterial vessels are very small and the margins of the liver tumor arblurd, further refining the segmentation of the vessels and liver is a challenging task. In order to allow the model to learn more detailed spatial feature features, we introduce the Edge Extraction Module, which is specially designed to learn the edge features of blood vessels and tumor regions of interest and fuse the edge features to the segmentation network. The structure of this module is shown in Figure 1 (d). The Edge Extraction Module takes the feature maps of feature extraction layers and the CT edge map(Figure 4(b)) extracted by the Canny algorithm as input, and predict the edge result $e\in\mathbb{R}^{H\times W}$. This module predicts edge information and fuses the predicted feature maps into the segmentation network. To accomplish this task, we process segmentation annotations to obtain edge annotations $e_{r}$(Figure 4(d)), which can be used as a supervision condition for this module.

In this module, we used Gated Excitation Convolution (GEC) layer. GEC is the most important unit in Edge extraction module and it can filter out some irrelevant information to help Edge extraction module focus on extracting image edge features. GEC is applied between the Edge extraction module and the feature extraction module. It uses gating mechanisms to deactivate its own activations that are not deemed relevant by the higher-level information contained in the extraction module[36]. At the same time, we introduce an excitation module in the gating activation layer to learn the importance between different feature maps.

We define $t_{i},c_{i}\in\mathbb{R}^{\frac{H}{2i}\times\frac{W}{2i}\times C}$ as the feature maps of the Transformer module and the local feature extraction module, and $i$ denote the number of locations. Before using the GEC module, $t_{i}$ and $c_{i}$ were fed into a convolutional layer $C_{1\times 1}$ to obtain the image-dimensional feature maps $t_{i}^{{}^{\prime}},c_{i}^{{}^{\prime}}\in\mathbb{R}^{H\times W}$. Let $e_{i}\in\mathbb{R}^{H\times W\times C}$ denote the feature map synthesized by the Edge Extraction Module. Given the feature maps $t_{i}^{{}^{\prime}},c_{i}^{{}^{\prime}}$ and $e_{i}$, an excitation convolutional layer is applied to generate the sigmoid activation $\alpha_{i}\in\mathbb{R}^{H\times W}$:

where $F_{se}$ denotes squeeze and excitation option as shown in Equation 3-5. Finally, $\alpha_{i}$ and $e_{i}$ is fed into Gated Convolution layer[37][38][36] and generate the $e_{i}^{{}^{\prime}}$. The Gated Convolution layer is compute as

Theoretically, GEC can be simply regard as a collection of attention for the spatial dimension and channel dimension of the feature map. Through GEC operation, the attention maps $\alpha_{i}$ selectively preserve the edge semantic features. We cancel the GEC operation on the shallow feature maps of the feature extractor, since the images feed to the convolution layer mainly learns the general low-level features, at the same time, the output feature map retains rich edge information. As the network deepens, the feature map will retain high-level features. Using GEC operation can effectively weight the useful edge information of high-level features in theory.

The Canny operator can effectively filter out the irrelevant features of the image to obtain the canny image as shown in Figure 4(b). We think it is applicable for medical image segmentation. Therefore, we firstly concatenate the canny image and the last GEC module output $e_{n}$. Then we feed them together with the output feature maps of the two feature extractors to the Fusion Module. At the same time, the edge extraction module uses edge loss as the loss function and edge label as the supervision to optimize the prediction edge map.

In the previous sections, we introduced two feature extraction structures that learn spatial features and semantically related features, respectively. At the same time, we introduce the edge feature extraction module to enhance the ability of the whole network to extract the target edge features. In this section, we introduce the multi-scale feature fusion decoding module to sample the semantic features learned by the three modules. This module takes the feature maps extracted by the three modules as the input and outputs the predicted category distribution map $\hat{y}\in\mathbb{R}^{K\times H\times W}$, where K represent the semantic classes.

We introduce a fusion module, which mainly fuses the feature maps of the three feature extraction modules. Figure 3 shows the structure of this module. We designed the module with reference to the spatial pyramid pooling(SPP). Firstly, the module uses $C_{1\times 1}$ and $C_{3\times 3}$ convolution to extract features from the concatenation result of the semantic feature map and spatial feature map respectively. Next, we feed it into the pooling layer and fused the edge feature map. Through the above operations, we obtained the feature maps of three different receptive fields. Finally, we sample and concatenate these three feature maps to output the fused feature map. Theoretically, the feature map output by these module can retain rich spatial features, semantic related features and edge features.

In the process of continuous feature extraction layer by layer in the coded network, the low-level information of the feature map is continuously filtered and the high-level information is extracted. UNet uses skip connections to conduct the feature maps of the encoding module of each stage to the decoding module of the corresponding stage, and the network can fully learn the feature maps of different levels of the image. We adopted the skip connection operation from UNet and introduced skip connections to different feature encoders to allow the whole network to better learn the feature information of different encoders at different levels. The skip connection introduced in the local feature extraction module is similar to the traditional UNet module, which combines the short-range skip connection (residual connection) and the long-range skip connection of SEBottleNet. As for the Transformer based feature extraction module, we first introduce skip connections in the encoding process of the backbone network to connect the intermediate feature maps in the forward propagation process of the backbone embedding network, which improves the low-level feature loss in the feature embedding process of the backbone network. Next, we add skip connections to the feature maps with global feature representations after Transformer feature encoding fusion to fuse the global low-level features. Eventually, after continuously fusing low-level feature maps of different scales, the decoder can learn the semantic information of images from coarse to fine.

We propose the edge feature extraction module to cooperate with the segmentation task of the model, so we train the model to complete semantic segmentation and edge information segmentation at the same time. We introduce the joint optimization of edge loss and segmentation loss respectively. At the same time, in order to better ensure the consistency of multi task learning optimization, we set up a regularization methods to balance the two losses.

We use Dice and Cross Entropy(CE) as the loss function of segmentation task to predict semantic segmentation $y$:

where $y\in\mathbb{R}^{H\times W}$ denote the real semantic label map of liver tumor and vessel. In the equation 8, $\lambda_{1}$ and $\lambda_{2}$ represent hyper parameters. As for edge prediction, we use Binary Cross Entropy(BCE) loss. In this experiment, the model mainly focuses on tumor and vascular segmentation. We extract their common edges to obtain edge label $\hat{e}\in\mathbb{R}^{H\times W}$ (Figure 4(d)) and take $\hat{e}$ as the loss supervision. Therefore, the edge loss can be expressed as:

where, $e$ represent the edge predict map of the edge extraction module. It is worth noting that in the optimization process, the parameters of feature extraction modules and edge extraction modules will be optimized based on loss. Next, we will input the feature map output by the edge extraction module into the fusion module to predict the segmentation results. Therefore, the prior knowledge learned by the edge extraction module was retained in $y$. At the same time, the segmentation loss will pay more attention to the edge features in the optimization process.

We introduce regularization methods to make the model cooperate better in the training process. As mentioned above, $y\in\mathbb{R}^{K\times H\times W}$ represents the predicted segmentation map and $e\in\mathbb{R}^{H\times W}$ represents the predicted edge graph. Therefore, we introduce shape regularization, which can be expressed as:

Where $\oplus$ represents the pixel-wise addition of $y$ without the background label map which can be implemented using kernel fixed convolution operator. This operation outputs a label map containing the predict region of the tumor and blood vessels. In particular, at the beginning of model training, because the edge extraction module cannot accurately predict the edge, loss does not play any role. We therefore introduced a dynamic adjustment strategy, which is set $\lambda_{4}$ to 0 before 100 epoch and $\lambda_{4}\geq 0$ after 100 epoch

Finally, the loss function of the model is:

We set the epoch to 300, the initial learning rate to 0.001 (using the cosine annealing learning rate decay method), and the batch size to 8. The model is trained using an SGD optimizer with a momentum of 0.9 and a weight decay of 1e-4.

The TransFusionNet can significantly learn full-resolution context feature information, and its segmentation effect in the public dataset of blood vessels and liver tumors is significant. However, due to the scarcity of the enhanced CT images of liver cancer after the screening, we only obtained CT images of 18 patients. Too little data will inevitably affect the performance of the model and deepen the over-fitting problem. For this purpose we introduce a transfer learning strategy, which does not require exactly representative training data and is able to take advantage of the similarity between datasets to capture specific prior knowledge during the training phase of the model in order to construct new segmentation models.

We first pre-trained the models using the publicly available datasets LITS and 3Dircadb to obtain a liver tumor segmentation model and a liver vascular segmentation model, respectively. Then, we use our liver tumor data and liver vascular data to retrain the two models obtained by pre-training, and obtain two liver tumor segmentation models and liver vascular segmentation models based on the training sample distribution of our dataset. When we need to perform segmentation of liver tumor and blood vessels of CT images, we only need to input one CT image, and the two models will segment the tumor and blood vessel parts of CT images respectively. The mask output from the two model segments is automatically fused into a single mask that contains the tumor and blood vessels with relative positions as the final output.

We choose 5 advanced segmentation models to compare with our method, the 5 models are SegNet[18], UNet[21], UNet++[23], UNet3+[39], TransUNet[40]. At the same time, we divide the proposed method into TransFusionNet(TFN) and TransFusionNet with edge module(TFNEdge), and evaluate them respectively. We first compare the segmentation effects of five models on blood vessels and tumors based on two public datasets:LITS(Tumor)and 3Dircadb(Vessel). Next, we use the LTBV dataset to fine-tune the five models and compare the segmentation effects of the five models.

From the above quantitative experimental methods, our model has the best performance in the segmentation of liver blood vessels and liver tumors. Next, we use TransFusionNet and other comparable models on a test case on the LTBV dataset to segment liver tumors and blood vessels and then visualize them. The first row of Figure 6 is the model’s segmentation of Vessel, the second row is the model’s segmentation of Tumor, and the third row is the result of fusing the first two rows of segmentation results in the same coordinate system.

From the perspective of visual analysis, SegNet and UNet are not accurate in segmenting the details of blood vessels. Although UNet++ can identify some details of blood vessels, the error rate is too high. TransFusionNet almost perfectly segmented the details of blood vessels, which is more accurate than TransUNet. This is attributed to SEBottleNet’s extraction of local receptive field information and the importance of different channels. In tumor segmentation, UNet++ has a great segmentation result for the edge of the tumor, while SegNet and UNet perform poorly in this respect. All comparison models have some wrong tumor segmentation, and TransFusionNet not only avoids these wrong segmentation but also can segment the edge and contour of the tumor accurately. We believe that after TransFusionNet extracts global and local information, the multi-scale feature fusion decoder almost perfectly restores the feature of the image, so that the segmentation accuracy is significantly improved, and the error rate is low.

In summary, the above comparison models are not accurate in segmenting tumors and blood vessels. They are easy to misclassify some areas that are not tumors, and they are not sensitive to the recognition of some fine blood vessels areas, resulting in incomplete blood vessel segmentation results. TransFusionNet can accurately segment the liver tumor regardless of its integrity or vascular continuity.

Medical image segmentation has a wide range of applications and research values in medical research and practice fields such as medical research, clinical diagnosis, pathological analysis, computer-assisted surgery, and three-dimensional simulations. In this experiment, we use knowledge distillation to generate a more lightweight model from the TransFusionNet and deploy it into JetsonTX2 embedded device. We fed the CT images to the JetsonTX2 embedded system to directly predict the results of 3D reconstruction of liver vessels and tumors.

We define segmentation model $\mathcal{F}:\mathbb{R}^{N\times H\times W}\rightarrow\mathbb{R}^{N\times C\times H\times W}$. this model has been fully trained. Next, the arterial phase CT-enhanced image $x\in\mathbb{R}^{N\times H\times W}$ feed into $\mathcal{F}$ to predict the liver and vessel label map $o\in\mathbb{R}^{N\times C\times H\times W}$ where $C=3$ in this segmentation task. We need to construct 3D reconstruction result according to the segmentation label map $o$. Thus, to obtain segmentation result $y\in\mathbb{R}^{N\times H\times W}$ from label map $o$, we use the following equation:

Where $\mathcal{G}$ denotes Gaussian filter. In order to better support the embedded platform, we refer to the method in [41] to distill the knowledge of the Transformer module. At the same time, we do the post-processing quantization operation for the middle layer of the trained model. After the above optimizations, we successfully deployed the model to the embedded processor

Finally, we save the reconstruction results in $.nrrd$ format and use 3Dslicer for visual display. The Figure 7 shows a comparison between the reconstruction results under embedded microprocessor and manual annotations of a typical patient. Except for some noise and loss of vessel details, the reconstruction results were very close to the actual annotation results. However, a detailed manual annotation requires a lot of time and effort, which significantly reflect the efficiency and accuracy of our proposed algorithm.