[Paper Review] TransNetR: Transformer-based Residual Network for Polyp Segmentation with Multi-Center Out-of-Distribution Testing

Debesh Jha, Nikhil Kumar Tomar+2 more|arXiv (Cornell University)|Mar 13, 2023

Colorectal Cancer Screening and Detection27 citations

TL;DR

TransNetR is an encoder-decoder polyp segmentation model that combines a pretrained ResNet50 encoder with a Residual Transformer block to achieve real-time performance and strong generalization to out-of-distribution (OOD) data across multi-center datasets."

ABSTRACT

Colonoscopy is considered the most effective screening test to detect colorectal cancer (CRC) and its precursor lesions, i.e., polyps. However, the procedure experiences high miss rates due to polyp heterogeneity and inter-observer dependency. Hence, several deep learning powered systems have been proposed considering the criticality of polyp detection and segmentation in clinical practices. Despite achieving improved outcomes, the existing automated approaches are inefficient in attaining real-time processing speed. Moreover, they suffer from a significant performance drop when evaluated on inter-patient data, especially those collected from different centers. Therefore, we intend to develop a novel real-time deep learning based architecture, Transformer based Residual network (TransNetR), for colon polyp segmentation and evaluate its diagnostic performance. The proposed architecture, TransNetR, is an encoder-decoder network that consists of a pre-trained ResNet50 as the encoder, three decoder blocks, and an upsampling layer at the end of the network. TransNetR obtains a high dice coefficient of 0.8706 and a mean Intersection over union of 0.8016 and retains a real-time processing speed of 54.60 on the Kvasir-SEG dataset. Apart from this, the major contribution of the work lies in exploring the generalizability of the TransNetR by testing the proposed algorithm on the out-of-distribution (test distribution is unknown and different from training distribution) dataset. As a use case, we tested our proposed algorithm on the PolypGen (6 unique centers) dataset and two other popular polyp segmentation benchmarking datasets. We obtained state-of-the-art performance on all three datasets during out-of-distribution testing. The source code of TransNetR will be made publicly available at https://github.com/DebeshJha.

Motivation & Objective

• Motivate real-time, accurate polyp segmentation suitable for clinical use.
• Address generalization gaps when models are tested on data from unseen centers or distributions.
• Propose a transformer-enhanced residual architecture that preserves speed and improves robustness to distribution shifts.

Proposed method

• Encoder-decoder architecture using a pretrained ResNet50 as the encoder.
• Four intermediate feature maps from the encoder are compressed with 1x1 convolutions and fed into a three-block decoder with skip connections.
• Residual Transformer (RT) blocks fuse convolutional features with transformer-based self-attention.
• Final decoder stage uses a residual block instead of RT to reduce parameters, followed by upsampling and a sigmoid 1x1 conv for segmentation.

Figure 1: Illustration of different scenarios expected to arise in real-world settings. The proposed work conducted both in-distribution and out-of-distribution validation process. C1 to C6 represent the different centers data present in PolypGen [ Ali et al.(2023)Ali, Jha, Ghatwary, Realdon, Canniz

Experimental results

Research questions

• RQ1 How well does TransNetR perform on in-distribution polyp segmentation benchmarks compared to state-of-the-art methods?
• RQ2 Does TransNetR generalize to out-of-distribution data from multiple centers/datasets (OOD testing)?
• RQ3 What is the impact of the Residual Transformer block on segmentation accuracy and model efficiency?
• RQ4 Can the model maintain real-time inference speed while achieving high segmentation quality across diverse datasets?

Key findings

• TransNetR achieves a Dice coefficient of 0.8706 and a mean IoU of 0.8016 on the Kvasir-SEG test set, with recall 0.8843 and precision 0.9073 at 54.60 FPS.
• Out-of-distribution testing shows TransNetR delivering state-of-the-art performance across PolypGen (6 centers), BKAI-IGH, and CVC-ClinicDB datasets.
• Ablation shows the Residual Transformer (RT) block improves metrics (e.g., +1.34% mIoU) over variants without RT, while retaining real-time speed.
• Across multiple OOD evaluations, TransNetR consistently outperforms competitors, including UACANet and UNeXt, in mIoU and DSC across center-wise and dataset-wise analyses.
• Center-wise results indicate robust performance across data from diverse centers, including small and multiple polyps, with accurate boundary delineation.

Figure 2: Block diagram of TransNetR along with the Residual Transformer block