[Paper Review] Neural ODEs for Image Segmentation with Level Sets

Rafael Valle, Fitsum A. Reda+4 more|arXiv (Cornell University)|Dec 25, 2019

Medical Image Segmentation Techniques43 references4 citations

TL;DR

This paper proposes a novel image segmentation method that combines Neural ODEs (NODEs) with the Level Set method to evolve initial contours or image embeddings into precise semantic segmentations. By learning the forcing function end-to-end via NODEs, the approach improves segmentation accuracy with fewer parameters, achieving state-of-the-art $F_\beta$ scores on kidney and salient object segmentation benchmarks.

ABSTRACT

We propose a novel approach for image segmentation that combines Neural Ordinary Differential Equations (NODEs) and the Level Set method. Our approach parametrizes the evolution of an initial contour with a NODE that implicitly learns from data a speed function describing the evolution. In addition, for cases where an initial contour is not available and 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. We evaluate our methods on kidney segmentation (KiTS19) and on salient object detection (PASCAL-S, ECSSD and HKU-IS). In addition to improving initial contours provided by deep learning models while using a fraction of their number of parameters, our approach achieves F scores that are higher than several state-of-the-art deep learning algorithms.

Motivation & Objective

• Address the challenge of accurate pixel-level segmentation near ambiguous object boundaries, especially in medical and salient object imaging.
• Overcome limitations of traditional level set methods that rely on hand-designed speed functions and lack data-driven adaptation.
• Reduce dependency on carefully designed initial contours or embedding functions by learning evolution dynamics via Neural ODEs.
• Explore the use of NODEs to directly evolve image embeddings into dense per-pixel semantic labels, bypassing explicit contour representation.
• Demonstrate that NODE-based evolution outperforms or matches state-of-the-art deep learning models with significantly fewer parameters.

Proposed method

• Parametrize the level set evolution of a contour using a Neural ODE, where the derivative of the level set function $\phi$ is governed by a learnable neural network $f_\theta$.
• Use a signed distance function as the initial embedding $\phi^0$, evolving it via $\frac{d\phi}{dt} = |\nabla\phi| F$, with $F$ implicitly learned by the NODE.
• Train the NODE end-to-end to minimize segmentation loss, allowing the model to learn optimal contour evolution dynamics from data.
• Introduce a second approach where the image embedding itself is evolved via a NODE into a dense pixel-wise semantic label map, eliminating reliance on initial contour design.
• Apply regularization terms such as curvature control and appearance-based constraints during training to guide evolution and improve generalization.
• Use architectures like NodeUNet (for contour evolution) and NodeStack (for image embedding evolution) to efficiently learn the forcing function with reduced parameter count.

Experimental results

Research questions

• RQ1 Can Neural ODEs effectively learn the forcing function for level set evolution in image segmentation, replacing hand-designed speed functions?
• RQ2 Does end-to-end learning of the forcing function via NODEs improve segmentation accuracy compared to standard deep learning models?
• RQ3 Can NODEs be used to evolve image embeddings directly into dense semantic label maps, eliminating the need for explicit initial contour design?
• RQ4 How does the performance of the proposed NODE-based method compare to state-of-the-art models in terms of $F_\beta$ scores and parameter efficiency?
• RQ5 To what extent can external constraints (e.g., curvature regularization, appearance features) be integrated into the NODE-based evolution process?

Key findings

• The NodeUNet model (5M parameters) improved $\omega$-$F_\beta$ scores on PASCAL-S by 5% relative to a 15M-parameter UNet baseline, demonstrating superior parameter efficiency.
• On ECSSD, the NodeStack model (41M parameters) achieved an $\omega$-$F_\beta$ score of 0.768, outperforming several state-of-the-art models including DLS (0.766) and ELD (0.756).
• The NodeStack model achieved an $\alpha$-$F_\beta$ score of 0.848 on ECSSD, surpassing all compared methods, including MTDS, MDF, and LEGS.
• On HKU-IS, the NodeStack model achieved an $\omega$-$F_\beta$ score of 0.734, exceeding DLS (0.748) and ELD (0.718), showing strong generalization on challenging salient object datasets.
• The method achieved state-of-the-art performance on PASCAL-S with an $\omega$-$F_\beta$ of 0.668, outperforming DRFI, MCDL, LEGS, MTDS, MDF, and ELD.
• The approach enables intervention during ODE solving, allowing for potential integration of shape priors or curvature regularization in future work.