Deep Convolutional Neural Networks Meet Variational Shape Compactness Priors for Image Segmentation

TL;DR

This work tackles image segmentation with a shape-compactness prior by formulating a nonconvex variational model that penalizes the perimeter-to-area ratio. It introduces two primal-dual optimization schemes, PD-TD and PD-STD, deriving equivalent models via a Gaussian-kernel TV surrogate and a dual variable, enabling efficient updates and differentiable network integration. The PD-STD variant is embedded as a neural-network block to enforce compact segmentation results, and experiments on Iris, fundus, and synthetic datasets show superior IoU, Dice, and compactness, especially under heavy noise, outperforming ADMM-based approaches. The proposed PD-STD layer can be plugged into popular segmentation networks (e.g., DeepLabV3, IrisParseNet) and is applicable to a wide range of architectures, offering improved robustness and accuracy in noisy imaging scenarios.

Abstract

Shape compactness is a key geometrical property to describe interesting regions in many image segmentation tasks. In this paper, we propose two novel algorithms to solve the introduced image segmentation problem that incorporates a shape-compactness prior. Existing algorithms for such a problem often suffer from computational inefficiency, difficulty in reaching a local minimum, and the need to fine-tune the hyperparameters. To address these issues, we propose a novel optimization model along with its equivalent primal-dual model and introduce a new optimization algorithm based on primal-dual threshold dynamics (PD-TD). Additionally, we relax the solution constraint and propose another novel primal-dual soft threshold-dynamics algorithm (PD-STD) to achieve superior performance. Based on the variational explanation of the sigmoid layer, the proposed PD-STD algorithm can be integrated into Deep Neural Networks (DNNs) to enforce compact regions as image segmentation results. Compared to existing deep learning methods, extensive experiments demonstrated that the proposed algorithms outperformed state-of-the-art algorithms in numerical efficiency and effectiveness, especially while applying to the popular networks of DeepLabV3 and IrisParseNet with higher IoU, dice, and compactness metrics on noisy Iris datasets. In particular, the proposed algorithms significantly improve IoU by 20% training on a highly noisy image dataset.

Figures (11)

• Figure 1: The network architecture of the proposed PD-STD Layers used for our proposed network as shown in Fig. \ref{['fig:architecture']}. The red rectangle denotes the dual variables in the regularization space. The blue rectangle denotes the dual variables in a compact-shaped space.
• Figure 2: The architecture of a segmentation network with PD-STD layer.
• Figure 3: Experimental results of different algorithms with different $\lambda$ for comparison: it is obvious that the segmentation region tends to be more circular as the weight parameter $\lambda$ in \ref{['eq:primal1']} becomes smaller (which means the weight of the introduced shape-compactness regularization $C(u)$ is bigger). This shows that the introduced shape-compactness regularization $C(u)$ does work properly for all the three algorithms of ADMM, PD-TD and PD-STD.
• Figure 4: Segmentation results by different algorithms on some real images. The proposed algorithms of PD-TD and PD-STD generate more compact region with smoother boundaries as segmentation results; in contrast, the compared ADMM algorithm sometimes fails to obtain a compact segmentation region as result.
• Figure 5: Segmentation results predicted by DeepLabV3 and our methods of PD-STD + DeepLabV3 trained on a clean image dataset. Noise level of the test image dataset from top to bottom: Gaussian noise level of $0.01$, $0.07$, salt and pepper noise level of $0.01$.

Theorems & Definitions (6)

• Proposition 1
• proof
• Lemma 1: Existence of minimizer
• proof
• Lemma 2: Convergence of minimizers
• proof