Graph Cuts Loss to Boost Model Accuracy and Generalizability for Medical Image Segmentation

Abstract

Segmentation accuracy and generalization ability are essential for deep learning models, especially in medical image segmentation. We present a novel, robust yet straightforward loss function to boost model accuracy and generalizability for medical image segmentation. We reformulate the graph cuts cost function to a loss function for supervised learning. The graph cuts loss innately focuses on a dual penalty to optimize the regional properties and boundary regularization. We benchmark the proposed loss on six public retinal vessel segmentation datasets with a comprehensive intra-dataset and cross-dataset evaluation. Results reveal that the proposed loss is more generalizable, narrowing the performance gap between different architectures. Besides, models trained with our loss show higher segmentation accuracy and better generalization ability than those trained with other mainstream losses. Moreover, we extend our loss to other segmentation tasks, e.g., left atrium and liver tumor segmentation. The proposed loss still achieves comparable performance to the state-of-the-art, demonstrating its potential for any N-D segmentation problem. The code is available at https://github.com/zzh_enggit/graphcutsloss.