[논문 리뷰] TransNetR: Transformer-based Residual Network for Polyp Segmentation with Multi-Center Out-of-Distribution Testing

Debesh Jha, Nikhil Kumar Tomar외 2명|arXiv (Cornell University)|2023. 03. 13.

Colorectal Cancer Screening and Detection인용 수 27

한 줄 요약

TransNetR은 다중 센터 데이터셋에서 실시간 성능과 OOD(Out-of-Distribution) 데이터에 대한 강한 일반화 성능을 달성하기 위해 사전 학습된 ResNet50 인코더와 잔차 트랜스포머(Residual Transformer) 블록을 결합한 인코더-디코더 폴립(segmentation) 모델이다.

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.

연구 동기 및 목표

• 실시간으로 임상 사용에 적합한 정확한 폴립 분할을 동기부여한다.
• 모델이 보지 못한 센터나 분포에서 데이터로 평가될 때의 일반화 격차를 해결한다.
• 속도를 유지하고 분포 이동에 대한 강인성을 향상시키는 트랜스포머 강화 잔차 구조를 제안한다.

제안 방법

• 사전 학습된 ResNet50을 인코더로 사용하는 인코더-디코더 아키텍처.
• 인코더의 네 개의 중간 피처 맵을 1x1 컨볼루션으로 압축하고 스킵 연결이 있는 세 블록 디코더로 입력한다.
• Residual Transformer(RT) 블록은 컨볼루션 피처를 트랜스포머 기반 셀프 어텐션과 융합한다.
• 최종 디코더 단계는 매개변수를 줄이기 위해 RT 대신 잔차 블록을 사용한 후 업샘플링과 시그모이드 1x1 컨볼루션으로 분할을 수행한다.

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

실험 결과

연구 질문

• RQ1 TransNetR이 현 배포 분포에서 최첨단 방법과 비교하여 폴립 분할 벤치마크에서 얼마나 잘 수행하는가?
• RQ2 TransNetR이 여러 센터/데이터셋의 OOD 데이터에 일반화하는가(OOD 테스트)?
• RQ3 Residual Transformer 블록이 분할 정확도와 모델 효율성에 미치는 영향은 무엇인가?
• RQ4 다양한 데이터셋에서 고품질 분할을 달성하면서 실시간 추론 속도를 유지할 수 있는가?

주요 결과

방법	mIoU	mDSC	Rec.	Prec.	F2	FPS	Para.	Flops
U-Net	0.7472	0.8264	0.8504	0.8703	0.8353	106.88	31.04	54.75
U-Net++	0.7420	0.8228	0.8437	0.8607	0.8295	81.34	9.16	34.65
ResU-Net++	0.5341	0.6453	0.6964	0.7080	0.6576	43.11	4.06	15.81
HarDNet-MSEG	0.7459	0.8260	0.8485	0.8652	0.8358	34.80	33.34	6.02
ColonSegNet	0.6980	0.7920	0.8193	0.8432	0.7999	73.95	5.01	62.16
UACANet	0.7692	0.8502	0.8799	0.8706	0.8626	25.85	69.16	31.51
UNeXt	0.6284	0.7318	0.7840	0.7656	0.7507	87.47	1.47	0.57
TransNetR (Ours)	0.8016	0.8706	0.8843	0.9073	0.8744	54.60	27.27	10.58

• TransNetR은 Kvasir-SEG 테스트 세트에서 Dice 계수 0.8706 및 평균 IoU 0.8016, 리콜 0.8843 및 정밀도 0.9073를 54.60 FPS에서 달성했다.
• OOD 테스트에서 TransNetR은 PolypGen(6센터), BKAI-IGH, CVC-ClinicDB 데이터셋에서 최첨단 성능을 제공한다.
• ablation에 따르면 Residual Transformer(RT) 블록은 RT가 없는 변형에 비해 지표가 향상되며(예: +1.34% mIoU), 실시간 속도는 유지된다.
• 다수의 OOD 평가에서 TransNetR은 중심별 및 데이터세트별 분석에서 UACANet 및 UNeXt를 포함한 경쟁자들을 상회하는 경향을 보이며, mIoU 및 DSC에서 우수하다.
• 센터별 결과는 소형 및 다수의 폴립을 포함한 다양한 센터 데이터에서도 강건한 성능과 정확한 경계 구분을 나타낸다.

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