Deep Distance Map Regression Network with Shape-Aware Loss for Imbalanced Medical Image Segmentation

Abstract

Small object segmentation, like tumor segmentation, is a difficult and critical task in the field of medical image analysis. Although deep learning based methods have achieved promising performance, they are restricted to the use of binary segmentation mask and suffer from the imbalance problem. In this research, we aim to tackle this limitation by adopting distance map as a novel ground truth and employing distance map regression as a proxy of the existing segmentation framework. Specially, we propose a new segmentation framework that incorporates the existing binary segmentation network and a light weight regression network (dubbed as LR-Net). Thus, the LR-Net can convert the conventional classification-based segmentation into a regression task and leverage the rich information of distance maps. Additionally, we derive a shape-aware loss by employing distance maps as penalty map to capture the complete shape of an object. We evaluated our approach on MICCAI 2017 Liver Tumor Segmentation (LiTS) Challenge dataset and a clinical dataset. Experimental results show that our approach outperforms the classification-based methods as well as other existing state-of-the-arts. Code is available at https://github.com/Huiyu-Li/Deep-Distance-Map-Regression.