Unbiased Shape Compactness for Segmentation

TL;DR

This paper introduces an unbiased, dimensionless shape compactness prior for segmentation, optimized efficiently with ADMM, and demonstrates its effectiveness in medical image segmentation tasks like abdominal aorta MRI segmentation.

Contribution

The paper presents a novel shape compactness prior that is unbiased and position-independent, along with an efficient ADMM-based optimization method for high-order segmentation problems.

Findings

Competitive performance on abdominal aorta MRI segmentation

Efficient optimization of high-order energy with dense graphs

Effective integration with CNN-based classifiers

Abstract

We propose to constrain segmentation functionals with a dimensionless, unbiased and position-independent shape compactness prior, which we solve efficiently with an alternating direction method of multipliers (ADMM). Involving a squared sum of pairwise potentials, our prior results in a challenging high-order optimization problem, which involves dense (fully connected) graphs. We split the problem into a sequence of easier sub-problems, each performed efficiently at each iteration: (i) a sparse-matrix inversion based on Woodbury identity, (ii) a closed-form solution of a cubic equation and (iii) a graph-cut update of a sub-modular pairwise sub-problem with a sparse graph. We deploy our prior in an energy minimization, in conjunction with a supervised classifier term based on CNNs and standard regularization constraints. We demonstrate the usefulness of our energy in several medical applications. In particular, we report comprehensive evaluations of our fully automated algorithm over 40 subjects, showing a competitive performance for the challenging task of abdominal aorta segmentation in MRI.

Figures

Tables

Table 1. Table 1: Quantitative evaluations of MR-T1 abdominal aorta segmentations.

Method	CNN [7]	CNN + GC [5]	TRIC [1]	Proposed
Dice	0.73 (0.09)	0.76 (0.09)	0.71 (0.07)	0.81 (0.08)
Time	$\sim$ 60 sec	$\sim$ 70 sec	$\sim$ 20 min	$\sim$ 80 sec

Equations

$$E(\mathbf{y})=E_p(\mathbf{y})+\lambda E_c(\mathbf{y}),$$

$$\operatorname*{arg\,min}_{\mathbf{y}\,\in\,\{0,1\}^{|\Omega|}}\ E(\mathbf{y})\ =\ {\mathbf{u}}^{\top}\mathbf{y}\,+\,\lambda\frac{\big{(}{\mathbf{y}}^{\top}L\mathbf{y}\big{)}^{2}}{{\mathbf{1}}^{\top}\mathbf{y}},$$

$$\operatorname*{arg\,min}_{\mathbf{y},\,\mathbf{z},\,s}\ {\mathbf{u}}^{\top}\mathbf{y}\,+\,\frac{\lambda}{s}\big{(}{\mathbf{y}}^{\top}L\mathbf{y}\big{)}\big{(}{\mathbf{z}}^{\top}L\mathbf{z}\big{)},\quad\textrm{s.t.}\ \mathbf{y}=\mathbf{z}\,\textrm{ and }\,s={\mathbf{1}}^{\top}\mathbf{z}.$$

$$\operatorname*{arg\,min}_{\begin{subarray}{c}\mathbf{y},\,\mathbf{z},\,s\\ \boldsymbol{\nu}_{1},\,\nu_{2}\end{subarray}}\ {\mathbf{u}}^{\top}\mathbf{y}\,+\,\frac{\lambda}{s}\big{(}{\mathbf{y}}^{\top}L\mathbf{y}\big{)}\big{(}{\mathbf{z}}^{\top}L\mathbf{z}\big{)}\,+\,\frac{\mu_{1}}{2}\big{\|}\mathbf{y}-\mathbf{z}+\boldsymbol{\nu}_{1}\big{\|}_{2}^{2}\,+\,\frac{\mu_{2}}{2}\big{(}s-{\mathbf{1}}^{\top}\mathbf{z}+\nu_{2}\big{)}^{2}.$$

$$\operatorname*{arg\,min}_{\mathbf{z}\,\in\,\mathbb{R}^{|\Omega|}}\ \alpha{\mathbf{z}}^{\top}L\mathbf{z}\,+\,\frac{\mu_{1}}{2}\big{\|}\mathbf{z}-(\mathbf{y}+\boldsymbol{\nu}_{1})\big{\|}_{2}^{2}\,+\,\frac{\mu_{2}}{2}\big{(}{\mathbf{1}}^{\top}\mathbf{z}-(s+\nu_{2})\big{)}^{2}.$$

$$\mathbf{z}\ =\ \big{(}\alpha L+\mu_{1}I+\mu_{2}\mathbf{1}{\mathbf{1}}^{\top}\big{)}^{-1}\big{(}\mu_{1}(\mathbf{y}+\boldsymbol{\nu}_{1})\ +\ \mu_{2}(s+\nu_{2})\mathbf{1}\big{)}.$$

$$\displaystyle\big{(}\alpha L+\mu_{1}I+\mu_{2}\mathbf{1}{\mathbf{1}}^{\top}\big{)}^{-1}$$

$$\displaystyle\,=\,Q^{-1}\,-\,\frac{1}{\mu_{1}}{\Big{(}\frac{\mu_{1}}{\mu_{2}}+|\Omega|\Big{)}}^{-1}\mathbf{1}{\mathbf{1}}^{\top}.$$

$$\operatorname*{arg\,min}_{s\geq 0}\ \frac{\beta}{s}\,+\,\frac{\mu_{2}}{2}\big{(}s-({\mathbf{1}}^{\top}\mathbf{z}-\nu_{2})\big{)}^{2}$$

$$s^{3}\,-\,\big{(}{\mathbf{1}}^{\top}\mathbf{z}-\nu_{2}\big{)}s^{2}\,=\,\frac{\beta}{\mu_{2}}\,>\,0.$$

$$\operatorname*{arg\,min}_{\mathbf{y}\,\in\,\{0,1\}^{|\Omega|}}\ {\mathbf{u}}^{\top}\mathbf{y}\,+\,\lambda{\mathbf{y}}^{\top}L\mathbf{y}\ +\ \frac{\mu_{1}}{2}\big{\|}\mathbf{y}-\mathbf{q}\big{\|}_{2}^{2}$$

$$\operatorname*{arg\,min}_{\mathbf{y}\,\in\,\{0,1\}^{|\Omega|}}\ \sum_{i\in\Omega}\big{(}u_{i}+\mu_{1}(\tfrac{1}{2}-q_{i})\big{)}y_{i}\,+\,\lambda\sum_{i,j\in\Omega^{2}}w_{ij}\big{|}y_{i}-y_{j}\big{|}.$$

Full text

1 Introduction

Several recent studies have shown that generic shape constraints such as convexity [10, 17], compactness [1, 6], axial symmetry [15], tubularity [13], skeleton consistency [12] and inter-region topology [2] can be very powerful in medical image segmentation. Such constraints can boost substantially the performances of state-of-the-art segmentation algorithms, including powerful supervised learning methods such as convolutional neural networks (CNN) [2]. Imposing these constraints can be beneficial in a breadth of medical applications, particularly when training data is limited or when the target segments undergo strong noise and poor contrasts/resolutions (see Fig. 1 for examples where target boundaries have varied shapes and low contrast). Unfortunately, these shape constraints are typically high-order functionals, which yield challenging optimization problems. For instance, the recent convexity term in [10] involves a large number of non-submodular triple-cliques, and the compactness in [1] is a high-order ratio, both requiring computationally expensive approximations and iterative schemes to reach a local minimum.

This study focuses on constraining segmentation functionals with shape compactness, a problem investigated previously in several works [1, 6, 14, 11]. In segmentation, the most common compactness functional is the ratio of boundary length to area (or surface to volume in 3D) [14, 11], which is related to the well-known isoperimetric graph partitioning problem (also referred to as the Cheeger problem). Unlike common length regularizers [4], this prior does not suffer from strong shrinking bias (i.e., bias towards small regions). Despite its simplicity, optimizing a discrete form of this ratio functional is an NP-hard problem [14]. However, it is possible to transform this problem into a family of parameterized sub-modular pairwise costs, which can be explored efficiently with parametric max-flow optimizers [14]. Unfortunately, length-to-area ratio varies significantly with region size111This is due to the fact that, in image segmentation, region size might be several orders of magnitude larger then boundary length.. As pointed out (and observed experimentally) in several studies [11, 14], the isoperimetric partitioning problem has a strong bias towards large regions, which is a serious limitation in practice.

The recent segmentation study in [1] proposed an alternative compactness measure based on a ratio of regional moment statistics, thereby ensuring size (or scale) invariance. This algorithm showed promising performances in the case of a vascular structure, and can handle arbitrary variations in region size. Unfortunately, this compactness term is pose-dependent, requiring the user to specify the center of mass of the target region in each 2D slice of a subject data set. This precludes its use for fully automated segmentation, particularly in the case of large 3D volumes. More importantly, it is strongly biased towards circular shapes, significantly limiting its applicability (see the examples in Fig. 1). Another related work is the semi-automatic segmentation algorithm in [6], whose compactness prior also favors large regions and requires user seeds.

The contributions of this study can be summarized as follows.

• •

General-purpose segmentation: We propose to constrain segmentation functionals with a dimensionless, unbiased and position-independent shape compactness prior, which we solve efficiently with an alternating direction method of multipliers (ADMM). To the best of our knowledge, the ratio of length-squared to area has not been used before in segmentation functionals, although it is widely accepted as a compactness measure in the area of shape metrics [16]. This is probably due to the fact that the ensuing high-order optimization problem is challenging. Involving a squared sum of pairwise potentials, our problem makes the use of powerful sub-modular optimization techniques (e.g., parametric max-flow) impractical for image segmentation problems, as it would result in dense (fully connected) graphs. We propose an efficient alternative by splitting the problem into a sequence of easier sub-problems. At each iteration, our ADMM solution alternates three steps, each can be performed efficiently: (i) a sparse-matrix inversion based on Woodbury identity, (ii) a closed-form solution of a cubic equation and (iii) a graph-cut update of a sub-modular pairwise sub-problem involving a sparse graph.

• •

Clinical applications: We first report a comprehensive evaluation of our algorithm on the challenging task of abdominal aorta segmentation in 3D MRI. Results obtained on a dataset of 40 subjects demonstrate a state-of-the-art performance of the proposed method, with an average Dice metric of 0.81. Except for the recent semi-automatic algorithms described in [8, 1], most existing aorta segmentation techniques have focused on CTA. Unlike these techniques, our algorithm allows the development of fully automated tools, which could be used retrospectively on large-scale MRI datasets to identify potentially fatal abnormalities like abdominal aortic aneurysms (AAA) [9]. Our experiments further show the usefulness of our method in three important and challenging segmentation problems: the delineation of the esophagus in CT, the right ventricle in MR, and the aorta in CT.

2 Formulation

Let $\Omega$ denote the image domain and $\mathbf{x}_{i}\in\mathbb{R}^{K}$ the input-feature vector of pixel (or voxel) $i\in\Omega$. Segmentation assigns to each $i$ a label $y_{i}\in\mathcal{L}$ ($\mathcal{L}$ is the set of possible labels). While our formulation could be easily extended to multi-class segmentation, we focus on a binary (two-region) statement for a clearer presentation. We minimize the following functional with respect to a discrete binary vector $\mathbf{y}\in\{0,1\}^{|\Omega|}$:

[TABLE]

where $E_{p}$ encodes unary-potential priors, and $E_{c}$ is our shape compactness. As is typical in segmentation, $E_{p}$ can be expressed in term of log-likelihoods: $E_{p}(\mathbf{y})\,=\,\sum_{i}u_{i}\,y_{i}$, where $u_{i}=\log p(y_{i}=0\,|\,\mathbf{x}_{i})-\log p(y_{i}=1\,|\,\mathbf{x}_{i})$. Although any (semi-) supervised technique can be used to learn $p(y_{i}\,|\,\mathbf{x}_{i})$, in this study, we employed a fully-convolutional neural network (FCNN) [7].

Our dimensionless, unbiased and position-independent shape compactness prior is based on a measure that is well accepted in the context of shape metrics [16], the ratio of length-squared to area, which we can write in discrete form (i.e., as a function of our segmentation variable): $E_{c}(\mathbf{y})=P(\mathbf{y})^{2}/A(\mathbf{y})$. The area can be expressed as $A(\mathbf{y})=\sum_{i}y_{i}$, and the length is proportional to the number of neighboring pixels with different labels, i.e., $P(\mathbf{y})\propto\sum_{i,j}w_{ij}(y_{i}-y_{j})^{2}$, where pairwise potential $w_{ij}=1$ if $i$ and $j$ are neighbors, otherwise $w_{ij}=0$. Note that this compactness measure can be easily extended so as to attract the solution towards strong edges in the image: $w_{ij}=\exp(-\sum_{k}\sigma_{k}\,(x_{ik}-x_{jk})^{2})$, with $\sigma_{k}$ controlling the relative importance of feature $k$ on the weight. With these definitions of $E_{p}$ and $E_{c}$, and re-writing length as $P(\mathbf{y})={\mathbf{y}}^{\top}L\mathbf{y}$, where $L$ is the Laplacian matrix corresponding to weights $w_{ij}$, our compactness model becomes:

[TABLE]

where $\mathbf{1}$ is a vector with the value one for each element.

2.1 ADMM optimization

The general principle of ADMM is to decompose a hard problem into easier-to-solve sub-problems, which are coupled together via equality constraints. In our case, we introduce auxiliary variable $\mathbf{z}\in\mathbb{R}^{|\Omega|}$ and $s\in\mathbb{R}_{+}$, and reformulate problem (2) as

[TABLE]

Note that since we impose $\mathbf{y}=\mathbf{z}$, we can relax the binarity constraints on $\mathbf{z}$. This strategy allows an efficient update of this variable.

In the next step, we move the constraints in the cost-function via a augmented Lagrange formulation:

[TABLE]

In this new formulation, $\boldsymbol{\nu}_{1}$ and $\nu_{2}$ are dual variables corresponding to the two constraints, and $\mu_{1}$, $\mu_{2}$ parameters controlling the trade-off between the primal objective and constraint satisfaction. The benefit of this formulation is that solving for each variable, considering all other variables fixed, can be carried out efficiently and to optimality. We thus adopt an iterative optimization method, where each variable is updated in turn until convergence.

Updating $\mathbf{z}$:

Let $\alpha=\frac{\lambda}{s}\,{\mathbf{y}}^{\top}L\mathbf{y}$, the task of updating $\mathbf{z}$ can be expressed from (4) as

[TABLE]

Minimizing this convex quadratic problem yields

[TABLE]

Note that computing $\mathbf{z}$ requires solving a large and dense linear system. Let $Q=\alpha L+\mu_{1}I$, following the Woodbury identity and using the fact that $Q^{-1}\mathbf{1}=\frac{1}{\mu_{1}}\mathbf{1}$, we can reformulate the matrix inversion as:

[TABLE]

Since $Q$ is very sparse, the resulting system can be solved efficiently via standard techniques like the preconditioned conjugate gradients method.

Updating $s$:

Let $\beta=\lambda({\mathbf{y}}^{\top}L\mathbf{y})({\mathbf{z}}^{\top}L\mathbf{z})$, updating $s$ amount to solving

[TABLE]

Deriving this equation w.r.t. $s$ and setting the result to zero gives the following cubic equation:

[TABLE]

It can be shown that this equation has a real valued root $s>{\mathbf{1}}^{\top}\mathbf{z}-\nu_{2}$. In practice, ${\mathbf{1}}^{\top}\mathbf{z}-\nu_{2}$ is positive and converges toward zero, and $s$ will be positive.

Updating $\mathbf{y}$:

Let $\gamma=\frac{\lambda}{s}{\mathbf{z}}^{\top}L\mathbf{z}$ and $\mathbf{q}=\mathbf{z}-\boldsymbol{\nu}_{1}$. We update $\mathbf{y}$ by considering the following problem:

[TABLE]

Since $\mathbf{y}$ is binary, we have that $y_{i}^{2}=y_{i}$ and $(y_{i}-y_{j})^{2}=|y_{i}-y_{j}|$. Hence, we can reformulate (10) as

[TABLE]

This corresponds to a simple graph-cut problem, which can be solved efficiently using the Boykov-Kolmogorov algorithm [5].

Finally, we update dual variables following the standard ADMM algorithm: $\boldsymbol{\nu}^{\prime}_{1}=\boldsymbol{\nu}_{1}+(\mathbf{y}-\mathbf{z})$ and $\nu^{\prime}_{2}=\nu_{2}+(s-{\mathbf{1}}^{\top}\mathbf{z})$. This iterative updating process is repeated until $\|\mathbf{y}-\mathbf{z}\|_{2}$ is less than a small epsilon.

3 Experiments

We first present a quantitative evaluation of our shape compactness on the task of segmenting the abdominal aorta in a MR-T1 data set of 40 subjects. The usefulness of the proposed method is then shown qualitatively for three additional segmentation scenarios: CT esophagus, MRI right ventricle and CT aorta segmentation. Our method’s accuracy is measured in terms of Dice coefficient, and compared to that of graph cuts (GC) [5] and TRIC222Note that TRIC needs a reference point to define the shortest-path distance. This point was defined as the centroid of the CNN-based segmentation in this work. compactness [1]. To measure the contribution of the pixelwise (i.e., unary potential) prior, we also report results obtained via a simple thresholding.

While any technique can be used to obtain the unary-potentials in Eq. (2), we considered the 3D fully-CNN (FCNN) architecture presented in [7]. This architecture is comprised of 9 convolutional layers with non-linear activation units, 3 fully-connected layers (converted into standard convolution operations), and a soft-max layer. The output of this model is a map indicating the probability of each pixel to belong to a given structure. The same unary potential was used for all three tested approaches. For all segmentation applications, the FCNN was trained using a k-fold validation strategy, with $k=4$.

For our method’s ADMM parameters, we used $\mu_{1}\!\approx\!2000$ and $\mu_{2}$=50 for all experiments. As mentioned in [3], ADMM algorithms are not overly sensitive to these parameters. To facilitate convergence, we increase the value of these parameters by 1% (i.e., 1.01 multiplication factor) at each iteration. In contrast, the compactness regularization $\lambda$ was tuned for each segmentation problem: $\lambda$ = 5000 for MR-T1 abdominal aorta, $\lambda$ = 1000 for CT esophagus, $\lambda$ = 15000 for MR right ventricle, and $\lambda$ = 3000 for CT abdominal aorta. The code has been made publicly available at https://github.com/josedolz/UnbiasedShapeCompactness.

MR-T1 abdominal aorta segmentation: As shown in the first row of Fig. 1, segmenting the aorta from MR-T1 images is a very challenging problem due to the noise and lack of visible boundaries. In this example, thresholding the FCNN probabilities yields parts of the background along with irregular contours. Because it regularizes the segmentation, GC provides a smoother contour. However, this contour follows the probability map of FCNN and does not reflect the compact shape of the target region. Increasing regularization weight in GC yields a compact region but decreases performance significantly due to the well-known shrinkage bias of standard pairwise length regularization (bias to small regions); note that regularization weight was tuned for an optimal GC performance. The next experimental examples will further highlight shrinkage bias. Unlike GC, TRIC finds a more compact region, due to its shape circularity prior. Nonetheless, this example illustrates this approach’s two main drawbacks: 1) it is biased to near-circular regions, and 2) it is not translation invariant and finds regions centered on the probability map’s center of mass. As seen in the figure, our pose-independent method did not suffer from a circularity/shrinkage bias, yielding a much more accurate segmentation, highly similar to the reference contours.

The bar plot of Fig. 2 gives the accuracy of tested methods obtained for each subject333Two subjects were excluded from the quantitative analysis, as the FCNN failed to generate a usable probability map.. The mean Dice and standard deviation computed across all subjects, is reported in Table 1. These results confirm that the proposed compactness term provides a considerable improvement in accuracy. In terms of runtime, our method takes about 80 seconds on average to segment a single subject, 60 seconds of time spend computing the unary potentials via the FCNN. In contrast, TRIC requires nearly 20 minutes per subject. Hence, our method is more suitable for the fully automated segmentation of large-scale datasets.

CT esophagus segmentation: Fig. 1 (second row) shows delineations of the esophagus in CT, also a challenging problem due to its complex shape and non-homogeneous appearance. In this particular example, we can observe how the well-known shrinkage bias problem of GC leads to an under-segmentation of the esophagus. Notice that GC yielded a compact region but did not improve FCNN’s performance. As in the previous case, we also observe how TRIC over-enforces shape circularity on the segmentation. On the other hand, our compactness term is able to preserve the target region’s shape.

MR Right Ventricle Segmentation: This qualitative result (Fig. 1, third row) shows how our compactness term can accommodate a more general class of shapes that differs significantly from tubular structures. Although TRIC [1] can also handle shapes different than a circle, multiple reference points, which form a skeleton, are required from the user in this case, a prohibitively time-consuming effort for 3D data. However, to keep the process fully automatic, only the centroid of the FCNN segmentation was provided. Having only a single reference point, and particularly in this scenario, led to failure of TRIC to achieve a satisfactory segmentation.

CT Aorta Segmentation: The application of the proposed method to CT aorta segmentation is illustrated in the last row of Fig 1. We see that our method can handle multi-region scenarios, i.e. bifurcations, in contrast to other compactness terms [1]. Note that, for this example, we applied TRIC separately on each of the two regions to avoid having a single contour centered in between these regions. This means that we fed TRIC with additional supervision information, unlike our method which is fully automatic. Notice that, to obtain compact and smooth regions, GC worsened FCNN performance due to shrinkage bias.

4 Conclusion

We presented an unbiased, fully-invariant and multi-region prior for the segmentation of compact shapes, based on the ratio of length-squared to area. An efficient ADMM strategy was proposed to solve the high-order energy minimization problem resulting from this formulation. Using an FCNN to obtain unary probabilities, the proposed method achieved high accuracy in four challenging segmentation problems.