RotCAtt-TransUNet++: Novel Deep Neural Network for Sophisticated Cardiac Segmentation

Mathematical methodologies encompass cluster-based algorithms like K-means and active contour models reliant on local and global intensities [general_attention]. Nonetheless, challenges such as variations in tissue appearance, low resolution, and indistinct boundaries undermine the robustness of these approaches against noise and diverse contrasts in medical imaging. Machine learning techniques, including model-based (e.g. active shape and appearance models) and atlas-based methods, have not exhibited superior efficacy in this domain as they frequently necessitate substantial feature engineering or pre-existing knowledge to attain acceptable accuracy [deep_review]. More recently, Deep Learning (DL) techniques have emerged triumphant in various computer vision applications, including object recognition and semantic segmentation.

Through meticulous experimentation and ablation studies, we observed the efficacy of the UNet++ [unet_plusplus] architecture coupled with dense downsampling and skip connections to preserve crucial information in achieving superior segmentation results. We are also inspired from pyramid pooling at different scales of Zhao at al [pyramid]. Furthermore, [transunet] also demonstrated that intensive incorporation of low-level features generally leads to a better segmentation accuracy. Thus, instead of the conventional CNN-based feature extraction approach, such as ResNet-50 in TransUNet, we embrace dense downsampling alongside nested skip connections, yielding four distinct feature maps $X_{1},X_{2},X_{3},X_{4}$ at varying resolutions.

Unlike TransUNet and its variants which only embeds the last lowest-resolution feature maps, we employs linear embedding for multi-scale feature maps. Specifically, the first three feature maps $X_{1},X_{2},X_{3}$ undergo linear embedding with a different patch size $p$ to produce different embedded vector $z_{i}^{j}\in Z_{i}$, which simultaneously go through transformer blocks to capture the interactions between patches and rotatory attention mechanisms to aggregate the information from adjacent slices. Within these transformer blocks, comprising $N$ transformer layers, the embedded sequence patches traverse self-attention mechanisms and multilayer perceptrons, facilitating robust intra-slice information capture and yielding $E_{1},E_{2},E_{3}$.

The rotatory attention block, conceived to treat the batch size as a continuous slice, selectively processes three consecutive slices—designating the first as the left, the second as the target, and the third as the right—culminating in the production of $R_{1},R_{2},R_{3}$ encapsulating information from adjacent slices in the volumetric data. Integration of interslice and intraslice information yields $F_{1},F_{2},F_{3}$, which are then reconstructed to their original resolution via upsampling techniques, resulting in $O_{1},O_{2},O_{3}$.

Finally, $X_{4}$ undergoes concatenation with $O_{3}$, perpetuating this iterative process until the final segmentation map is obtained post $1\times 1$ convolution.

The input is structured as $(B,1,H,W)$, representing the batch size, the number of channels (typically 1 in medical segmentation), height, and width, respectively. The batch size is also considered here since it also represents the number of adjacent slices whose information would be aggregated in rotatory attention block. This input undergoes convolutional operations to yield $X_{1}^{1}$, with a shape of $(B,C,H,W)$, where $C$ is set to 64 in our network. Subsequently, the resulting feature maps are downsampled to obtain $X_{2}^{1}$, with dimensions of $(B,C\times 2,\frac{H}{2},\frac{W}{2})$. This $X_{2}^{1}$ is then upsampled to match the shape of $(B,C\times 2,H,W)$. Following this, $X_{2}^{1}$ and $X_{1}^{1}$ are concatenated along the $C$ axis, resulting in a shape of $(B,C\times 3,H,W)$, which then undergoes further convolution to produce $X_{1}^{2}$. This resultant tensor, $X_{1}^{2}$, shares the same shape as $X_{1}^{1}$ but encompasses aggregated information from $X_{2}^{1}$. This iterative process continues through subsequent lower resolution images. If we designate the desired number of different-resolution outputs as $D$, we then have $X_{i}^{j}$ $\forall i\in\{1,\dots,D-1\}$ and $\forall j\in\{1,\dots,D-i\}$, where $X_{i}^{j}$ has a shape of $(B,C\times 2^{i-1},\frac{H}{2^{i-1}},\frac{W}{2^{i-1}})$. Notably, the $D$-th resolution map has a shape of $(B,C^{D-2},\frac{H}{2^{D-1}},\frac{W}{2^{D-1}})$, same depth as $D-1$-th resolution map and bypasses both the Transformer block and Rotatory Attention block but is instead utilized for the decoder. Specifically, when choosing $D=4$, as in our case, the resulting feature maps are $X_{1}^{3}$, $X_{2}^{2}$, and $X_{3}^{1}$. For simplicity, these three $X$ tensors are denoted as $X_{i}$ for all $i\in\{1,2,3\}$. Subsequently, they are linearly embedded via convolution operations $E$ to produce patches represented as embedded vectors $z_{j}^{p_{i}}\in Z_{i}$ where $Z_{i}$ has shape $(B,n_{i},d_{f}^{i})$ and $1\leq j\leq n_{i}$. The number of tokens or sequence length and the feature dimension of $Z_{i}$ are denoted as $n_{i}={\frac{H_{i}\times W_{i}}{p_{i}^{2}}}$ and $d^{f}_{i}$ represent , respectively. Ensuring uniformity across $n_{i}$ for all $i$, we establish $D-1$ patch sizes $p_{i}=2^{D-i+1}$, where $i$ ranges from 1 to $D-1$, implying that $p=\{2^{4},2^{3},2^{2}\}$ and the smallest patch size is $2^{2}=4$, given $D=4$. The multiscale feature extraction is illustrated in Figure 1 2.

Patch Embedding involves transforming vectorized patches $\hat{z}_{j}^{p_{i}}\in Z_{i}$ into a latent space of $d_{i}$ dimensions using a trainable linear projection. To preserve the spatial information of the patches, we incorporate position embedding specific to each patch, which are then combined with the patch embeddings.

where $E_{i}$ is the convolution operation to perform patch embbeding on $X_{i}$ and produce $\hat{Z}^{i}$, while $E_{\text{pos}}^{i}\in(B,n,d^{i}_{f})$ denotes the position embedding, $Z_{i}$ is the linear embedding projection after adding vectors $\hat{z}_{j}^{p_{i}}\in(B,1,d^{i}_{f})$ with positional vectors $e_{j}^{i}\in(B,1,d_{f}^{i})$. The linear embedding and positional embedding is displayed in Figure 1 2.

The Transformer encoder consists of $N$ layers of Multihead Self-Attention (MSA) and Multi-Layer Perceptron (MLP) blocks. Therefore the output of the $l$-th $\in N$ layer can be written as follows:

where $LN(\dot{)}$ denotes the layer normalization operator and $Z_{i}^{l}$ is the encoded image representation at scale $i$. The structure of Transformer layer is illustrated in Figure 1 2. In each layer $l$-th, the encoded image representation $Z_{i}$ undergo a self-attention mechanism, enabling encoded patches to learn how to attend to each other. Mathematically, the attention scores $A_{i}=\text{Attention}(Q_{i},K_{i},V_{i})$ for $Z_{i}$ are computed as follows:

where $Q_{i}=W_{q}(Z_{i}),K_{i}=W_{k}(Z_{i}),V_{i}=W_{v}(Z_{i})$ and $Q_{i},K_{i},V_{i}\in(B,n,d_{f}^{i})$. Additionally, the Multilayer Perceptron contains a fully connected layer of size $d_{i}\times 4$ in the middle. The resulting $E_{i}$ maintains the same shape as $Z_{i}$, which learns the intraslice information or the relationship between patches in one 2D image slice.

This technique is typically used in natural language processing, namely text sentiment analysis [rot, general_attention] where there three inputs involved. The phrase for which the sentiment needs to be determined (target phrase), the text before target phrase (left context), text after target phrase (right context). This greedy method assumes that adjacent phrases would contribute the most to the current center/targer phrase. In our context, if we denote the current encoded input representation as $Z_{i}\in(B,n,d_{f}^{i})$, we can separate this into multiple images $\{Z_{i}^{1},\dots,Z_{i}^{k},\dots,Z_{i}^{B}\}$. Therefore, three consecutive encoded slices/images can be selected as $\{Z^{k-1}_{i},Z^{K}_{i},Z^{K+1}_{i}\}$ or $\{Z^{l},Z^{t},Z^{r}\}$ to follow the left-target-right manner. For simple mathematical representation, we temporally disregard the scale $i$. These 3 encoded images are represented as:

The idea is to extract a single vector $r\in d_{f}$ and add this vector to $Z^{t}$ to adjust the hidden states or transform the position of each embedded patch $z_{j}^{t}$ in the semantic dimensional space by referring to the information from adjacent slices. In detail, we need to represent $Z^{t}$ as a single vector $r^{t}$ and incorporate necessary information from left and right context by attention mechanism to avoid noise and redundant information. Firstly, a single target representation is created by using pooling layer that takes the average over rows of $Z^{t}$:

Then similar to self-attention mechanism in Transformer layers, the key and value matrices are extracted from left context:

The $r_{t}$ now is used as a query to create the context vector out of left context. The scores are calculated with activated general score function with tanh activation function and the attention scores are calculated with softmax function:

A weighted combination of patch embedding is considered as the component representation for left contexts:

In Figure 2 3, we denote the above process as Single Attention (SA), which is represented as:

The vector $r^{l}$ is then used as query to create context out of target context to integrate information back into the center encoded slice/image to produce $r^{l/r}=SA(Z^{t},r^{l})$. An analogous procedure can be performed to obtain the right-aware target representation $r^{r}=SA(Z^{r},r^{t})$ and $r^{r/t}=SA(Z^{t},r^{r})$. Finally, to obtain the full representation vector $r$, we perform concatenation: $r^{k}=\text{concat}([{r}^{l},{r}^{r},{r}^{l/t},{r}^{r/t}])$ with $r^{k}\in\mathbb{R}^{1\times d_{f}\times 4}$. This $r$ vector contains the aggregated information between 3 consecutive slices, thus we have $B-2$ vectors $r^{k}$ with $1<k<B$ where $B$ is batch size since we perform the dense rotatory attention as illustrated in Figure 2 3. The final vector $R$ is achieved as: $R=W_{r}(\text{mean}(r^{k}|1<k<B))$. But this is only one $i$-th level output, thus we have $R_{i}$ output. This interslice-informational vector is added to encoded intraslice-informational $E_{i}$ to retrieve more optimized vectorized patch embeddings $F_{i}$.

In order to better fuse features of inconsistent semantics between the Channel Transformer and U-Net decoder, we propose a channel-wise cross attention module, which can guide the channel and information filtration of the Transformer features and eliminate the ambiguity with the decoder features.

Mathematically, we take the $i$-th level output $F_{i}$ after Transformer and Rotatory blocks to reconstruct or decode the encoded image representations to get $O_{i}\in\mathbb{R}^{C\times H\times W}$. The reconstructed $O_{i}$ are taken with $i$-th level decoder feature map ${D_{i}}\in\mathbb{R}^{C\times H\times W}$ as the inputs of Channel-wise Cross Attention.

Spatial squeeze is performed by a global average pooling (GAP) layer, producing vector $\mathcal{G}({X})\in\mathbb{R}^{C\times 1\times 1}$ with its $k^{th}$ channel $\mathcal{G}({X})=\frac{1}{H\times W}\sum_{i=1}^{H}\sum_{j=1}^{W}{X}^{k}(i,j)$. We use this operation to embed the global spatial information and then generate the attention mask:

where ${L}_{1}\in\mathbb{R}^{C\times C}$ and ${L}_{2}\in\mathbb{R}^{C\times C}$ and being weights of two Linear layers and the ReLU operator $\delta(\cdot)$. This operation encodes the channel-wise dependencies. Followed ECA-Net [eca] which empirically showed avoiding dimensionality reduction is important for learning channel attention, we use a single Linear layer and sigmoid function to build the channel attention map. The resultant vector is used to recalibrate or excite ${O}_{i}$ to ${\hat{O}}_{i}=\sigma({M}_{i})\cdot{O_{i}}$, where the activation $\sigma({M}_{i})$ indicates the importance of channels. Finally, the masked ${\hat{O}}_{i}$ is concatenated with the up-sampled features of the $i$-th level decoder.

In our experimental phase, we delved into both binary segmentation and multi-class segmentation tasks across a diverse range of datasets divided into two types: one abdominal dataset and four cardiac datasets. Here’s a detailed breakdown of the datasets used:

We used NVIDIA RTX 4090 1X GPU with 24GB memory, 81.4 TFLOPS for the training process. For our experiments, we utilized the NVIDIA RTX 4090 1X GPU, with 24GB of memory, 81.4 TFLOPS for our training tasks. We implemented our network RotCAtt-TransUNet++ with 8 different networks: TransUNet, Swin-unet, Attention Swin-UNet, UNet, UNet Attention, UNet++, UNet++ Attention, ResUNet. Across 5 diverse datasets, we evaluated the performance of 9 different networks using essential metrics such as Dice Coefficient Score (DSC), Intersection over Union (IoU) scores, and Hausdorff Distance. For a detailed class-wise analysis, we provided supplementary class-wise DSC and class-wise IoU scores. Nine networks were implemented using PyTorch, employing a fixed configuration for patch size $p_{i}$ and embedding dimension $d^{i}_{f}$, where $i$ signifies distinct feature map scales. Specifically, we utilized $p=[16,8,4]$ and $d_{f}=[64,128,256]$. Consequently, for input image sizes of $128,256,512$, we had token counts of $64,256,1024$ respectively. Additionally, we saved the matrices of self-attention weights, context weights, and rotatory attention vectors denoted as $A,C,R$ respectively, for visualization and ablation study purposes. We consciously avoided employing any data augmentation techniques to maintain the synthetic nature of our data and to prevent the introduction of extraneous artifacts that could potentially bias performance comparisons between models.For optimization, we chose the Stochastic Gradient Descent (SGD) optimizer with an initial learning rate set at $0.01$ and a weight decay of $0.0001$. However, our code implementation also provides an option for the Adam optimizer.

We employed a 3 loss functions: Cross Entropy Loss, Dice Loss, and IoU Loss, leveraging the combined loss of Dice Coefficient (DSC) and Intersection over Union (IoU) for efficient backpropagation. The mathematical formulations are as follows:

Cross Entropy Loss quantifies the disparity between the predicted probability distribution ($P_{ij}$) and the ground truth labels ($G_{ij}$). It calculates the average negative logarithm of the predicted probabilities assigned to the correct classes. This loss function is commonly employed in classification tasks to guide the model towards minimizing classification errors.

Dice Loss measures the dissimilarity between the predicted segmentation ($P$) and the ground truth ($G$) by computing the Dice coefficient. It assesses the overlap between the two sets, emphasizing regions of agreement while penalizing inconsistencies. This loss is particularly effective in scenarios where class imbalances exist, as it provides a robust measure of segmentation accuracy.

IoU Loss, or Intersection over Union Loss, evaluates the spatial overlap between the predicted and ground truth segmentation masks. It quantifies the ratio of the intersection area to the union area of the two sets, providing a comprehensive measure of segmentation accuracy. By penalizing deviations from ideal overlap, IoU Loss guides the model towards producing segmentation maps that closely align with ground truth annotations.

Here, $P$ and $G$ represent the predicted segmentation map and ground truth respectively, while $c$ denotes the class. The exclusion of $c\neq 0$ ensures the avoidance of unreal DSC and IoU scores stemming from dominant background pixels. Our composite loss function is defined as: