Blending variational approaches and deep learning to enforce prior constraints in medical image segmentation

Overview & Synopsis

Image segmentation plays a crucial role in computer vision, particularly in medical image analysis. While convolutional neural networks (CNNs) are the current state-of-the-art for segmentation, they may generate artifacts, such as isolated pixels or holes, when object contours are not well defined. Therefore, incorporating prior knowledge into the segmentation process is essential, whether through topological prescriptions (e.g., the number of connected components, the partial convexity of an object boundary) or geometrical constraints (e.g., penalizing volume through constraints), particularly when preserving contextual relationships between objects and achieving a segmentation that is homeomorphic to a known a priori is desired. In this paper, we propose two hybrid variational/deep learning frameworks that include geometric and topological constraints in the training of CNNs, in the form of two penalty terms in the loss function. The first model incorporates geometric constraints via a regularization based on weighted total variation, a volume/area penalty, and a Mumford-Shah like term. In the second model, we view the segmentation process as a registration task that pairs the ground truth and the image to be labelled, based on non-linear elasticity principles. Incompressibility conditions stated in terms of the Jacobian matrix determinant of the deformation ensures volume and topology preservation, without any material self-intersection. We evaluate the proposed losses in the challenging problem of thoracic organs at risk segmentation in Computed Tomographic scans, where some organs lack contrast, and their borders are ill-defined. Our ablation study demonstrates that our methods provide significant improvements over existing unconstrained approaches, in terms of both quantitative criteria such as the measurement of region overlap and qualitative assessment, particularly when dealing with unbalanced classes.