[논문 리뷰] SAM3D: Segment Anything Model in Volumetric Medical Images

Nhat-Tan Bui, Dinh-Hieu Hoang외 6명|arXiv (Cornell University)|2023. 09. 07.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 3

한 줄 요약

SAM3D는 2D 슬라이스를 먼저 처리하고 나서 3D 디코더를 적용하여 깊이 방향의 관계를 포착함으로써, 사전 훈련된 SAM 인코더를 볼륨 데이터에 활용하는 경량이고 효율적인 3D 의료 영상 분할 모델을 소개한다. 오직 188만 개의 파라미터로도 최신 기술 수준(SOTA) 성능을 달성하며, 다른 SAM 기반 및 SOTA 모델보다 분할 정확도에서 뛰어나면서 모델 복잡도를 크게 감소시킨다.

ABSTRACT

Image segmentation remains a pivotal component in medical image analysis, aiding in the extraction of critical information for precise diagnostic practices. With the advent of deep learning, automated image segmentation methods have risen to prominence, showcasing exceptional proficiency in processing medical imagery. Motivated by the Segment Anything Model (SAM)-a foundational model renowned for its remarkable precision and robust generalization capabilities in segmenting 2D natural images-we introduce SAM3D, an innovative adaptation tailored for 3D volumetric medical image analysis. Unlike current SAM-based methods that segment volumetric data by converting the volume into separate 2D slices for individual analysis, our SAM3D model processes the entire 3D volume image in a unified approach. Extensive experiments are conducted on multiple medical image datasets to demonstrate that our network attains competitive results compared with other state-of-the-art methods in 3D medical segmentation tasks while being significantly efficient in terms of parameters. Code and checkpoints are available at https://github.com/UARK-AICV/SAM3D.

연구 동기 및 목표

• 기존 3D 의료 분할 모델이 복잡한 아키텍처와 높은 파라미터 수에 의존하는 데서 비롯되는 한계를 해결하기 위해.
• 전체 모델을 미세조정하지 않고도 Segment Anything Model(SAM)의 능력을 볼륨 영상에 확장하기 위해.
• 각 슬라이스를 독립적으로 처리하는 것이 아니라 슬라이스 간의 3D 공간적 관계를 포착함으로써 분할 성능을 향상시키기 위해.
• 높은 정확도를 유지하면서도 계산 비용이 낮고 파라미터 수가 적은 프레임워크를 개발하기 위해.

제안 방법

• 자연 이미지에서 가장 낮은 수준의 특징(예: 윤곽선, 경계선)을 유지하기 위해 사전 훈련된 SAM 이미지 인코더(ViT-B)를 동결한다.
• 3D 볼륨의 각 2D 슬라이스를 동결된 SAM 인코더를 통해 독립적으로 처리하여 $ \frac{H}{16} \times \frac{W}{16} \times D \times 256 $ 크기의 3D 슬라이스 임베딩을 생성한다.
• 슬라이스 임베딩 간의 깊이 방향 관계를 모델링하기 위해 가벼운 3D 컨볼루션 디코더를 적용하여 3D 맥락 이해를 가능하게 한다.
• 모달리티에 특화된 편향을 방지하고 다양한 의료 영상 모달리티 간의 일반화를 확보하기 위해 SAM에서 프ompt 인코더를 제거한다.
• SAM 인코더를 동결한 채로 오직 3D 디코더 헤드만 엔드 투 엔드로 훈련함으로써 파라미터 업데이트와 계산 비용을 최소화한다.
• 세분화된 특징을 유지하고 분할에서 가장자리 및 경계선의 정확도를 향상시키기 위해 3D 디코더에 스킵 연결을 통합한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 2D 기초 모델인 SAM이 전체 미세조정 없이도 3D 볼륨 영상 분할에 효과적으로 적용될 수 있는가?
• RQ2 3D 디코더를 통해 슬라이스 간 상호 관계를 포착하면 슬라이스별로 처리하는 것보다 분할 성능이 향상되는가?
• RQ3 가벼운 3D 디코더가 기존 SOTA 3D 분할 모델보다 훨씬 적은 파라미터로 경쟁 가능한 성능을 달성할 수 있는가?
• RQ4 다양한 모달리티를 가진 3D 의료 데이터셋에서 SAM3D는 다른 SAM 기반 및 Transformer/CNN 하이브리드 모델과 비교해 성능과 효율성 면에서 어떻게 성과를 내는가?

주요 결과

• Synapse 데이터셋에서 SAM3D는 188만 개의 파라미터를 사용함에도 불구하고 SAMed_s(632만 개 파라미터)보다 DSC 1.78% 높은 77.78%를 기록했다.
• ACDC 심장 데이터셋에서 SAM3D는 파라미터 수가 1/50 미만인 TransUNet보다 DSC 0.41% 높게 기록했다.
• BraTS 뇌 종양 데이터셋에서 SAM3D는 파라미터 수가 1/20 미만인 UNETR보다 DSC 1.8% 높게 기록했다.
• 제거 분석 결과, 3D 디코더 내 스킵 연결이 성능 향상에 기여함을 확인하였으며, 이는 경계 세부 정보 유지에 중요한 역할을 한다는 것을 시사한다.
• SAM3D는 심장, 뇌, 폐 분할 작업을 포함해 다양한 모달리티와 해부학적 구조에 대해 강력한 일반화 성능을 보였다.
• 가장 작은 SAM 버전(ViT-B)을 사용하더라도 성능이 경쟁 가능했으며, 이는 더 큰 백본 버전(ViT-L, ViT-H 등)을 사용할 경우 향후 성능 향상 가능성을 시사한다.