Boundary Aware U-Net for Medical Image Segmentation

Automatic medical image segmentation plays an integral role in the health care system as it facilitates the cancer detection process and provides a basis to analyze and monitor cancer progress. Convolutional neural networks have proven to be an effective approach to automate medical image segmentation tasks. These networks perform a set of convolutional layers followed by the activation and pooling operations to represent the object of interest in terms of texture and semantic information. Although the texture information can reveal the disorders in medical images, it pays less attention to the anatomical structure of the human tissue and is consequently less precise in the boundary area. To compensate for the boundary representation, we propose to incorporate the Vision Transformer (ViT) model on top of the bottleneck layer. In our design, we seek to model the distribution of the boundary area using the global contextual representation deriving from the ViT module. In addition, by fusing the boundary representation generated by the ViT module to each decoding block, we preserve the anatomical structure for the boundary-aware segmentation. Throughout a comprehensive evaluation of several medical image segmentation tasks, we demonstrate the effectiveness of our model. Particularly our method achieved ISIC2017: 0.905, ISIC2018: 0.898, PH2: 0.944 and the Lung segmentation task with 0.990 dice scores.