Image Segmentation via Variational Model Based Tailored UNet: A Deep Variational Framework

TL;DR

This work tackles image segmentation by marrying variational PDE priors with deep learning. It introduces VM_TUNet, a hybrid framework that couples a fourth-order Cahn-Hilliard based model with a UNet backbone, and learns the boundary-preserving operator F(f) end-to-end. The Tailored Finite Point Method is employed to maintain sharp boundaries, while the data-driven F(f) eliminates manual parameter tuning. Experiments on standard benchmarks show improved accuracy and Dice scores over several strong baselines, though challenges remain for large-scale, dense instance segmentation. Overall, VM_TUNet offers an interpretable, data-driven pathway for precise segmentation in diverse imaging domains.

Abstract

Traditional image segmentation methods, such as variational models based on partial differential equations (PDEs), offer strong mathematical interpretability and precise boundary modeling, but often suffer from sensitivity to parameter settings and high computational costs. In contrast, deep learning models such as UNet, which are relatively lightweight in parameters, excel in automatic feature extraction but lack theoretical interpretability and require extensive labeled data. To harness the complementary strengths of both paradigms, we propose Variational Model Based Tailored UNet (VM_TUNet), a novel hybrid framework that integrates the fourth-order modified Cahn-Hilliard equation with the deep learning backbone of UNet, which combines the interpretability and edge-preserving properties of variational methods with the adaptive feature learning of neural networks. Specifically, a data-driven operator is introduced to replace manual parameter tuning, and we incorporate the tailored finite point method (TFPM) to enforce high-precision boundary preservation. Experimental results on benchmark datasets demonstrate that VM_TUNet achieves superior segmentation performance compared to existing approaches, especially for fine boundary delineation.

Figures (10)

• Figure 1: Traditional image segmentation using variational PDE methods involves formulating an energy functional whose minimization corresponds to an optimal segmentation, and solving the associated partial differential equations to find this minimum. Here are the results of different number of iterations of the Chan-Vese algorithm; see in Appendix \ref{['chan']}.
• Figure 2: VM_TUNet architecture. (a) VM_TUNet process. (b) VM_TUNet block, which is the process of solving Cahn-Hilliard equation. (c) UNet class architecture which approximates $F(f)$ in (a).
• Figure 3: Comparison results of cicada and eggs (From top to bottom) of HKU-IS between the proposed models: UNet, UNet++, DeepLabV3+, DN-I, and VM_TUNet. The pictures from left to right are: Image; Ground Truth; and the results of UNet, UNet++, DeepLabV3+, DN-I, and VM_TUNet, respectively.
• Figure 4: Comparison results of family, gentleman, plane, and bench (From top to bottom) of DUT-OMRON between the proposed models: UNet, UNet++, DeepLabV3+, DN-I, and VM_TUNet. The pictures from left to right are: Image; Ground Truth; and the results of UNet, UNet++, DeepLabV3+, DN-I, and VM_TUNet, respectively.
• Figure 5: Illustration of UNet type network with input of size $N_1\times N_2\times D$. The left branch is the encoding part, the right branch is the decoding part, and the bottom rectangle denotes the bottleneck. Green arrows represent downsampling operations. Transparent arrows represent upsampling operations. Horizontal red dashed arrows represent skip connections. The orange rectangles denote the outputs of the encoding part that are passed to the decoding part via the skip connections. The length and width of the rectangle represent the output resolution and number of channels, respectively. Every network in this class can be fully characterized by the channels vector $\boldsymbol{c}$: given a $\boldsymbol{c}\in\mathbb{R}^S$, the corresponding network has $S+1$ resolution levels, $c_s$ channels at resolution level $s$ for $1\leq s\leq S$, and $2c_S$ channels at resolution level $S+1$liu2024double.