[논문 리뷰] MaskSAM: Towards Auto-prompt SAM with Mask Classification for Volumetric Medical Image Segmentation

Bin Xie, Hao Tang외 4명|arXiv (Cornell University)|2024. 03. 21.

Advanced Neural Network Applications인용 수 4

한 줄 요약

MaskSAM은 3D 체적 의료 영상 분할을 위한 프롬프트 불필요, 마스크 분류 인식 기반의 Segment Anything Model(SAM)의 변형을 제안한다. 학습 가능한 프롬프트 생성기로 보조 마스크, 박스, 분류자 토큰을 생성하고, SAM의 트랜스포머 블록에 3D 깊이 합성곱 및 깊이-MLP 어댑터를 통합함으로써, AMOS2022에서 SOTA 성능인 90.52%의 Dice 점수를 달성하여 nnUNet보다 2.7% 향상되었다.

ABSTRACT

Segment Anything Model (SAM), a prompt-driven foundation model for natural image segmentation, has demonstrated impressive zero-shot performance. However, SAM does not work when directly applied to medical image segmentation, since SAM lacks the ability to predict semantic labels, requires additional prompts, and presents suboptimal performance. Following the above issues, we propose MaskSAM, a novel mask classification prompt-free SAM adaptation framework for medical image segmentation. We design a prompt generator combined with the image encoder in SAM to generate a set of auxiliary classifier tokens, auxiliary binary masks, and auxiliary bounding boxes. Each pair of auxiliary mask and box prompts can solve the requirements of extra prompts. The semantic label prediction can be addressed by the sum of the auxiliary classifier tokens and the learnable global classifier tokens in the mask decoder of SAM. Meanwhile, we design a 3D depth-convolution adapter for image embeddings and a 3D depth-MLP adapter for prompt embeddings to efficiently fine-tune SAM. Our method achieves state-of-the-art performance on AMOS2022, 90.52% Dice, which improved by 2.7% compared to nnUNet. Our method surpasses nnUNet by 1.7% on ACDC and 1.0% on Synapse datasets.

연구 동기 및 목표

• SAM의 의료 영상 분할 한계를 해결할 것: 의미적 레이블 예측 부재 및 수동 프롬프트 의존성.
• 사전 학습된 2D SAM을 사용해 0-샷, 프롬프트 불필요 추론을 가능하게 할 것.
• 2D 자연 이미지와 3D 의료 볼륨 간 도메인 갭을 해소하기 위해 SAM의 아키텍처를 3D 공간 이해에 맞게 적응시킬 것.
• 사용자 제공 프롬프트 없이 정밀한 마스크와 해당 마스크의 의미적 레이블을 동시에 예측하는 통합 프레임워크를 제안할 것.
• 사용자 지정 태스크에 맞는 재학습 없이도 주요 3D 의료 분할 벤치마크에서 SOTA 성능을 달성할 것.

제안 방법

• SAM의 이미지 인코더에서 다중 수준 특징 맵으로부터 보조 이진 마스크, 경계 상자, 분류자 토큰을 생성하는 프롬프트 생성기를 설계할 것.
• 마스크 디코더에 학습 가능한 전역 분류자 토큰을 도입하여 보조 분류자 토큰과 합산함으로써 각 마스크에 대한 의미적 레이블을 예측할 것.
• 이미지 인코더의 각 트랜스포머 블록에 3D 깊이 합성곱 어댑터(DConvAdapter)를 통합하여 3D 인도크티브 바이어스를 가진 3D 이미지 임베딩을 처리할 것.
• 마스크 디코더에 3D 깊이-MLP 어댑터(DMLPAdapter)를 도입하여 3D 프롬프트 임베딩을 처리할 것. 스킵 연결을 갖춘 역-보틀넥 구조를 사용해 기능 정제를 향상시킬 것.
• 이미지 및 마스크 인코더 양쪽에서 3D 공간 모델링을 향상시키기 위해 깊이 위치 임베딩(DPosEmbed)을 적용할 것.
• 예측된 세그먼트와 진짜 세그먼트 간 이항 매칭을 통해 손실을 계산하여, 다양한 길이의 다중 클래스 마스크를 갖는 엔드 투 엔드 학습을 가능하게 할 것.

실험 결과

연구 질문

• RQ1 프롬프트 불필요한 SAM의 변형이 수동 포인트 또는 상자 프롬프트 없이도 고성능 3D 의료 영상 분할을 달성할 수 있는가?
• RQ2 SAM의 원래 이진 마스크 예측 전용 설계를 바탕으로, 분할된 마스크에 대한 의미적 레이블 예측을 어떻게 확장할 수 있는가?
• RQ3 2D 사전 학습된 SAM이 체적 의료 영상의 3D 공간 정보를 효과적으로 추출하고 활용할 수 있도록 하기 위해 어떤 아키텍처 수정이 필요한가?
• RQ4 3D 인식 어댑터(DConvAdapter 및 DMLPAdapter)는 표준 어댑터에 비해 3D 일반화 및 분할 정확도 향상에 얼마나 기여하는가?
• RQ5 학습 가능한 프롬프트와 분류자 토큰을 갖춘 통합형 엔드 투 엔드 프레임워크는 nnUNet과 같은 기존 SOTA 방법보다 3D 의료 분할 벤치마크에서 얼마나 뛰어난 성능을 내는가?

주요 결과

• MaskSAM은 AMOS2022 데이터셋에서 90.52%의 Dice 점수를 기록하여 nnUNet보다 2.7% 향상되었다.
• ACDC 데이터셋에서 MaskSAM은 Dice 점수로 nnUNet을 1.7% 초월하여 심장 MRI 분할 분야에서도 뛰어난 일반화 성능을 입증했다.
• Synapse 데이터셋에서 MaskSAM은 nnUNet 대비 1.0% 향상되어 다양한 해부학적 구조에 대한 강건성을 확인했다.
• 제거 실험 결과, 평균 경계 상자(MaskAvgBBoxPG)를 사용하는 제안된 프롬프트 생성기가 마스크 전용 생성기보다 1.9% 향상된 최고의 성능을 기록했다.
• DMLPAdapter와 DConvAdapter의 통합은 기준 모델 대비 0.3% 향상되었으며, DMLPAdapter는 다른 깊이-MLP 배치 방식에 비해 더 뛰어난 기능 정제 능력을 보였다.
• 깊이 위치 임베딩(DPosEmbed)의 추가로 기준 모델 대비 0.15% 이상의 성능 향상이 있었으며, 이는 3D 공간 모델링에서의 기여도를 입증했다.