VOTENET+ : AN IMPROVED DEEP LEARNING LABEL FUSION METHOD FOR MULTI-ATLAS SEGMENTATION Zhipeng Ding1, Xu Han1, Marc Niethammer1;2 1Department of Computer Science, UNC Chapel Hill, USA 2Biomedical Research Imaging Center, UNC Chapel Hill, USA ABSTRACT to use DL to predict trust-worthy atlases for MAS and to use DL-based image registration. This allowed matching DL per- In this work, we improve the performance of multi-atlas formance while retaining spatial consistency. segmentation (MAS) by integrating the recently proposed VoteNet model with the joint label fusion (JLF) approach. Related Work. There are various recent approaches combin- Specifically, we first illustrate that using a deep convolutional ing MAS with machine learning. E.g., [6] proposed to use neural network to predict atlas probabilities can better dis- random forests for patch-based label fusion. Further, [5] uses tinguish correct atlas labels from incorrect ones than relying a CNN to locally select the best atlases and [7] uses a CNN to on image intensity difference as is typical in JLF. Motivated predict patch-based similarity in the JLF framework. Most re- by this finding, we propose VoteNet+, an improved deep lated to our work are [5, 7]. However, our approach differs in network to locally predict the probability of an atlas label to the following ways: a) we design VoteNet+, a network which differ from the label of the target image. Furthermore, we improves over VoteNet in [5]; b) we use JLF for label fusion show that JLF is more suitable for the VoteNet framework while [5] use plurality voting; c) we predict the probability for as a label fusion method than plurality voting. Lastly, we an entire image instead of focusing on patch centers as in [7]; use Platt scaling to calibrate the probabilities of our new d) we propose to use Platt scaling to correct probabilities pre- model. Results on LPBA40 3D MR brain images show that dicted by our CNN while [7] use a heuristic to approximate our proposed method can achieve better performance than probabilities. Sec. 4 demonstrates that our proposed approach VoteNet. can indeed achieve improved performance. Contributions. (1) New network architecture: We propose Index Terms— multi-atlas segmentation, joint label fu- a new deep convolutional network (VoteNet+), which locally sion, VoteNet, probability calibration identifies sets of trustworthy atlases more accurately than VoteNet. (2) Probability Calibration: We calibrate the prob- 1. INTRODUCTION abilities of our network resulting in more accurate segmenta- tions. (3) We show the advantages of combining VoteNet and Image segmentation, i.e. assigning pixel-wise or voxel-wise Joint Label Fusion via Oracle experiments. (4) We further labels, is important for image-based diagnosis and analy- improve the final segmentation performance by combining sis [1]. Thus much effort has been spent on developing VoteNet+ with a U-Net-based segmentation network. fast and accurate segmentation algorithms. Recently, deep- learning (DL) approaches [2, 3, 4] have started to dominate arXiv:1911.00582v2 [cs.CV] 13 Jan 2020 2. METHODOLOGY the field of medical image analysis for many tasks, including image segmentation, due to their good performance. 2.1. MAS Overview Prior to the ascent of DL approaches, multi-atlas segmen- tation (MAS) techniques have been widely successful. MAS Assume we have n atlas images and their corresponding man- 1 1 1 2 2 2 n utilizes deformable image registration to transfer atlas labels ual segmentations: A = (AI ;AS);A = (AI ;AS); :::; A = n n to a target image to be segmented. As label maps are de- (AI ;AS). Let TI be the target image that needs to be seg- formed in a controlled manner via deformable registration ap- mented. MAS first uses a reliable deformable image regis- proaches local structure and topology stay well behaved in tration method to warp all atlases to the target image space, ~1 ~1 ~1 ~2 ~2 ~2 ~n ~n ~n the target image space, thus helping to retain spatial consis- i.e. A = (AI ; AS); A = (AI ; AS); :::; A = (AI ; AS); tency (e.g., to avoid adding unrealistic structures or missing and then uses a label fusion method, F , to combine all the structures). However, besides these desirable MAS behav- candidate segmentations to produce the final segmentation iors, MAS is slow as it relies on costly image registrations T^S for TI , i.e. and has been outperformed for many segmentation tasks by ^ ~1 ~2 ~n DL approaches. Recently, VoteNet [5] proposed an approach TS = F (A ; A ; :::; A ;TI ): (1) Fig. 1. VoteNet+ architecture. We use a customized attention U-Net with multi-scale input. Feature dimensions (channels) of each convolutional layer are listed on the top of each block. VoteNet+ takes the target image and a warped atlas image as inputs and outputs a voxel-wise probability indicating how likely the label from the warped atlas image is correct for the target image. 2.2. Joint Label Fusion is a key aspect of the VoteNet [5] framework. Still building on top of the 3D-Unet [2], we improved the network architecture Joint Label Fusion (JLF) [8] is a label fusion method taking (VoteNet+) to better predict such probabilities. We use an into account correlated errors of atlases. It models binary seg- image pyramid to provide structure details at different scales. mentations, but can be extended to multi-label settings. The This image pyramid is injected into the second and third following is a brief introduction. encoder blocks and concatenated with maxpooling features Binary segmentation errors can be modeled as from previous encoding blocks, as illustrated in Fig. 1(left). ~i i Further, inspired by [3, 4], we use customized soft atten- TS(x) = AS(x) + δ (x) (2) tion gates (AGs) to help identify where label mis-assignment ~i where TS(x) and AS(x) are the true target segmentation and might occur. AGs produce attention coefficients α 2 [0; 1] the ith atlas segmentation at position x; δi(x) is the label at each voxel to scale the input feature maps xl of layer l i difference. When TS(x) = 1, δ (x) 2 f0; 1g; otherwise to output salient features x^l. As illustrated in Fig. 1(bottom i when TS(x) = 0, δ (x) 2 {−1; 0g. The consensus segmen- right), a gating tensor g, which is used to determine focus ^ ^ Pn ~i tation TS is modeled as TS(x) = i=1 !i(x)AS(x), where regions, is first upsampled to the same shape as the features 0 ≤ !i(x) ≤ 1 is the weight assigned to the ith atlas and xl. Additive attention is then formulated as follows: Pn i=1 !i(x) = 1. JLF tries to find the optimal weights !i(x) T^ (x) minimizing the expected error between S and the true l T T l T q = σ1(A x + A g + bg) + b segmentation TS(x): att x g l l l h i α = σ2(qatt(x ; g;Θatt)) (6) E (T (x) − T^ (x))2 : (3) S S x^l = αl xl Ignoring that !i(x) 2 [0; 1] the JLF weights can be computed in closed-form as where σ2 is a sigmoid function; σ1 is a ReLU activation func- −1 Mx 1n tion; linear transformations Ax, Ag and are computed using wx = (4) t −1 channel-wise 1×1×1 convolutions; bg and b are bias terms; 1nMx 1n all the parameters in AG are represented as Θatt. where wx is the weight vector that contains weight of all at- lases. M 2 n×n is the dependency matrix, whose entries Let the network be P. Inputs are the target image TI and x R ~ are the joint probabilities of both the ith atlas (row) and the the warped atlas image AI ; the output of the network is the ~ jth atlas (column) producing the wrong label for the target probability p(AI = TS). Thus, to approximate the joint prob- image, i.e., ability in Sec. 2.2, we have i j Mx(i; j) = p(δ (x)δ (x) = 1): (5) i ~i ~i p(δ 6= 0) = 1 − p(AI = TS) = 1 − P(TI ; AI ); i j i j 2.3. Network Architecture p(δ δ = 1) ≈ p(δ 6= 0)p(δ 6= 0); (7) ~i ~j Designing a good network to predict the probability of = (1 − p(AI = TS))(1 − p(AI = TS)); whether an atlas locally has the same label as the target image ~i ~j = (1 − P(TI ; AI ))(1 − P(TI ; AI )): 2.4. Probability Calibration One-sided t score. In [8], the probability of the ith atlas hav- ing the wrong label at position x is defined as Probabilities predicted from deep convolutional neural net- works are often not well-calibrated [9]. We use Platt scal- 2 3β ing [10, 11] to calibrate the probabilities of our deep neu- i X ~i 2 p(δ (x) 6= 0) = 4 (TI (y) − AI (y)) 5 (8) ral network. Specifically, we fixed the learned parameters y2N (x) of VoteNet+ (Sec. 2.3) and optimized two scalar parameters R a; b 2 and output p^ = σ(az + b) as the calibrated proba- while in VoteNet, it is simply bility. Here σ is the sigmoid function, z is the output logit of i ~i VoteNet+ before going into the sigmoid function. Parameters p(δ (x) 6= 0) = 1 − P(TI (x); AI (x)): (9) a and b are optimized using the negative log likelihood (NLL) loss over the validation dataset of VoteNet+. To compare which approach is better at distinguishing cor- rect labels from incorrect ones we compute a one-sided t- score.