LMS-Net: A Learned Mumford-Shah Network For Few-Shot Medical Image Segmentation
Shengdong Zhang, Fan Jia, Xiang Li, Hao Zhang, Jun Shi, Liyan Ma, Shihui Ying
TL;DR
The paper tackles the challenge of few-shot medical image segmentation by building an interpretable framework that unifies pixel-to-prototype matching with deep priors via a learned Mumford–Shah (LMS) variational model. It introduces LMS-Net, an end-to-end deep unfolding network that decomposes optimization into a prototype-update pathway (via a Momentum Update Transformer) and a mask-update pathway (via a primal–dual with a learned denoiser, PD-Net). The approach provides clear interpretability and leverages structural priors to capture regional geometry, reporting improved Dice scores across three medical datasets and through extensive ablations. The work suggests that integrating variational priors with deep learning in a staged, interpretable architecture can enhance robustness and accuracy for FSS in clinically scarce data scenarios, with code forthcoming at the project repository.
Abstract
Few-shot semantic segmentation (FSS) methods have shown great promise in handling data-scarce scenarios, particularly in medical image segmentation tasks. However, most existing FSS architectures lack sufficient interpretability and fail to fully incorporate the underlying physical structures of semantic regions. To address these issues, in this paper, we propose a novel deep unfolding network, called the Learned Mumford-Shah Network (LMS-Net), for the FSS task. Specifically, motivated by the effectiveness of pixel-to-prototype comparison in prototypical FSS methods and the capability of deep priors to model complex spatial structures, we leverage our learned Mumford-Shah model (LMS model) as a mathematical foundation to integrate these insights into a unified framework. By reformulating the LMS model into prototype update and mask update tasks, we propose an alternating optimization algorithm to solve it efficiently. Further, the iterative steps of this algorithm are unfolded into corresponding network modules, resulting in LMS-Net with clear interpretability. Comprehensive experiments on three publicly available medical segmentation datasets verify the effectiveness of our method, demonstrating superior accuracy and robustness in handling complex structures and adapting to challenging segmentation scenarios. These results highlight the potential of LMS-Net to advance FSS in medical imaging applications. Our code will be available at: https://github.com/SDZhang01/LMSNet
