Neural ODEs for Image Segmentation with Level Sets

1 INTRODUCTION . Image segmentation is the task of delineating pixels belonging to semantic labels . The ability to automatically segment objects is important because accurate labeling is expensive and hard ( Vittayakorn & Hays , 2011 ; Zhang et al. , 2018 ) . Automatic image segmentation can have large impact in many domains , e.g . obstacle avoidance in autonomous driving and treatment planning in medical imaging . Accurate classification of pixels in close proximity to inter-class boundaries remains a challenging task in image segmentation . Object boundaries can have high curvature contours or weak pixel intensity that complicate separating the object from surrounding ones . In deep CNNs ( Simonyan & Zisserman , 2014 ; Zeiler & Fergus , 2014 ; Szegedy et al. , 2015 ; He et al. , 2016 ; Chen et al. , 2017 ) , the object-of-interest and surrounding competing objects can provide equal context to a receptive field of a boundary pixel , which can make accurate classification difficult . Humans also find it difficult to accurately label pixels near object boundaries . Level Set methods ( Zhao et al. , 1996 ; Brox et al. , 2006 ) and Active Shapes ( Paragios & Deriche , 2000 ; Chan & Vese , 2001 ) have been proposed to incorporate shape and image priors to mitigate boundary ambiguities ( Tsai et al. , 2003 ; Rousson & Paragios , 2002 ) . The Level Set method for image segmentation evolves an initial contour of an object-of-interest along the normal direction with a forcing function . A contour is represented by an embedding function , typically a signed distance function , and its evolution amounts to solving a differential equation ( Osher & Sethian , 1988 ) . In this work , we extend the formulation of the level set method . Inspired by the recent progress in Neural Ordinary Diferential Equations ( NODEs ) ( Chen et al. , 2018 ; Dupont et al. , 2019 ; Gholami et al. , 2019 ) , we propose to use NODEs to solve the level set formulation of the contour evolution , thus learning the forcing function in an end-to-end data driven manner . Unlike earlier attempts in combining the level set method with CNNs , we benefit from NODEs parametrization of the derivative of the contour because it allows us to incorporate external constraints that guide the contour evolution , e.g . by adding a regularization penalty to the curvature of the front or exploiting images at the evolving front by extracting appearance constraints in a non-supervised way . Finally , similar to experiments in ( Chen et al. , 2018 ) , to alleviate the need for careful choice or design of contour embedding functions , we propose a NODE-based method that evolves an image embedding into a dense per-pixel semantic label space . To the best of our knowledge , this work is the first to apply Neural ODEs to real world problems . We validate our methods on two 2D segmentation tasks : kidney segmentation in transversal slices of CT scans and salient object segmentation . Given an initial estimate of kidney via existing algorithms , our method effectively evolves the initial estimates and achieves improved kidney segmentation , as we show in Figure 1 . On real life salient objects , in addition to contour evolution , we use our method to directly evolve the embedding of an input image into a pixel-wise dense semantic label . Following ( Hu et al. , 2017 ) , we compare against the results in ( Wang et al. , 2017 ; Li et al. , 2016 ; Li & Yu , 2015 ; Zhao et al. , 2015 ; Lee et al. , 2016 ; Wang et al. , 2015 ; Hu et al. , 2017 ) and achieve ω-Fβ scores , PASCAL-S 0.668 and ECSSD 0.768 , that are higher than several state-of-the art algorithms . Our results suggest the potential of utilizing NODEs for solving the contour evolution of level set methods or the direct evolution of image embeddings into segmentation maps . We hope our findings will inspire future research in using NODEs for semantic segmentation tasks . We foresee that our method would allow for intervention on intermediate states of the solution of the ODE , allowing for injection of shape priors or other regularizing constraints . In summary , our contributions are : • We propose to use NODEs to solve the level set formulation of the contour evolution . • We propose using NODEs to learn the forcing function in an end-to-end data driven way . • We show NODEs can also evolve image embeddings directly into dense per-pixel semantic label spaces , which may alleviate the need for careful choice or design of contour embedding functions . 2 METHODS . Suppose I is a 2D image , S is the contour of an object we want to segment , and φ is a contour embedding function , defined as a distance map , such that S = { ( x , y ) |φ ( x , y ) = 0 } . We assume an initial but rough contour of the object is given by a human operator or by an existing algorithm . A level set segmentation ( Osher & Sethian , 1988 ) solves a differential equation to evolve a contour along its normal direction with a speed function F as : dφ dt = |∇φ|F for t ∈ [ 0 , 1 ] , ( 1 ) where the initial value φ0 ( x , y ) is defined as a signed Euclidean distance from ( x , y ) to the closest point on the initial contour S0 . The speed function F is often modelled to be a function of the target image I , the shape statistics of the object contour ( derived from training shapes ) , or a regularizing curvature term ( ∇ ∇φ‖∇φ‖ ) . In Neural ODEs , we parametrize the derivative of the hidden state h using a neural network fθ parametrized by θ : dh dt = fθ ( h , t ) . ( 2 ) The relationship between Eq . 1 and Eq . 2 is immediate . In the next section , we propose two approaches that adapt NODEs to the level set method for image segmentation . 2.1 CONTOUR EVOLUTION WITH NODES . We propose to solve a more general form of Eq . 1 to evolve an initial contour estimate φ̂ for image segmentation . We define the state of the NODE to be φ̂ augmented with the input image ’ s embedding , h. We then advance the augmented state , γ = ( φ̂ , h ) , using NODEs , which can be interpreted as estimating the speed function F described in Eq . 1 . Mathematically , γ = ( φ̂ , h ) , dγ dt = fθ ( γ , t ) for t ∈ [ 0 , 1 ] , γ ( 0 ) = ( φ̂ ( 0 ) , h ( 0 ) ) , φ̃ = φ̂ ( 1 ) + ψ ( γ ( 1 ) ) , ( 3 ) where t is the time step in the evolution , γ is the augmented state of the NODE , f is a neural network parametrized by θ , φ̂ ( 0 ) is the initial value of the distance map , h ( 0 ) is the initial value of the image embedding , ψ is a learned function and φ̃ is the dense per-pixel distance map prediction . Figure 2a schematically illustrates our initial contour evolution approach . Throughout this paper , we will refer to this method as Contour Evolution . 2.2 IMAGE EVOLUTION WITH NODES . In our first approach , we obtain a final optimal contour by evolving an initial estimate . In our second approach , inspired by Chen et al . ( 2018 ) , we evolve an image embedding h and project it into a dense per-pixel distance map φ̃ , whose zero level set defines the final segmentation map , S ( t ) = { ( x , y ) |φ ( t ) ( x , y ) = 0 } . Mathematically , dh dt = fθ ( h , t ) for t ∈ [ 0 , 1 ] , h ( 0 ) = λ ( I ) , φ̃ = ψ ( h ( 1 ) ) , ( 4 ) where t is the time step in the evolution , f is a neural network parametrized by θ , I is an image , λ is a learned image embedding function and ψ is a learned function that maps an image embedding to a distance map . Figure 2b schematically illustrates our direct image evolution approach . Throughout this paper , we will refer to this method as Image Evolution . 3 IMPLEMENTATION . In the following subsections , we describe our design choices in loss function and their regularization terms , architectures , strategies for emphasizing the evolution of the contour on a region of interest . We also detail our model initialization strategies to prevent drifting from the sub-optimal initial value , and choices of error tolerances and activation normalization . 3.1 LOSS FUNCTION AND REGULARIZATION TERMS . We optimize the parameters of our NODE models , described in Figures 2a and 2b , to minimize the empirical risk computed as the Mean Squared Error ( MSE ) between the target ( φ ) and predicted ( φ̃ ) distance maps . We remind the reader that although our techniques can access intermediate NODE states , which could allow injection of priors or other constraints , we do not explore this in our current experiments , and relay it to future work . 3.2 NARROW BAND AND RE-INITIALIZATION . In the level set formulation , all levels that describe the propagating contour are tracked . Adalsteinsson & Sethian ( 1995 ) proposed limiting the evolution to the subset of levels within a narrow band of the zero level contour . In our approach , we obtain the equivalent of a narrow band by applying a hyperbolic tangent nonlinearity on the evolved distance map . It effectively attenuates the contribution of levels in the optimization process . This transformation is especially valuable in refinement setups because it weights the gradients of the loss according to the proximity to the contour1 . Re-initialization of φ is another common practice in classical level set methods . It ensures the states in the trajectory of the numerical solution remain valid distance maps . ( Sussman et al. , 1994 ; Hartmann et al. , 2010 ) propose to first extract a zero level set of an evolving state , and re-calculate a distance map of that contour . In our experiments , we found that our optimization is not sensitive to non valid distance maps , and we did not find it necessary to reinitialize φ . 3.3 PARAMETER INITIALIZATION AND LEARNING RATE RAMPUP . In tasks where the initial value is already close to the desired solution , not initializing the model parameters to represent the identity function and not using learning rate ramp up can slow down the optimization process as the model predictions can immediately drift away from the initial value . In addition to using learning rate rampup , we prevent this issue by setting the weights and biases on the last layer of the NODE and Postnet layers to zero . This approach has been successfully used in normalizing flow models ( Kingma & Dhariwal , 2018 ; Prenger et al. , 2019 ) . 3.4 ADAPTIVE SOLVERS AND ERROR TOLERANCES . In ordinary differential equations , adaptive step solvers vary the step size according to the error estimate of the current step and the error tolerance . If the error estimate is larger than the threshold , the step will be redone with a smaller size until the error is smaller than the error tolerance . The error tolerance eitol given the current state i is the sum of the absolute error tolerance atol and the infinity norm of the current state h weighted by the relative error tolerance rtol : eitol = atol + rtol ∗ ∥∥hi∥∥∞ . ( 5 ) 1For the hyperbolic tangent , the gradients decrease as it moves away from zero , which represents the contour at the zero level set . Given that we do not know in advance the infinity norm of hi , which in our case contains the image embedding as described in Equations 3 and 4 , we set the contribution of the relative error tolerance term to zero and adjust the absolute error tolerance .

Summary: This paper proposes to utilize Neural ODEs (NODEs) and the Level Set Method (LSM) for the task of image segmentation.  The argument is that the NODE can be used to learn the force function in an LSM and solve the contour evolution process. The authors propose two architectures and demonstrate promising performance on a few image segmentation benchmarks.