Segmentation then Prediction: A Multi-task Solution to Brain Tumor Segmentation and Survival Prediction

Accurate brain tumor segmentation and survival prediction are two fundamental but challenging tasks in the computer aided diagnosis of gliomas. Traditionally, these two tasks were performed independently, without considering the correlation between them. We believe that both tasks should be performed under a unified framework so as to enable them mutually benefit each other. In this paper, we propose a multi-task deep learning model called segmentation then prediction (STP), to segment brain tumors and predict patient overall survival time. The STP model is composed of a segmentation module and a survival prediction module. The former uses 3D U-Net as its backbone, and the latter uses both local and global features. The local features are extracted by the last layer of the segmentation encoder, while the global features are produced by a global branch, which uses 3D ResNet-50 as its backbone. The STP model is jointly optimized for two tasks. We evaluated the proposed STP model on the BraTS 2020 validation dataset and achieved an average Dice similarity coefficient (DSC) of 0.790, 0.910, 0.851 for the segmentation of enhanced tumor core, whole tumor, and tumor core, respectively, and an accuracy of 65.5% for survival prediction.