Learning Directional Feature Maps for Cardiac MRI Segmentation

We first detail the notation of the direction field. As shown in Fig. 3(a-b), for each foreground pixel $p$, we find its nearest pixel $b$ lying on the cardiac tissue boundary and then normalize the direction vector $\overrightarrow{bp}$ pointing from $b$ to $p$ by the distance between $b$ and $p$. We set the background pixels to $(0,0)$. Formally, the direction field $DF$ for each pixel $p$ in the image domain $\Omega$ is given by:

We propose a simple yet effective DF module to learn the above direction field, which is depicted in Fig. 3. This module is made up of a $1\times 1$ convolution, whose input is the 64-channel feature extracted by U-Net and output is the two-channel direction field. It is noteworthy that we can obtain the ground truth of the direction field from the annotation easily by distance transform algorithm.

The direction field predicted by the DF module reveals the directional relationship between pixels and provides a unique direction vector that points from the boundary to the central area for each pixel. Guided by these direction vectors, we propose a Feature Rectification and Fusion (FRF) module to utilize the characteristics of the central area to rectify the errors in the initial segmentation feature maps step by step. As illustrated in Fig. 4, the N-steps improved feature maps $F^{N}\in\mathbb{R}^{C\times H\times W}$ are obtained with the initial feature maps $F^{0}\in\mathbb{R}^{C\times H\times W}$ and the predicted direction field $DF\in\mathbb{R}^{2\times H\times W}$ step by step. Concretely, the improved feature of the pixel $p$ is updated iteratively by the feature of the position that $DF(p)$ points to, which is calculated by the bilinear interpolation. In other words, $F(p)$ is rectified by the features of the central area gradually. The whole procedure is formalized as below:

where $1\leq k\leq N$ denotes the current step, $N$ is the total steps (set to 5 if not stated otherwise), and $p_{x}$ (resp. $p_{y}$) represents the $x$ (resp. $y$) coordinate of the pixel $p$.

After performing the above rectification process, we concatenate $F^{N}$ with $F^{0}$, and then apply the final classifier on the concatenated feature maps to predict the final cardiac segmentation.

The proposed method involves loss function on the initial segmentation $L_{CE}^{i}$, final segmentation $L_{CE}^{f}$, and direction field $L_{DF}$. We adopt the general cross-entropy $L_{CE}$ as the segmentation loss to encourage class-seperate feature, which is commonly used in semantic segmentation. Formally, $L_{CE}$ is given by $L_{CE}=-\sum_{i}{y_{i}log(\hat{y}_{i})}$, where $y_{i}$ and $\hat{y}_{i}$ denote the ground truth and the prediction, respectively. For the loss to supervise the direction field learning, we choose the $L_{2}$-norm distance and angle distance as the training objective:

where $\hat{DF}$ and $DF$ denote the predicted direction field and the corresponding ground truth, respectively, $\alpha$ is a hyperparameter to balance the $L_{2}$-norm distance and angle distance, and is to 1 in all experiments, and $w(p)$ represents the weight on pixel $p$, which is calculated by:

where $|C_{i}|$ denotes the total number of pixels with label $i$ and $N_{cls}$ is the number of classes. The overall loss $L$ combines $L_{CE}$ and $L_{DF}$ with a balance factor $\lambda=1$:

The network is trained by minimizing the proposed loss function in Eq. (5) using ADAM optimizer [8] with the learning rate set to 10^-3. The network weights are initialized with [4] and trained for 200 epochs. Data augmentation is applied to prevent over-fitting including: 1) random translation with the maximum absolute fraction for horizontal and vertical translations both set to 0.125; 2) random rotation with the random angle between -180^∘ and 180^∘. The batch size is set to 32 with the resized $256\times 256$ inputs.

We first evaluate the proposed method on the ACDC dataset (train/val split described in Sec. 3.1) and the self-collected dataset. Tab. 1 presents the performance of the proposed DFM on two datasets. LV, RV and MYO represent the left ventricle, right ventricle and myocardium, respectively. Our approach consistently improves the baseline (U-Net), demonstrating its effectiveness. To further make comparison with the current state-of-the-arts methods, we submit the results to the ACDC leadboard. As shown in Tab. 2, compared with those methods that rely on well-designed networks (e.g. 3D network in [5] and Grid-like CNN in [15]) or multi-model fusion, the proposed DFM achieves competitive performance with only two simple yet effective modules added to the baseline U-Net. We also compare our approach with other methods dedicated to alleviate the inter-class indistinction and intra-class inconsistency. As depicted in Tab. 3, the proposed DFM significantly outperforms these methods. Some qualitative comparisons are given in Fig. 5.

To analyze the generalization ability of the proposed DFM, we performed a cross-dataset segmentation evaluation. Results listed in Tab. 4 show that the proposed DFM can consistently improve the cross-dataset performance compared with the original U-Net, validating its generalization ability and robustness. We also conduct ablation study on ACDC dataset to explore how the number of steps $N$ involved in the FRF module influences the performance. As shown in Tab. 5, the setting of $N=5$ gives the best performance.