[논문 리뷰] Slide-SAM: Medical SAM Meets Sliding Window

Quan Quan, Fenghe Tang외 3명|arXiv (Cornell University)|2023. 11. 16.

AI in cancer detection인용 수 4

한 줄 요약

Slide-SAM는 사전 훈련된 Segment Anything Model(SAM)을 사용하여 효율적이고 다중 슬라이스 3D 의료 영상 분할을 가능하게 하는 슬라이딩 윈도우 접근법을 제안한다. 중심 슬라이스에 대한 단일 프롬프트만 필요로 하며, 사전 훈련된 가중치를 활용하고 슈퍼픽셀 기반 프롬프트를 통한 가짜 레이블 생성 및 3D와 2D 레이블을 결합한 하이브리드 손실 함수를 사용함으로써, 최소한의 프롬프트로 최신 기술 수준의 성능을 달성한다. 이는 소수의 프롬프트 설정에서 2D SAM 기반 모델들을 능가한다.

ABSTRACT

The Segment Anything Model (SAM) has achieved a notable success in two-dimensional image segmentation in natural images. However, the substantial gap between medical and natural images hinders its direct application to medical image segmentation tasks. Particularly in 3D medical images, SAM struggles to learn contextual relationships between slices, limiting its practical applicability. Moreover, applying 2D SAM to 3D images requires prompting the entire volume, which is time- and label-consuming. To address these problems, we propose Slide-SAM, which treats a stack of three adjacent slices as a prediction window. It firstly takes three slices from a 3D volume and point- or bounding box prompts on the central slice as inputs to predict segmentation masks for all three slices. Subsequently, the masks of the top and bottom slices are then used to generate new prompts for adjacent slices. Finally, step-wise prediction can be achieved by sliding the prediction window forward or backward through the entire volume. Our model is trained on multiple public and private medical datasets and demonstrates its effectiveness through extensive 3D segmetnation experiments, with the help of minimal prompts. Code is available at \url{https://github.com/Curli-quan/Slide-SAM}.

연구 동기 및 목표

• 2D SAM을 3D 의료 볼륨에 적용할 때 슬라이스 단위 추론으로 인한 불연속성과 높은 주석 비용 문제를 해결하기 위해.
• 단일 초기 프롬프트만을 사용하면서도 SAM의 사전 훈련된 가중치를 유지한 채 효과적인 3D 분할을 가능하게 하기 위해.
• 임계치 범위와 슈퍼픽셀 기반 프롬프트에서 유래한 풍부한 2D 가짜 레이블을 생성함으로써 데이터 효율성을 향상시키기 위해.
• 3D 진짜 레이블과 2D 가짜 레이블을 조합한 하이브리드 손실을 통해 훈련 및 추론의 안정성을 향상시키기 위해.
• 2D SAM과의 앙상블 학습을 통해 상호 보완적인 강점을 활용해 성능을 향상시키기 위해.

제안 방법

• Slide-SAM는 세 개의 연속된 슬라이스를 윈도우로 처리하며, 중심 슬라이스에 프롬프트를 적용하여 세 슬라이스 모두에 마스크를 예측한다.
• 상단 및 하단 슬라이스의 예측 결과를 기반으로 다음 윈도우의 새로운 프롬프트를 생성함으로써 볼륨 전체를 순차적이고 단계적으로 추론할 수 있도록 한다.
• 의료 영상에서 다양한 임계치 범위에 걸쳐 슈퍼픽셀 기반 프롬프트를 적용하여 SAM을 통해 가짜 레이블을 생성한다.
• 하이브리드 손실 함수는 훈련 중 3D 진짜 레이블과 2D 가짜 레이블을 선택적으로 적용하여 단일 슬라이스 및 다중 슬라이스 지도 학습을 동시에 최적화한다.
• 2D SAM과의 앙상블 전략은 안정성 점수와 IoU 임계치 기반으로 예측 결과를 필터링하고 재평가함으로써 성능을 향상시킨다.
• Slide-SAM는 SAM의 사전 훈련된 시각 인코더 및 디코더를 유지하면서도 3D 컨텍스트 모델링을 위한 경량 어댑터를 도입한다.

실험 결과

연구 질문

• RQ1 초기 훈련 없이 2D 사전 훈련된 SAM을 3D 의료 영상 분할에 효과적으로 적응시킬 수 있는가?
• RQ2 최소한의 인간 프롬프트를 어떻게 활용하여 전체 볼륨에 걸쳐 고품질의 3D 분할을 달성할 수 있는가?
• RQ3 슈퍼픽셀 기반 프롬프트에서 유래한 가짜 레이블이 저자료 환경에서 모델 일반화 능력을 향상시킬 수 있는가?
• RQ4 Slide-SAM와 2D SAM의 통합이 분할 정확도 및 안정성을 어떻게 향상시키는가?
• RQ5 슬라이딩 윈도우 메커니즘이 공간 일관성을 유지하면서도 주석 비용을 얼마나 줄일 수 있는가?

주요 결과

• Slide-SAM는 AbdomenCT-1K 및 CHAOS 데이터셋에서 해부학적 구조당 평균 5개의 프롬프트만으로도 경쟁적인 3D 분할 성능를 달성하며, 2D SAM 기반 모델들이 평균 약 40개의 프롬프트가 필요한 것에 비해 크게 감소시켰다.
• 2D SAM 방법 대비 90% Dice 점수에 도달하기 위해 필요한 프롬프트 수를 50% 이상 감소시켜, 훨씬 뛰어난 소수 샘플 효율성을 입증했다.
• 시각적 결과에서는 슬라이스 간에 매끄럽고 일관된 분할이 이루어져, 슬라이스 단위로 추론하는 기존 SAM 방식에서 흔히 발생하는 불연속성을 제거했다.
• Slide-SAM와 2D SAM의 앙상블은 모든 해부학적 구조에서 성능 향상을 보였으며, 여러 테스트 세트에서 Dice 점수 향상이 두드러졌다.
• 슈퍼픽셀 기반 프롬프트와 SAM 추론을 통해 생성된 가짜 레이블은 실재 레이블만으로 미세조정하는 것보다 성능 향상에 크게 기여했다.
• 특히 이동 및 스케일링 변형에 노이즈가 포함된 상자 프롬프트에 대해 강건성을 보이며, 정확도 유지에 기여했다.