Attention U-Net ensemble for interpretable polyp and instrument segmentation

Abstract

The difficulty associated with screening and treating colorectal polyps alongside other gastrointestinal pathology presents an opportunity to incorporate computer-aided systems. This paper develops a deep learning pipeline that accurately segments colorectal polyps and various instruments used during endoscopic procedures. To improve transparency, we leverage the Attention U-Net architecture, enabling visualisation of the attention coefficients to identify salient regions. Moreover, we improve performance by incorporating transfer learning using a pre-trained encoder, together with test-time augmentation, softmax averaging, softmax thresholding and connected component labeling to further refine predictions.