[논문 리뷰] Enhancing the Reliability of Segment Anything Model for Auto-Prompting Medical Image Segmentation with Uncertainty Rectification

Yichi Zhang, Shiyao Hu외 4명|arXiv (Cornell University)|2023. 11. 17.

Brain Tumor Detection and Classification인용 수 8

한 줄 요약

UR-SAM은 SAM을 프롬프트 증강과 불확실성 보정 모듈로 확장하여 추가 학습 없이도 신뢰할 수 있는 자동 프롬프트 의학 영상 세분화를 가능하게 한다. 3D HaN 및 복부 데이터셋에서 Dice 점수의 주목할 만한 향상을 달성한다.

ABSTRACT

The Segment Anything Model (SAM) has recently emerged as a groundbreaking foundation model for prompt-driven image segmentation tasks. However, both the original SAM and its medical variants require slice-by-slice manual prompting of target structures, which directly increase the burden for applications. Despite attempts of auto-prompting to turn SAM into a fully automatic manner, it still exhibits subpar performance and lacks of reliability especially in the field of medical imaging. In this paper, we propose UR-SAM, an uncertainty rectified SAM framework to enhance the reliability for auto-prompting medical image segmentation. Building upon a localization framework for automatic prompt generation, our method incorporates a prompt augmentation module to obtain a series of input prompts for SAM for uncertainty estimation and an uncertainty-based rectification module to further utilize the distribution of estimated uncertainty to improve the segmentation performance. Extensive experiments on two public 3D medical datasets covering the segmentation of 35 organs demonstrate that without supplementary training or fine-tuning, our method further improves the segmentation performance with up to 10.7 % and 13.8 % in dice similarity coefficient, demonstrating efficiency and broad capabilities for medical image segmentation without manual prompting.

연구 동기 및 목표

• 신뢰할 수 있는 자동 의료 영상 세분화를 프롬프트 기반 모델로 가능하게 하려는 동기 부여.
• 높은 애매함을 가진 세부분 영역을 식별하고 보정하기 위해 불확실성 추정 활용.
• 의료 데이터에서 SAM의 성능을 개선하기 위한 프롬프트 증강 및 클래스 특정 신뢰도 메커니즘 개발.
• 머리-목 및 복부 장기를 다루는 3D 의료 데이터셋에서 프레임워크를 검증.
• 미세 조정 없이 기본 자동 프론핑 대비 우월성을 입증하고 수동 프롬핑에 근접한 성능 달성.

제안 방법

• SAM을 위한 3D 바운딩 박스 프롬프트를 자동 생성하는 랜드마크 로컬라이제이션 모듈 사용.
• 초기 프롬프트를 섭동으로 증강하여 다수의 프롬프트를 생성하고 분포형 세분 결과를 도출.
• 다중 프롬프트 출력으로부터 예측 엔트로피를 사용해 픽셀 수준의 불확실성 추정.
• 높은 불확실 영역을 보정하기 위해 클래스 특이적 신뢰도 기반 임계값 적용.
• 이미지 내 강도 관계를 활용하여 불확실한 영역을 재할당하고 세분화를 보정.
• StructSeg(HaN 장기)와 FLARE 22(복부 장기)에서 SAM 및 MedSAM 백본으로 UR-SAM 평가.

실험 결과

연구 질문

• RQ1 추가 학습 없이도 프롬프트 증강과 불확실성 보정이 의료 영상 세분화의 SAM 기반 자동 프롬프팅을 향상시킬 수 있는가?
• RQ2 의료 데이터세트에서 고정 임계값 보정 전략과 비교했을 때 클래스 특이적 불확실성 기반 보정의 성능은 어떠한가?
• RQ3 UR-SAM이 3D CT 데이터의 작고 큰/복잡한 장기에 대한 세분 정확도에 미치는 영향은?
• RQ4 UR-SAM이 자동 프롬핑의 성능을 수동 프롬핑에 더 가깝게 만들까?
• RQ5 UR-SAM에서 SAM 대 MedSAM 중 어떤 기반 모델 사용이 효과적인가?

주요 결과

• UR-SAM은 3D 의료 데이터셋에서 기본 자동 프롬핑 대비 세분화 성능을 향상시키며, 머리-목 세분화에서 Dice가 최대 10.7%, 복부 세분화에서 최대 13.8% 향상을 달성한다.
• 섬세한 섭동으로 프롬프트를 증강하면 강건성 향상이 나타나지만, 증강 수가 많아지거나 과도한 섭동은 성능에 부정적일 수 있다.
• 클래스 특이적 불확실성 기반 필터링은 고정 임계값 접근법 및 앙상블 방법보다 DSC를 일관되게 개선한다.
• UR-SAM은 SAM 및 MedSAM 백본에서 자동 프롬핑을 능가하며, 여러 대상에서 수동 프롬핑 성능에 근접할 수 있다.
• 보정된 세분 결과는 앙상블 결과에 비해 오차를 줄이는 경향을 보이며, 불확실성 기반 보정의 효과를 시사한다.
• MedSAM은 미리 보지 못한 데이터셋에서 SAM에 비해 성능이 떨어질 수 있어 UR-SAM의 불확실성 보정 전략의 가치를 강화한다.