Bottleneck Supervised U-Net for Pixel-wise Liver and Tumor Segmentation

Abstract: In this paper, we propose a bottleneck supervised (BS) U-Net model for liver and tumor segmentation. Our main contributions are: first, we propose a variation of the original U-Net that incorporates dense modules, inception modules and dilated convolution in the encoding path; second, we propose a bottleneck supervised (BS) U-Net that contains an encoding U-Net and a segmentation U-Net. To train the BS U-Net, the encoding U-Net is first trained to get encodings of the label maps that contain the anatomical information (shape and location). Subsequently, this information is used to guide the training of the segmentation U-Net so as to reserve the anatomical features of the target objects. More specifically, the loss function for segmentation U-Net is set to be the weighted average of the dice loss and the MSE loss between the encodings and the bottleneck feature vectors. The model is applied to a public liver and tumor CT scan dataset. Experimental results show that besides achieving excellent overall segmentation performance, BS U-Net also works great in controlling shape distortion, reducing false positive and false negative cases.