RelationalUNet for Image Segmentation

Medical image segmentation is one of the most classic applications of machine learning in healthcare. A variety of Deep Learning approaches, mostly based on Convolutional Neural Networks (CNNs), have been proposed to this end. In particular, U-Shaped Network (UNet) have emerged to exhibit superior performance for medical image segmentation. However, some properties of CNNs, such as the stationary kernels, may limit them from capturing more in-depth visual and spatial relations. The recent success of transformers in both language and vision has motivated dynamic feature transforms. We propose RelationalUNet (RelationalUNet) which introduces relational feature transformation to the UNet architecture. RelationalUNet models the dynamics between visual and depth dimensions of a 3D medical image by introducing Relational Self-Attention blocks in skip connections. As the architecture is mainly intended for the semantic segmentation of 3D medical images, we aim to learn their long-range depth relations. Our method was validated on the Multi-Atlas Labeling Beyond The Cranial Vault (BTCV) dataset for multi-organ segmentation. Robustness to distribution shifts is a particular challenge in safety-critical applications such as medical imaging. We further test our model performance on realistic distributional shifts on the Shifts 2.0 White Matter Multiple Sclerosis Lesion Segmentation. Experiments show that our architecture leads to competitive performance. The code is available at https://github.com/ivaxi0s/runet.