Combining UNet 3+ and Transformer for Left Ventricle Segmentation via Signed Distance and Focal Loss

Abstract

Left ventricle (LV) segmentation of cardiac magnetic resonance (MR) images is essential for evaluating cardiac function parameters and diagnosing cardiovascular diseases (CVDs). Accurate LV segmentation remains a challenge because of the large differences in cardiac structures in different research subjects. In this work, a network based on an encoder–decoder architecture for automatic LV segmentation of short-axis cardiac MR images is proposed. It combines UNet 3+ and Transformer to jointly predict the segmentation masks and signed distance maps (SDM). UNet 3+ can extract coarse-grained semantics and fine-grained details from full scales, while a Transformer is used to extract global features from cardiac MR images. It solves the problem of low segmentation accuracy caused by blurred LV edge information. Meanwhile, the SDM provides a shape-aware representation for segmentation. The performance of the proposed network is validated on the 2018 MICCAI Left Ventricle Segmentation Challenge dataset. The five-fold cross-validation evaluation was performed on 145 clinical subjects, and the average dice metric, Jaccard coefficient, accuracy, and positive predictive value reached 0.908, 0.834, 0.979, and 0.903, respectively, showing a better performance than that of other mainstream ones.

1. Introduction

The World Health Organization (WHO) showed that in 2019, almost 17.9 million people died of cardiovascular diseases (CVDs), accounting for 32% of fatalities worldwide [1]. Early diagnosis of CVD can help improve cardiac function and reduce patient mortality [2]. Cardiovascular magnetic resonance (MR) imaging is harmless and has become the most commonly used technique for evaluating cardiovascular system structure and function. Left ventricle (LV) segmentation, as a key step in the treatment of CVD, can provide better visual aid during the diagnosis of CVD. Most CVDs affect the physiological shape of the cardiac LV, and LV dysfunction is the cause of many heart diseases, such as ventricular hypertrophy and myocardial infarction, making the examination of the LV an important prerequisite for determining whether the heart is diseased. LV segmentation can accurately delineate the boundaries of the LV on cardiac MR images so that physicians can better understand some clinical parameters, such as the patient’s ventricular volume, ejection fraction, LV mass, and stroke volume [3,4].

In early clinical work, medical images are usually annotated by experts to mitigate the subjective bias caused by the level of a particular expert or possible negligence of subtle symptoms [5]. However, for most professional clinicians, manual segmentation is a cumbersome and time-consuming task. In general, it takes a clinician about 20 min to segment a patient’s cardiac MR image. Additionally, in Figure 1, as you can see from the left ventricle’s structure, the intensity and shape similarity between the LV and other organs, boundary inaccuracies, and the inherent noise of cardiac MR imaging have all posed obstacles to LV segmentation [6,7,8].

Some excellent segmentation algorithms spawned in the past decades have been caught in certain dilemmas that make them difficult to apply in clinical settings. Among them, many of the algorithms are traditional methods based on machine learning [9,10,11], such as thresholding [12], clustering [13], active contours [14,15], and split–merge. Figure 2 displays the thresholding, clustering, and active contours methods. Most of them use semi-automatic methods that heavily rely on the initialization step [16,17,18,19], resulting in their failure to achieve the desired results. Meanwhile, because of the rapid development of the accessibility of vast training data and computer hardware, medical image segmentation algorithms based on deep learning are becoming increasingly prevalent. In particular, Convolutional Neural Networks (CNNs) have achieved excellent results in various computer vision (CV) fields, such as image segmentation [20], object detection [21], and image classification [22]. In such a trend, CNN-based LV segmentation models such as Densely Connected Convolutional Networks (DenseNet) [23] and Fully Convolutional Neural Networks (FCN) [24] have been suggested, and have achieved good results in clinical trials.

Although these networks are representative, CNN-based approaches still have some limitations in capturing global information and recovering weak texture details. To overcome this limitation, some studies have suggested Transformer-based designs that use the self-attention mechanism to construct the contextual representations. Transformers, unlike traditional CNNs, can model global contextual information by modeling the relationships between spatially distant pixels [25]. However, both models have their drawbacks. CNN is insensitive to global features, and Transformer has the issues of high computation cost and a lack of ability to capture regional features, so the applicability of both models needs to be further improved.

To solve the above problems, this study proposes a fast and automated method of cardiac LV segmentation to help facilitate the diagnosis of CVD. The network has the following two main contributions: (1) The backbone uses a combined framework of UNet 3+ [26] and Transformer, which can efficiently acquire low-level spatial features as well as high-level semantic information and can also model the global context. (2) The shared backbone network is used to jointly predict the segmentation masks and signed distance maps (SDM) to study segmentation targets’ different representations from different perspectives.

1.1. Traditional Segmentation Methods

Currently, the methods for LV segmentation can usually be divided into traditional segmentation methods and deep learning methods. However, the early traditional segmentation methods have some obvious drawbacks. Several of these algorithms can only obtain relatively accurate segmentation results when the pixel intensities of the LV and other tissues reach a high level of contrast. For example, the threshold-based segmentation method was used by Goshtasby et al. [27] in 1995 to extract the LV contour. It adaptively determines the threshold grayscale according to the global or local grayscale histogram of the image. Therefore, this method can only obtain better segmentation results on the condition that the grayscale of the region is significantly different from the background. Since its appearance, the K-Means clustering algorithm has been extensively applied in the fields of image analysis and data mining. However, because many regions in cardiac MR images are similar or even connected to the LV, the K-means clustering algorithm cannot achieve the expected results and needs further improvement. In 2006, Katouzian et al. [12] employed the idea of split–merge to segment the LV, which could successfully extract the epicardium and endocardium of the LV under the condition of manually annotating the epicardium and endocardium of the first slice. This manual process obviously increases the clinical application complexity. On the whole, these traditional methods have the characteristic of relying on manual design, which contradicts the idea of automatic segmentation.

1.2. Deep Learning

Unlike traditional segmentation methods, deep learning methods prefer to train with large-scale data to find the intrinsic patterns and representation levels in images and get more representative feature information [28]. They naturally describe image features without relying on the manual extraction of features, addressing the limitations of traditional segmentation methods. J. Long et al. [24] proposed a fully convolutional neural network (FCN) based on a CNN to recover the feature map to the original image size via transposed convolution, thus achieving image segmentation in 2015. P. V. Tran [7] segmented both the left and right ventricles in 2016 using FCN. However, the recovery of the LV contour was poor due to only a single upsampling. Ronneberger et al. [29] suggested the U-Net network based on FCN. U-Net effectively integrates low-resolution information and high-resolution information to learn better feature representation and improve generalization performance. However, the utilization of feature maps is poor and does not work well for object boundary segmentation. SegNet [30] proposed a decoder to perform nonlinear upsampling, the input feature maps of which were the maximum pooling index received from the corresponding encoder. In spite of avoiding the learning of upsampling and improving the precision of image boundary localization via this method, the segmentation accuracy is not high enough to satisfy the real-time requirement. On the other hand, despite the excellent representation capability of all these networks, CNN-based approaches have difficulty learning global semantic information because of the inherent locality of convolution operations. As a result, these networks usually yield weak segmentation performance, particularly for target structures that exhibit large inter-patient variations in size, shape, and texture.

1.3. Transformers

Transformer was first proposed by Vaswani et al. [31] as the main method for machine translation. It has now been introduced as a new model for image recognition [32], semantic segmentation [33], and many other computer vision tasks [34,35]. In contrast with previous CNN-based approaches, Transformer is powerful at modeling global contexts and demonstrates superior transferability for downstream tasks under large-scale pretraining. However, Transformer concentrates on modeling the global context at all stages, leading to generating low-resolution features. Further, due to the lack of detailed localization information for these low-resolution features, which cannot be efficiently recovered by direct upsampling to full resolution, the ultimate obtained segmentation results are coarse. In order to make the model focus on both regional and global features in segmentation tasks, recent studies have tended to combine CNN and Transformer, such as TransUNet [36] and TransFuse [37], which yield satisfactory performance in segmentation tasks. These works show that the combination model has great potential in the field of CV.

2. Method

The proposed network in this paper aims to automatically segment the LV in cardiac MR images, thus reducing the tedious manual segmentation and improving the disease diagnosis efficiency. The new medical image segmentation framework proposed in this work is shown in Figure 3, and the backbone part uses a combined architecture of UNet 3+ and Transformer. The network takes the same cardiac MR images as input and predicts both pixel probability maps and SDM. A loss function composed of two main components is designed to train the segmentation network. One is for computing the pixel probability map and the other one is for computing the SDM.

2.1. Segmentation Network

Feature Extraction—First, feature maps are generated for the input images using the encoder structure of UNet 3+ as a feature extractor. We are given an image $x\in {\mathbb{R}}^{H\times W\times C}$ with the spatial resolutions H, W and C. The preprocessed image x is fed into the network and the encoder applies a series of convolutional blocks to model the pixel-level contextual representations, with features progressively downsampled to $\frac{H}{16}\times \frac{W}{16}$.

Transformer—To extract the global features, the Transformer module is applied in the encoder design. Because Transformer is not as efficient as UNet 3+ when it comes to capturing regional features, patch embedding is directly used for patches generated from the UNet 3+ feature maps rather than from raw images. Then patches $x_p$ are mapped into a K-dimensional embedding space by a trainable linear projection. To preserve location information, specific position embeddings $E_{pos}$ are added to the patch embeddings when encoding the patch spatial information. Details are as follows:

$$z_0=\left[x_p^{1}E;x_p^{2}E;\cdots ;x_p^{N}E\right]+E_{pos},$$

(1)

where $E\in {\mathbb{R}}^{\left(P^2·C\right)\times K}$ is the patch embedding projection, $E_{pos}\in {\mathbb{R}}^{N\times K}$ denotes the position embedding, and $x_p^1,\cdots ,x_p^N$ are image patches.

Then, a stack of Transformer blocks consisting of alternating layers of multi-head self-attention (MSA) and multi-layer perceptron (MLP) blocks are used to learn the long-range context representation. Layer normalization (LN) is used before each block, and residual connections are used after each block. The following can be expressed as the output of the i-th layer:

$$z_i^{{}^{\prime }}=MSA\left(LN\left(z_{i-1}\right)\right)+z_{i-1},$$

(2)

$$z_i=MLP\left(LN\left(z_i^{{}^{\prime }}\right)\right)+z_i^{{}^{\prime }},$$

(3)

where $LN\left(·\right)$ represents the layer normalization operator, $i\in \left\{1,2,\cdots ,L\right\}$ where L is the number of Transformer layers and $z_i$ denotes the image representation of the encoded output of the i-layer Transformer.

Decoder—To generate segmentation masks and SDM in the raw image space, a decoder for UNet 3+ is introduced to perform feature upsampling. Since the output of the Transformer is sequential data, its output should first be recovered to spatial order. The encoded feature representation $z_i\in {\mathbb{R}}^{N\times K}$ is reshaped into $z_i\in {\mathbb{R}}^{\frac{H}{P}\times \frac{W}{P}\times K}$ where p is the size of each patch, and then the channel dimension is reduced by a $3\times 3$ convolution block. In addition, each decoder layer in the decoder combines the feature maps of all encoders. Additionally, the full-scale deep supervision proposed by UNet 3+ is used to learn hierarchical representations from the full-scale aggregated feature maps, and the output from each decoder stage is supervised by the ground truth (GT). To achieve deep supervision, a $3\times 3$ convolution block, a bilinear upsampling, and a sigmoid function are added to the last layer of each decoder stage in the network. To generate the SDM, the network adds an SDM head at the last decoder stage, which consists of a convolution block and tanh activation.

2.2. Loss Function

Based on the above architecture design, the segmentation network generates pixel-level segmentation maps and SDM, and the following functions are introduced to convert GT to SDM. The SDM assigns each pixel a value, indicating its signed distance to the nearest boundary of the target object.

$$\mathcal{D}\left(a\right)=\left\{\begin{array}{cc}0,& a\in S\\ \underset{b\in \mathcal{C}}{-inf}{∥a-b∥}_2,& a\in {\mathcal{C}}_{in}\\ \underset{b\in \mathcal{C}}{+inf}{∥a-b∥}_2,& a\in {\mathcal{C}}_{out}\end{array}\right.$$

(4)

where $∥a-b∥$ is the Euclidean distance between pixels a and b, S represents the boundary of the target object, and ${\mathcal{C}}_{in}$ and ${\mathcal{C}}_{out}$ represent the region inside and outside of the target object, respectively. Typically, SDM takes negative values inside the target and positive values outside the target, with the absolute value indicating the distance from the point to the nearest point on the target object’s surface.

In the network training, for the regression task branch, a ${\mathcal{L}}_2$ loss is used between the SDM of the network output $P_a$ and the transformed GT map $\mathcal{D}\left(Y\right)$.

$${\mathcal{L}}_{sdm}\left(P_a,Y\right)={∥P_a-D\left(Y\right)∥}^2,$$

(5)

where Y denote the GT map.

For the segmentation task branch, the combination ${\mathcal{L}}_{seg}$ of dice loss and focal loss is applied as the loss function between the segmentation mask and the GT for each decoder output, and then their average is taken as the final segmentation loss.

$${\mathcal{L}}_{seg}\left(P_b,Y\right)={\mathcal{L}}_{Dice}\left(P_b,Y\right)+{\mathcal{L}}_{FL}\left(P_b,Y\right),$$

(6)

where ${\mathcal{L}}_{Dice}$ denotes the dice loss and ${\mathcal{L}}_{FL}$ denotes the focal loss. $P_b$ and Y denote the prediction partition map and label, respectively.

The final loss is defined as:

$$\mathcal{L}={\mathcal{L}}_{sdm}\left(P_a,Y\right)+{\mathcal{L}}_{seg}\left(P_b,Y\right).$$

(7)

3. Experiments

3.1. Datasets

The dataset from MICCAI 2018 is used to train and evaluate the proposed model [38]. The dataset contains 2900 short-axis cardiac MR images from 145 subjects at three hospitals attached to two healthcare centers (London Healthcare Center and St. Joseph’s Healthcare). The age range of the study subjects is 16 to 97, with a mean of 58.9. In the 1.5625 mm/pixel mode, the pixel spacing of the MR images ranges from 0.6836 mm/pixel to 2.0833 mm/pixel. Pathologies such as myocardial hypertrophy, LV dysfunction, atrial septal defect, regional wall motion abnormalities, mildly enlarged LV, etc., are present. During the entire cardiac cycle, twenty frames are acquired for each subject. According to the standard AHA prescription, in each frame, the LV is divided into equal thirds (basal, medial, and parietal) perpendicular to the long axis of the heart. Before manual annotating GT, all cardiac MR images need to be performed with landmark labeling, rotation, ROI cropping, and resizing. After preprocessing, all images are cropped and resized to the dimension of $80\times 80$ and normalized. After manual contouring, two experienced cardiac radiologists (A. Islam and M. Bhaduri) obtain the epicardial and endocardial boundary and perform a double examination. The labels of the data are approved by industry doctors.

3.2. Implementation Details

In the experiments, the backbone of the main framework is the combination of UNet 3+ and Transformer. The network is implemented in PyTorch (1.11.0), with the runtime platform processor of Inter(R) Core (TM) i9-10850K CPU, NVIDIA GeForce RTX 3080 Ti. All training and test images are uniformly adjusted to $80\times 80$ dimensions. In the training stage, random rotation and flipping operations are applied as data augmentations. The model is evaluated and compared by the five-fold cross-validation. The dataset is split into five groups of 29 subjects each. One group (580 images) is selected for testing, and the remaining four groups (2320 images) are utilized as the training set. The final evaluation result is calculated using the average of five times this process. The proposed network is trained end-to-end with the Adam optimizer with a weight decay of 1 $\times {10}^{-5}$ and an initial learning rate of 2 $\times {10}^{-4}$. The model is trained in 100 epochs in a batch size of 16. The segmentation result utilized in testing is the output of the segmentation task branch.

3.3. Evaluation Metric

The goal of the network is making sure the model can accurately segment the LV from cardiac MR images. For objective evaluation of the proposed model, the region-based dice metric (DM) and Jaccard coefficient (JC) are employed as metrics. DM and JC are explained in detail as follows.

Dice Metric—DM calculates the overlap between the manual segmentation area and the automatic segmentation contour area obtained using the proposed method. DM lies in the [0, 1] range. The better the match between manual and predicted segmentation, the larger the DM value is. DM is defined as:

$$DM\left(A,B\right)=\frac{2\left|A\cap B\right|}{\left|A\right|+\left|B\right|},$$

(8)

where A and B represent the area of manual and automatic contour, respectively.

Jaccard Coefficient—The Jaccard coefficient, also named the Intersection over Union (IoU), is used to calculate the degree of dissimilarity between the manual segmentation area and the automatic segmentation contour area using the proposed method. Similar to DM, the JC is between [0, 1]. The larger the JC, the lower the similarity. The formula for JC is as follows:

$$J\left(A,B\right)=\frac{\left|A\cap B\right|}{\left|A\cup B\right|}=\frac{\left|A\cap B\right|}{\left|A\right|+\left|B\right|-\left|A\cap B\right|},$$

(9)

where A and B represent the area of manual and automatic contour, respectively.

In addition, we also use accuracy (ACC) and positive predictive value (PPV) to evaluate the results of pixel classification. They are defined by:

$$ACC=\frac{TP+TN}{TP+FN+FP+TN},$$

(10)

where $TP$, $TN$, $FP$, and $FN$ represent true positive, true negative, false positive, and false negative.

$$PPV=\frac{TP}{TP+FP},$$

(11)

where $TP$, $TN$, $FP$, and $FN$ represent true positive, true negative, false positive, and false negative.

4. Experimental Results

4.1. Performance of the Network

Figure 4 displays the predicted segmentation masks, GT, and contours for four subjects’ cardiac MR images from the MICCAI 2018. It is clear from Figure 4a–d that the proposed method can accurately segment LV in cardiac MR images. In Figure 4d, it can be found that the automated segmentation contours marked by the red curve almost overlap on the GT (marked by the green curve), indicating that the model can accurately segment LV in diverse shapes. In summary, the model shows great potential for cardiac MR image segmentation with high accuracy, and thus may provide a visual aid to clinicians for qualitative diagnosis.

4.2. Performance Comparison

On the MICCAI 2018 test set, comparison is made between the method and other prevalent segmentation methods such as FCN, Conv–Deconv [39], U-Net, Indices-JSQ [40], and SegNet to evaluate its effectiveness. As shown in

Table 1. The comparison of the suggested method with several popular segmentation methods.

Methods	DM	JC	ACC	PPV
FCN	0.873	0.778	0.972	0.878
SegNet	0.843	0.728	0.964	0.845
U-Net	0.887	0.800	0.975	0.887
Conv-Deconv	0.846	0.733	0.965	0.840
Indices-JSQ	0.870	− 1	−	−
Cardiac-DeepIED [41]	0.890	0.801	0.976	0.891
Our Method	0.908	0.834	0.979	0.903

¹—Indicates that the results are not provided.

, the DM and JC of the proposed method reach 0.908 and 0.834, respectively, with a certain extent of improvement compared to that of other segmentation methods. These results suggest that the model is capable of more accurately determining the class of each pixel point and then achieving higher segmentation performance.

4.3. Ablation Studies

In this section, ablation studies are performed on each component of the suggested method. As shown in

Table 2. Ablation experiments of the model on the MICCAI 2018 test set.

UNet 3+	Transformer	${\mathcal{L}}_{\mathbf{seg}}$	${\mathcal{L}}_{\mathbf{sdm}}$	DM
✓	×	✓	×	0.896
✓	✓	✓	×	0.907
✓	✓	✓	✓	0.908

, as the proposed modules are sequentially added on the UNet 3+ baseline, the model performance is gradually improved. It is clear from Table 2 that when Transformer is integrated into the UNet 3+ baseline, the DM reaches 0.907, with a 1.1% improvement compared to that of the baseline (0.896). This is because the incorporation of Transformer compensates for the inability of UNet 3+ to model the global context. Simultaneously, the combination also solves the problem of Transformer ignoring low-resolution detail features compared to using itself directly.

To further study the impact of the loss function in the model, an ablation study is conducted for the segmentation loss (${\mathcal{L}}_{seg}$) and SDM loss (${\mathcal{L}}_{sdm}$). By adding an additional regression head to the segmentation network’s end and using the SDM loss, the dice metric value of the model is further enhanced to 0.908, with a slight increase of 0.1% compared to that of the model only with ${\mathcal{L}}_{seg}$. This result suggests that joint SDM training can implicitly force the model to learn shape information compared to traditional training using only segmented masks.

5. Conclusions

In this study, a network for automatic LV segmentation from cardiac MR images is proposed, providing an effective solution that allows physicians to diagnose CVDs. The proposed method extensively experimented on cardiac MR image data from 145 subjects, and the DM and JC reached 0.908 and 0.834, respectively, on the MICCAI2018 test set. The proposed module and loss function both improve the segmentation accuracy, as verified by ablation experiments. The method also outperforms the current mainstream methods in the comparison experiments suggesting that it can be considered an effective automatic LV segmentation task that will reduce the workload of radiologists during clinical diagnosis. In future research, more possibilities for Transformer application to medical image segmentation networks will be explored to provide a better technique for LV segmentation.