Adaptive Thresholding based on Multi-task Learning for Refining Binary Medical Image Segmentation
Qin Lei, Jiang Zhong, Chen Wang, Qizhu Dai, Rongzhen Li · 2023
Binary medical image segmentation plays a pivotal role in the diagnosis and treatment of a wide range of diseases. However, the performance of the segmentation model is closely related to the choice of the binarization threshold (default 0.5), which is used to binarize the output probability mask. In this paper, we introduce an innovative multi-task learning framework featuring an Adaptive Thresholding Module (ATM) designed to predict the optimal threshold for each image. Within this multi-task learning framework, the segmentation task is divided into two distinct subtasks. The first subtask focuses on the original segmentation task to output the probabilistic mask. Simultaneously, the second subtask leverages spatial features extracted from the segmentation network, with ATM learning and applying these features in a regression task to derive the optimal threshold for each image. Subsequently, a binarization process is enacted using these optimal thresholds, leading to a marked improvement in segmentation accuracy. The crux of ATM’s contribution to enhanced segmentation accuracy lies in its ability to optimize the binarization process, striking a well-balanced equilibrium between the labels of true positive and false positive. We integrated ATM into various medical segmentation models and subjected it to evaluation on datasets encompassing diverse binarized medical image patterns. The results underscore the effectiveness of ATM in elevating the accuracy of pre-existing medical segmentation models.