Variational Wasserstein Model on Riemannian Manifolds for Image Segmentation

The data fidelity term plays a crucial role in variational image segmentation. Recently, optimal transport (OT)-based variational models have improved segmentation by globally matching predicted feature distributions with reference distributions. However, existing OT-based methods usually rely on color or intensity similarity, which can be insufficient when foreground and background regions share similar appearance. We address this limitation by incorporating geodesic distance into the OT cost on a Riemannian feature manifold, allowing the transport cost to jointly encode spatial proximity, appearance similarity, and boundary information. The proposed model is well-defined: we prove existence of minimizers and convexity of the energy functional. It is also tractable: the OT terms admit differentiable transport-map representations, avoiding storage of the full pairwise transport matrix and enabling efficient gradient-based optimization. Experiments on BraTS and RESECT show that our method consistently outperforms traditional variational models, OT-based baselines, and prompt-based deep segmentation models, particularly when foreground and background have overlapping intensity or color patterns.