[논문 리뷰] Shallow Attention Network for Polyp Segmentation

Jun Wei, Yiwen Hu외 4명|arXiv (Cornell University)|2021. 08. 02.

Advanced Neural Network Applications참고 문헌 22인용 수 9

한 줄 요약

이 논문은 색상 변동, 소형 폴립 열화, 배경-전경 클래스 불균형 문제를 해결하기 위해 폴립 세그멘테이션을 위한 얕은 주의망(SANet)을 제안한다. 색상 교환을 통해 콘텐츠와 색상을 분리하고, 고해상도 특징을 향상시키기 위한 얕은 주의 모듈을 도입하며, 추론 시 확률 보정 전략을 적용함으로써 SANet는 다섯 가지 벤치마크에서 72 FPS의 추론 속도로 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Accurate polyp segmentation is of great importance for colorectal cancer diagnosis. However, even with a powerful deep neural network, there still exists three big challenges that impede the development of polyp segmentation. (i) Samples collected under different conditions show inconsistent colors, causing the feature distribution gap and overfitting issue; (ii) Due to repeated feature downsampling, small polyps are easily degraded; (iii) Foreground and background pixels are imbalanced, leading to a biased training. To address the above issues, we propose the Shallow Attention Network (SANet) for polyp segmentation. Specifically, to eliminate the effects of color, we design the color exchange operation to decouple the image contents and colors, and force the model to focus more on the target shape and structure. Furthermore, to enhance the segmentation quality of small polyps, we propose the shallow attention module to filter out the background noise of shallow features. Thanks to the high resolution of shallow features, small polyps can be preserved correctly. In addition, to ease the severe pixel imbalance for small polyps, we propose a probability correction strategy (PCS) during the inference phase. Note that even though PCS is not involved in the training phase, it can still work well on a biased model and consistently improve the segmentation performance. Quantitative and qualitative experimental results on five challenging benchmarks confirm that our proposed SANet outperforms previous state-of-the-art methods by a large margin and achieves a speed about 72FPS.

연구 동기 및 목표

• 다양한 내 endoscopic 조건으로 인한 폴립 영상의 색상 분포 불일치로 인한 과적합 문제를 해결하기 위해.
• 심층 네트워크에서 반복적인 특징 다운샘플링으로 인해 열화되는 소형 폴립의 세그멘테이션 성능 향상하기 위해.
• 모델 훈련을 왜곡시키고 소형 폴립 성능을 떨어뜨리는 심각한 전경-배경 픽셀 불균형을 완화하기 위해.
• 재학습이 필요 없이도 높은 정확도를 유지하면서도 경량이며 빠른 추론 속도를 갖춘 모델 개발하기 위해.
• 최소한의 계산 부담으로 다수의 공개 폴립 세그멘테이션 벤치마크에서 최신 기술 수준 성능 달성하기 위해.

제안 방법

• 입력 이미지에 임의의 이미지의 색상을 전달하는 색상 교환(CE) 연산을 제안하여 콘텐츠와 색상을 분리하고, 비합리적인 색상 기반 상관관계를 감소시킨다.
• 깊은 특징을 사용해 고해상도 얕은 특징에서 배경 노이즈를 걸러내는 얕은 주의 모듈(SAM)을 도입하여 소형 폴립의 세밀한 세부 사항을 유지한다.
• 훈련 과정을 수정하지 않고도 편향된 예측을 재보정하기 위해 추론 중 확률 보정 전략(PCS)을 적용한다.
• 잔차 연결과 다중 척도 감시를 사용한 U-Net 유사 인코더-디코더 아키텍처를 사용하여 강력한 특징 학습을 가능하게 한다.
• 표준 데이터 증강(임의의 플립, 회전, 다중 척도 훈련)과 BCE + Dice 손실을 사용한 엔드 투 엔드 훈련을 적용하여 최적화한다.
• 소형 폴립 세그멘테이션에 필수적인 공간적 세부 정보를 유지하기 위해 고해상도 얕은 특징을 활용한다.

실험 결과

연구 질문

• RQ1 이미지 콘텐츠를 색상 분포에서 분리함으로써 폴립 세그멘테이션 모델의 과적합을 줄일 수 있는가?
• RQ2 주의 메커니즘을 사용해 얕은 특징을 향상시키면 소형 폴립의 세그멘테이션 성능이 향상되는가?
• RQ3 재학습 없이도 후처리 보정 전략이 편향된 모델의 세그멘테이션 성능을 효과적으로 향상시킬 수 있는가?
• RQ4 색상 교환, 얕은 주의, 확률 보정을 조합하면 다양한 데이터셋에서 일관된 성능 향상을 이끌 수 있는가?
• RQ5 제안된 방법이 실세계 폴립 세그멘테이션 벤치마크에서 고속 추론(예: 72 FPS)을 유지하면서 최신 기술 수준 성능을 달성할 수 있는가?

주요 결과

• Kvasir 데이터셋에서 SANet는 0.904의 mDice와 0.847의 mIoU를 기록하여 여섯 가지 SOTA 방법을 모두 초월한다.
• CVC-ClinicDB에서 SANet는 mDice 0.916과 mIoU 0.859를 달성하여 PraNet(0.898 및 0.840)을 크게 앞서며 성능을 뛰어넘는다.
• ETIS에서 SANet는 mDice 0.750과 mIoU 0.654를 기록하여 이전 SOTA인 PraNet(0.628 및 0.567)을 능가한다.
• 절단 실험 결과, 백본, SAM, CE, PCS 모두가 필수적임을 확인하였으며, 전체 SANet는 ColonDB에서 0.753 mDice, Kvasir에서 0.904 mDice를 달성한다.
• RTX 2080Ti에서 SANet는 실시간 추론 속도 72 FPS를 기록하여 PraNet의 64 FPS보다 더 빠르며, 높은 효율성을 입증한다.
• 시각적 비교 결과, SANet는 특히 소형 및 도전적인 폴립에서 더 선명하고 정확한 마스크를 생성하며 배경 억제 성능도 뛰어나다.