Multi-class medical image segmentation using one-vs-rest graph cuts and majority voting

Abstract.

Purpose: Semi-automatic image segmentation is still a valuable tool in clinical applications since it retains the expert oversights legally required. However, semi-automatic methods for simultaneous multi-class segmentation are difficult to be clinically implemented due to the complexity of underlining algorithms. We purpose an efficient one-vs-rest graph cut approach of which the complexity only grows linearly as the number of classes increases.

Approach: Given an image slice, we construct multiple one-vs-rest graphs, each for a tissue class, for inference of a conditional random field (CRF). The one-vs-rest graph cut is to minimize the CRF energy derived from regional and boundary class probabilities estimated from random forests to obtain a one-vs-rest segmentation. The final segmentation is obtained by fusing from those one-vs-rest segmentations based on majority voting. We compare our method to a well-used multi-class graph cut method, alpha-beta swap, and a fully connected CRF (FCCRF) method, in brain tumor segmentation of 20 high-grade tumor cases in 2013 MICCAI dataset.

Results: Our method achieved mean Dice score of 0.83 for whole tumor, compared to 0.80 by alpha-beta swap and 0.79 by FCCRF. There was a performance improvement over alpha-beta swap by a factor of five.

Conclusions: Our method utilizes the probabilistic-based CRF which can be estimated from any machine learning technique. Comparing to traditional multi-class graph cut, the purposed one-vs-rest approach has complexity that grows only linearly as the number of classes increases, therefore, our method can be applicable for both online semi-automatic and offline automatic segmentation in clinical applications.

1. Introduction

Many clinical applications utilizing medical images require careful and thorough labeling of anatomical structures or tissue types to design an optimal treatment plan, or to evaluate or predict treatment outcomes. Human oversight is legally and ethically necessary for decision-making in the healthcare industry that is highly regulated. Thus, segmentation done by expert is still the gold standard. However, manual segmentation is prone to human error and introduces variations existing between physicians with different training backgrounds and within the same physician from time to time. This inconsistency may imperil successful patient treatment. Automatic or semi-automatic segmentation approaches can diminish human fatigue, accelerate workflow, and provide consistency.

Most commercial solutions employ atlas-based approaches^1–5 to automatic normal organ segmentation. Atlas-based methods assume that organ shape and location are similar among the population. Thus, once a spatial correspondence is established between an atlas (or multiple atlases) and a target image, the atlas’ pre-segmented labels can be transferred to the target image. However, when the anatomy of a new case is very different from the selected atlases, the image registration between the target image and the atlases might not be sufficiently accurate to obtain a good segmentation.⁶ To address variations of organ shapes among the population, one may introduce more atlases to increase the coverage, but as the number of atlases increases, it becomes more computationally expensive to find the proper atlases as more registrations are required.

With the advanced chip development of graphics processing units (GPUs), the neural network has taken a quantum leap in recent years in computer vision and natural language processing applications. Convolutional neural networks (CNNs)⁷ have shown remarkable performance in object recognition and have been extended to medical image segmentation^8,9 with equally impressive capability. A network needs to go deep to achieve robustness, i.e., it includes many layers of neural nodes and convolution filters. As a result, it requires an enormous amount of training data to fit a massive number of network parameters. For real-world images, much effort, such as ImageNett,¹⁰ has been made publicly available to provide very large-scale curated and labeled image datasets for the training of neural networks. While it is possible to acquire annotations of those real-world images via crowdsourcing, it is impossible to do so in the field of medical images. Obtaining a large-scale medical image training dataset remains a challenge. First, well-curated data from medical expert annotation are demanding. Second, the training dataset is usually institution-specific. A dataset from one institution may not apply to another institution. Lastly, certain disease sites and certain image modalities may be lacking sufficient images.

We propose a novel and efficient one-vs-rest graph-based approach for multi-class and multi-modal medical image segmentation. We do not aim to compete with state-of-the-art fully automatic methods that require tremendous effort in training. Instead, we examined the performance of the proposed graph-based segmentation method with limited on-line training and showed compatible results comparing to state-of-the-arts.

2. Background of Graph-Based Segmentation

One of the approaches to image segmentation is to treat a discrete image grid as a graph and utilize widely studied graph partition algorithms. We handle multi-class image segmentation as a minimum $K$-cut problem in graph theory to partition the image (graph) into $K$ sub graphs (classes) such that the cost to separate the graph is minimum. The cost commonly includes a data term to penalize disagreement between the assigned image label of a pixel and its observation, plus a smoothness term to penalize discontinuity if neighboring image pixels have similar features. The cost is formulated as an energy function of labeling $f$ with $K$ labels:

$$E(f)=\sum _{p\in P}{D}_p(f_p)+\sum _{p,q\in Nbr}{V}_{p,q}(f_p,f_q),$$ (1)

where $D$ and $V$ are unary data and pairwise smoothness terms, respectively, $P$ is the set of image pixels and $Nbr$ is the set of neighboring pixels. The energy function resembles the energy from conditional random field (CRF)¹¹ which models the conditional probability of labeling (segmentation) given the observation. The conditional probability is maximized when the energy is minimized. It is well known that the $K$-cut problem of finding the minimum of the energy function is NP-hard when $K>=3$ so it is computationally intractable. Boykov et al.¹² demonstrated that the minimum of the energy function in Eq. (1) can be approximated via a sequence of two-class $s-t$ cuts, namely, $\alpha -\beta$ swap and $\alpha$-expansion algorithms for $V$ being a semi-metric and metric, respectively. Both algorithms use $s-t$ cuts iteratively to change the labels of pixels for approaching the minimum of energy until convergence. At each iteration, the $\alpha -\beta$ swap algorithm considers each pair of labels to determine if pixels assigned to one label in the pair can be swapped to another label ($swap$-move) using graph cut, while $\alpha$ expansion algorithm iterates through each label to determine if pixels assigned to the label can be expanded, that is, allowing any set of pixels to change their label to the examined label ($expansion$-move.) Although these algorithms are more efficient than traditional approximation method such as simulated annealing, they are still very time consuming due to the iterative process, thus not suitable for on-line segmentation.

Recently, the graph structure of CRF model has been extended to a fully connected graph [fully connected CRF (FCCRF)]^13,14 to establish pairwise potentials on all pairs of pixels in the image. The resulting graph may have billions of edges, making the conventional inference algorithms impractical. To efficiently solve the minimization of CRF energy for inference, mean-field approximation was employed, nevertheless, the edge (boundary) potentials were limited to Gaussian kernels.

Our method, although maintaining traditional CRF graph structure, does not make assumptions on boundary potentials, and addresses the inference problem using the concept of ensemble classification.¹⁵ We perform $K$ one-vs-rest graph cuts, each to minimize CRF energy of a one-vs-rest permutation, to obtain $K$ segmentations. The final segmentation is decided by label fusion based on majority voting from the $K$ segmentations. Our method only increases the complexity of $s-t$ graph cut by a factor of $O(K)$. The overview of our method is shown in Fig. 1.

Fig. 1.

The overview of the proposed segmentation workflow. $K$ one-vs-rest graphs are constructed with edge costs assigned from probabilities estimated by trained regional and boundary classifiers. Individual $s-t$ cut is performed for each graph and the segmentations from cuts are combined through a majority vote to obtain the final segmentation.

3. Methods

3.1. Conditional Random Fields

Let $G=(V,E)$ be the graph representing a 2D or 3D volumetric image with $N$ voxels. $V=\{i|1\le i\le N\}$ are nodes representing image voxels and $E=\{\{i,j\}|i\in V,j\in Nb{r}_i\}$ are edges connecting neighboring nodes, where $Nb{r}_i$ is the set of neighbors of voxel $i$ in a neighborhood system (four-connected system in this study.) We treat image segmentation as a classification problem for assigning a class label to each voxel. Let $L=\{\mathrm{1,2},\dots ,K\}$ be the set of $K$ labels for the classification, the label assignment be a random variable $Y={(Y_i)}_{i\in V,Y_i\in L}$ and observed image features be $X={(X_i)}_{i\in V}$. The most plausible segmentation $\widehat{y}$ given an image $x$ is when conditional probability $P(Y=\widehat{y}|X=x)$ is maximum. In a CRF, $P(Y|X)$ factorizes according to $G$ as defined in the following form

$$P(y|x)\propto \exp (-\sum _{i\in V}{r}_i(y_i,x)+\lambda \sum _{\{i,j\}\in E}{u}_{ij}(y_i,y_j,x)),$$ (2)

where $r_i$ is a unary term describing how likely that the node $i$ is assigned to label $y_i$ given the observed image feature $x$ and $u_{ij}$ is a pairwise term for an edge $\{i,j\}$ describing how likely that nodes $i$ and $j$ are assigned to labels $y_i$ and $y_j$, respectively. We define the two terms to reflect the probabilities estimated from a regional classifier for classifying voxels to certain tissue types and a boundary classifier for classifying an edge to be a non-boundary or boundary type. Let $B=\{(\alpha ,\beta )|\alpha \ne \beta ,\alpha ,\beta \in L\}$ be the set of boundary types. We define

$$r_i(y_i,x)=-\ln p(y_i|x_i),$$ (3)

and

$$u_{ij}(y_i,y_j,x)=\{\begin{array}{ll}0,& \text{if}{y}_i=y_j\\ -\ln p(y_i,y_j|x_i,x_j),& \text{if}(y_i,y_j)\in B\end{array}.$$ (4)

3.2. One-vs-Rest $s-t$ Cuts

The exponent in Eq. (2) can be minimized by graph cut¹⁶ via $\alpha -\beta$ swap or $\alpha$ expansion. Since multi-cut ($K$-cut) problem is NP hard, those approximation methods are expensive to compute. We propose a one-vs-rest approach that combines the results of $Ks-t$ cuts. Since multiple labels can be assigned to a voxel after $Ks-t$ cuts, the final class label assignment is determined in a fashion similar to label fusion strategy commonly seen in atlas-based segmentation. Let $\alpha -R_{\alpha }$ cut be an $s-t$ cut that determines the voxels to be assigned to class $\alpha$ or the rest of the classes, $R_{\alpha }=L\setminus \{\alpha \}$. The construction of graphs for $s-t$ cut is similar to that used in Boykov and Jolly¹⁷ for bi-class segmentation. For each voxel node, $t-link$ edges are added to connect the $source$ to the node and the node to the $\sin k$, and $n-link$ edges are added to connect the node to its neighbors. For a node $i$ and its neighbor $j$, let the cost of the edges being cut be $t_{(source,i)}$ and $t_{(i,sink)}$ for $t-link$ edges and $n_{(i,j)}$ for $n-link$. They are assigned as follows for an $\alpha -R_{\alpha }s-t$ cut that minimizes CRF energy defined in Eq. (2),

$$t_{(source,i)}:\min \{r_i(\beta ,x)|\beta \in R_{\alpha }\}{t}_{(i,sink)}:r_i(\alpha ,x)n_{(i,j)}:\min \{u_{ij}(\alpha ,\beta ,x)|\beta \in R_{\alpha }\}$$

Once the edge costs are assigned, then a regular minimum $s-t$ cut is performed. After the cut, any voxel (node) that remains connected to source is assigned to class label $\alpha$ for the $\alpha -R_{\alpha }$ cut at this time and otherwise it is assigned to $R_{\alpha }$. If a link with $n_{(i,j)}$ is cut, then there is a boundary $(\alpha ,\beta )$ for an unknown $\beta$ in $R_{\alpha }$.

3.3. Majority Votes

After $K\alpha -R_{\alpha }$ cuts, each voxel has $K$ class assignments from the cuts, the final class assignment, i.e., segmentation, is determined by the majority votes. Let $C_k(i)$ be the class assignment for voxel $i$ from the $k$’th cut, then we define a counting function $T$

$$T_{\alpha }(i)=\sum _{k=1}^K\delta (C_k(i),\alpha ),$$ (5)

and a $\delta$ function

$$\delta (C_k(i),\alpha )=\{\begin{array}{ll}1,& \text{if}{C}_k(i)=\alpha \text{or}{C}_k(i)=R_{\gamma },\gamma \ne \alpha \\ 0,& \text{otherwise}\end{array}.$$ (6)

$T_{\alpha }(i)$ is the number of times of that voxel $i$ is assigned to $\alpha$. We then determine the final class assignment for voxel $i$ using the following rules:

• •Majority rule: If $T_{\alpha }(i)>T_{\beta }(i)$, for all $\beta \ne \alpha$, assign voxel $i$ to class $\alpha$. If there is no majority, tie breaker rule is applied.
• •Tie breaker rule: If there is a subset $L^{\prime }\subset L$ such that $T_{\alpha }(i)=T_{\beta }(i)$, where $\alpha ,\beta \in L^{\prime }$, assign voxel $i$ to class $\widehat{\alpha }$, classified from a regional classifier. That is, $\widehat{\alpha }={\max }_{\alpha \in L^{\prime }}(p(\alpha |x_i))$

To summarize the procedure of one-vs-rest cuts, the algorithm is listed in Algorithm 1.

Algorithm 1.

Algorithm for One-vs-Rest Cuts

Input:$G=\{V,E\}$, $L$, $M_R$ (regional classifier), $M_B$ (boundary classifier)

Output:$S$ (segmentation)

1: for$k=1$ to $|L|$do

2:  initialize class assignment $C_k$

3:  $\alpha =L(k)$

Construct$\alpha -R_{\alpha }$graph:

4:  $G^{\prime }=G$

5:  for$i$ in $V$do

6:   add edge (source, $i$) and assign edge cost $t_{(source,i)}$ estimated from $M_R$

7:   add edge ($i$, sink) and assign edge cost $t_{(i,sink)}$ estimated from $M_R$

8:  end for

9:  for$(i,j)$ in $E$do

10:   assign edge cost $n_{(i,j)}$ estimated from $M_B$

11:  end for

12: perform $s-t$ cut on $G^{\prime }$

13:  for$i$ in $V$do

14:   if$i$ remains connected to $source$then

15:    $C_k(i)=\alpha$

16:   else

17:    $C_k(i)=R_{\alpha }$

18:   end if

19:  end for

20: end for

Majority votes:

21: for$i$ in $V$do

22:  for$\alpha$ in $L$do

23:   Calculate $T_{\alpha }$ from $C_k(i)$ using Eq. (5)

24:  end for

25:  $l_{\max }$ = the set of labels that have $\max \{T_{\alpha }\}$

26:   if$|l_{\max }|=1$then

27:    $S(i)=l_{\max }$

28:   else

29:    $S(i)$ = the label predicted from $M_R$

30:  end if

31: end for

32: return$S$

4. Experiments and Results

4.1. Dataset

We evaluated our method with 20 MRI high-grade (HG) cases from MICCAI Brain Multi-Modal Tumor Segmentation Challenge (BRATS 2013).¹⁸ Four MRI sequences: Flair, T1, T1 contrast (T1c), and T2, as well as ground truth are provided in the dataset. The ground truths are consensual manual segmentations from experts and contain the following five classes of tissue type: 0- Normal organs, 1- Necrosis, 2- Edema, 3- Non-Enhancing tumor, and 4- Enhancing Tumor. In this experiment, the evaluation was done for Whole tumor (1 + 2 + 3 + 4), Active tumor (1 + 3 + 4) and Edema (2), considering applications in clinical practice.

4.2. Training

To demonstrate the performance of our method with limited training samples available during on-line segmentation, we performed case-specific training. For each case, we used 3 to 5 image slices for training: an axial slice in the middle of the tumor, an axial slice in the superior part of the tumor, an axial slice in the inferior part of the tumor, and when necessary, two additional axial slices to ensure that training samples are available for all classes. Depending on the case, tumors appear on 32 to 98 image slices. Image intensities from all four MR sequences were the only image feature used to train regional and boundary classifiers to estimate the probabilistic terms in the CRF energy. Intensity values are raw data directly from image files without any pre-processing. We chose random forest¹⁹ (RF) for both regional and boundary classifiers in this study. Voxel and boundary samples of all five labels from the training slices were used to train the regional and boundary RF classifiers. It was possible, however, that samples from a boundary class were lacking, i.e., two types of tissues may not have attached to each other at all. In this case, a very small epsilon value ($1e-6$) was assigned to the probability.

4.3. Results

We did voxel-wise classification using the regional RF classifiers to generate baseline segmentations for evaluation of our one-vs-rest method and the classic multi-class graph cut method, $\alpha -\beta$ Swap. In addition, we compared our method to the most recent state-of-art FCCRF. Our method and $\alpha -\beta$ Swap shared the same RF regional and boundary classifiers for CRF energy whereas, for FCCRF, the RF regional classifier was used to establish unary potentials and Gaussian kernels-based on image intensities and positions of paired voxels were used to establish boundary potentials for the mean-field approximation as proposed by Krahenbuhl et al.¹³

We used Dice similarity coefficient (DSC) as the measure of accuracy for comparison to the ground truth from the MICAAI dataset. The average scores of different methods are shown in Table 1 and Fig. 2. Our method one-vs-rest performed best in all three tissue categories: whole tumor, active tumor, and edema with mean DSC in 0.83, 0.79, and 0.71, respectively. Our method was significantly better than all three other methods in the segmentation of whole tumor ($p<0.002$). Our method also was significantly better ($p<0.05$) than RF and FCCRF in both active tumor and edema segmentations. For those two tasks, when compared to the $\alpha -\beta$ swap method, our method was not significantly better. This was because of the larger variation of Dice score in the $\alpha -\beta$ swap method. It was commonly understood that the Dice coefficient is too sensitive for small structures and tends to have larger variation. Figure 3 shows the segmentation results of five example cases from the MICCAI 2013 HG cases. The baseline RF segmentation is the result of voxel-wise classification and therefore had many isolated mis-segmented voxels. FCCRF removed most of the mis-segmented voxels, however, did not preserve piece-wise continuity of the segmentation. The $\alpha -\beta$ swap preserved the continuity well but in the process some fine details were lost. In contrast, the segmentation from our one-vs-rest method shows advantages in both continuity and fine details. These were quantified using the following approach. We found the connected components in 3D (26-connected neighborhood) and calculated the percentage of number of components that has volume less than 26 voxels (isolated components either from over-segmentation or fine details.) For whole tumor, the ground truth is 92%, compared to 98% in RF, 98% in FCCRF, 67% in $\alpha -\beta$ swap, and 97% in our method (Fig. 4). The numbers confirm the significant under-segmentation in $\alpha -\beta$ swap.

Table 1.

Comparison of DSC of segmentations from our one-vs-rest method and other methods.

	Whole tumor		Active tumor		Edema
	Mean	Std.	Mean	Std.	Mean	Std.
RF	0.67	0.14	0.65	0.19	0.56	0.15
FCCRF	0.79	0.10	0.74	0.15	0.68	0.14
$\alpha -\beta$ swap	0.80	0.18	0.77	0.21	0.69	0.22
One-vs-rest	0.83*	0.11	0.79**	0.14	0.71**	0.16

Note: Wilcoxon signed rank test.

^*$p<0.05$, compared to RF and FCCRF

^**$p<0.002$, compared to all other methods

Fig. 2.

Dice scores. Our one-vs-rest method has the highest mean in all three tissue categories: whole tumor (0.83), active tumor (0.79), and edema (0.71).

Fig. 3.

Segmentation results from five HG cases. Each row is an example slice from a different case. (a) T2 weighted MRI image; (b) results using only the regional classifier; (c) results from FCCRF; (d) results from $\alpha -\beta$ swap; (e) results of our proposed one-vs-rest method; and (f) the ground truth.

Fig. 4.

Comparison of percentage of isolated components in whole tumor segmentation. Our one-vs-rest method is closest to the ground truth while $\alpha -\beta$ swap has significant under-segmentation.).

4.4. Performance

The implementation of the algorithms in this study is in Python with a runtime library implemented in C/C++ for a fast version of $s-t$ cut¹⁷ used in both $\alpha -\beta$ swap and our one-vs-rest method. On average, the segmentation time per slice is 82 s for $\alpha -\beta$ swap and 18 s for one-vs-rest on a Windows^® 7 workstation equipped with an Intel^® Xeon^® E2-2620 2GHz processor.

5. Conclusion

We present a novel one-vs-rest graph cut algorithm under the CRF framework for multi-class image segmentation. Our algorithm performs multiple minimum $s-t$ cuts to minimize CRF energy for obtaining $K$ segmentations of the $K$ one-vs-rest class permutations. The final segmentation is combined from the $K$ one-vs-rest segmentations via majority votes. Unlike the previous graph-based approach $\alpha -\beta$ swap, the complexity of our method grows only linearly with the number of classes, in addition to the complexity of the $s-t$ algorithm chosen. Our CRF unary (regional) and pairwise (boundary) potentials are purely probabilistic and can be estimated from any chosen and trained classification model. We demonstrated that our algorithm is suitable for semiautomatic segmentation with light-weight and case-specific training, which can be done interactively. In our experiment with MICCAI brain tumor dataset, our method yielded the highest Dice scores in segmentation of whole tumor, active tumor, and edema.

Future work will investigate other image features to further improve segmentation results and feasibility of using this method as a post-editing tool for deep learning-based segmentation.

Biographies

Yu-chi Hu is a software specialist at Memorial Sloan Kettering Cancer Center. His research interests are focused on developing and incorporating machine learning techniques in image segmentation, image registration, and longitudinal analysis for applications in radiotherapy.

Gikas Mageras is an emeritus medical physicist at Memorial Sloan Kettering Cancer Center. His research interests are in detection and management of patient internal motion by means of image-guided radiation therapy. His current research is in the application of deep learning to deformable CT-to-cone-beam-CT registration for localizing organs-at-risk during radiation treatment of pancreatic cancer.

Michael Grossberg is an associate professor of computer science at City College of New York. His research spans data science, computer science, statistics, and mathematics. This has focused on data visualization for big and high dimensional data, using tools from statistical/machine learning, and using web technology for exploration. Recent domains have included remote sensing, climate science, medical imaging, computer vision, and finance. He is also codirector of the CCNY Master’s Program in Data Science and Engineering.

Disclosures

No conflicts of interest, financial, or otherwise, are declared by the authors.