[논문 리뷰] Prototype Learning for Out-of-Distribution Polyp Segmentation

Nikhil Kumar Tomar, Debesh Jha외 1명|arXiv (Cornell University)|2023. 08. 07.

Colorectal Cancer Screening and DetectionMedicine인용 수 3

한 줄 요약

이 논문은 다양한 내시경 영상 모odalities (WLI, BLI, LCI, FICE)에서 분포 외(out-of-distribution, OOD) 폴립 분할에 대해 강건성을 향상시키는 프로토타입 학습 기반 분할 모델인 PrototypeLab을 제안한다. Coarse Mask Generation Module (CMGM), Prototype Generation Module (PGM), 그리고 Prototype Mask Generation Module (PMGM)을 활용하여, 형태, 질감, 색상 특징을 포괄하는 클래스별 프로토타입을 학습한다. 이로 인해 OOD 데이터셋에서 DSC ≥ 90% 및 mIoU ≥ 85%의 성능을 달성하며, 거의 실시간 추론 속도를 확보한다.

ABSTRACT

Existing polyp segmentation models from colonoscopy images often fail to provide reliable segmentation results on datasets from different centers, limiting their applicability. Our objective in this study is to create a robust and well-generalized segmentation model named PrototypeLab that can assist in polyp segmentation. To achieve this, we incorporate various lighting modes such as White light imaging (WLI), Blue light imaging (BLI), Linked color imaging (LCI), and Flexible spectral imaging color enhancement (FICE) into our new segmentation model, that learns to create prototypes for each class of object present in the images. These prototypes represent the characteristic features of the objects, such as their shape, texture, color. Our model is designed to perform effectively on out-of-distribution (OOD) datasets from multiple centers. We first generate a coarse mask that is used to learn prototypes for the main object class, which are then employed to generate the final segmentation mask. By using prototypes to represent the main class, our approach handles the variability present in the medical images and generalize well to new data since prototype capture the underlying distribution of the data. PrototypeLab offers a promising solution with a dice coefficient of $\geq$ 90\% and mIoU $\geq$ 85\% with a near real-time processing speed for polyp segmentation. It achieved superior performance on OOD datasets compared to 16 state-of-the-art image segmentation architectures, potentially improving clinical outcomes. Codes are available at https://github.com/xxxxx/PrototypeLab.

연구 동기 및 목표

• 다른 임상 센터에서의 분포 외(OOD) 데이터셋에서 기존 폴립 분할 모델의 낮은 일반화 성능를 해결하기 위해.
• 작고 평평하거나 부분적으로만 노출되거나, 경계가 명확하지 않으며, 다양한 영상 조건에서 나타나는 폴립으로 인한 과제를 극복하기 위해.
• 다양한 영상 모달리티(WLI, BLI, LCI, FICE)와 다중 센터 데이터셋에서 잘 일반화되는 강건한 분할 모델을 개발하기 위해.
• 내재된 데이터 분포 특징을 포착하기 위해 프로토타입 학습을 통합하여 임상 환경에서의 모델 일반화 및 신뢰도를 향상시키기 위해.
• 실제 컴퓨터 보조 진단(CADx) 시스템에 구현 가능한 높은 정확도와 거의 실시간 추론 속도를 확보하기 위해.

제안 방법

• 입력 내시경 영상에서 다중 척도 특징을 추출하기 위해 피라미드 비전 트랜스포머(PVT) 인코더를 활용한다.
• 초기 거친 마스크를 생성하여 프로토타입 학습을 안내하기 위해 Coarse Mask Generation Module (CMGM)를 구현한다.
• 기본 데이터 분포를 포착하여 클래스당 다수의 프로토타입을 학습함으로써 특징 표현을 향상시키기 위해 Prototype Generation Module (PGM)를 설계한다.
• 다중 척도 특징 통합을 가능하게 하고 특징의 풍부함을 향상시키기 위해 Encoder Feature Fusion Module (EFFM)를 도입한다.
• 학습된 프로토타입을 디코더 특징과 융합하여 최종 분할 마스크를 생성하기 위해 Prototype Mask Generation Module (PMGM)를 활용한다.
• 내부 분포 및 OOD 데이터에서의 분할 성능 최적화를 위해 교차 엔트로피 손실과 Dice 손실을 사용하여 엔드 투 엔드 아키텍처를 훈련시킨다.

Figure 1 : The overall architecture of the proposed PrototypeLab. The input image is fed to the PVT-encoder to generate a coarse mask, which is combined with encoder features in the prototype generation module to generate various prototypes. Subsequently, the decoder is employed, which produces a fe

실험 결과

연구 질문

• RQ1 프로토타입 학습이 다른 임상 센터에서의 분포 외(OOD) 데이터셋에서 폴립 분할 모델의 일반화 성능 향상에 기여하는가?
• RQ2 제안된 PrototypeLab은 다양한 영상 모달리티에서 어려운 폴립 유형(소형, 평평하거나 부분적으로 노출된 병변 등)을 다루는 데 얼마나 효과적인가?
• RQ3 CMGM, PGM, PMGM 구성 요소가 기준 아키텍처 대비 분할 정확도 및 강건성 향상에 기여하는 정도는 어느 정도인가?
• RQ4 다중 척도 특징 융합( EFFM)과 프로토타입 기반 추론의 통합이 노이즈가 많거나 저대비 내시경 영상에서 성능 향상에 기여하는가?
• RQ5 임상 적용에 적합한 거의 실시간 추론 속도를 확보하면서도 높은 정확도를 유지할 수 있는가?

주요 결과

• BKAI-IGH OOD 데이터셋에서 PrototypeLab은 DSC 0.9243, mIoU 0.8744를 기록하여 16개의 최신 기술(SOTA) 방법보다 뚜렷이 뛰어난 성능을 보였다.
• Kvasir-SEG 데이터셋에서 PrototypeLab은 DSC 0.9086, mIoU 0.8544, HD 3.71를 달성했으며, SSFormer-L 대비 DSC 0.26% 향상되고 mIoU 0.44% 향상되었다.
• Kvasir-SEG에서 훈련하고 다중 센터 PolypGen 데이터셋에서 테스트한 결과, PrototypeLab은 DSC 0.7583, mIoU 0.6957, F2 0.7563을 기록했으며, Polyp-PVT 및 SSFormer-L보다 우수한 성능을 보였다.
• 절단 실험 결과, CMGM, PGM, PMGM, EFFM, LKDC 블록을 모두 포함한 전체 아키텍처는 기준 모델 대비 DSC 1.15% 향상되고 mIoU 1.45% 향상됨을 확인하였다.
• PrototypeLab는 23.85 FPS(약 30 FPS)로 이미지를 처리하여 거의 실시간 추론 속도를 확보하면서도 고정확도를 유지함을 입증하였다.
• 정성적 결과에서는 DuAT 및 Polyp-PVT와 같은 최신 기술 기반 모델 대비 소형, 평평하거나 노이즈가 많은 폴립에서 과도 및 부족 분할을 줄이는 데 성공하였다.

Figure 2 : Qualitative results of models trained BKAI-IGH and tested on PolypGen. It can be observed that PrototypeLab produces a more accurate segmentation map in all the centers from C1 to C6.