Segment anything in medical images 리뷰

https://arxiv.org/abs/2304.12306

Segment Anything in Medical Images

https://github.com/bowang-lab/MedSAM

GitHub - bowang-lab/MedSAM: Segment Anything in Medical Images

 

Medical image segmentation 분야의 foundation model을 만들기 위해 시도한 논문입니다. 기존 SAM 대비 굉장히 높은 성능을 보여주고, medical segmentation의 후속 연구를 위한 초석을 마련했다 해도 과언이 아닌 논문입니다.

 

Abstract

Medical 분야에서는 segmentation은 정확한 진단과 질병 모니터링을 용이하게 하기 위한 핵심적인 요소다. 그러나, 일반적은 segmentation은 질병의 종류가 다양한 medical 분야에 적용시키기에 일반화 성능이 부족한 문제가 있다.

 

그래서 본 논문에서는 Medical image segmentation에 특화된 Foundation model로 MedSAM을 소개하고 있다. MedSAM은 다양한 내외부 dataset을 통해 일반화 성능이 검증되었고, 다양한 의료 작업에서 강건한 성능을 보여준다.

 

Introduction

Deep learning-based models은 medical image segmentation에서 큰 발전을 이루고 있다. 그러나 현존하는 model의 가장 큰 문제점은 각각의 구체적인 task에 특화되어 있어, 이를 다양한 task에 일반화 하기란 굉장히 어렵다는 점이다.

 

최근 Foundation segmentation model인 SAM의 등장으로 이러한 문제점을 해결할 가능성을 충분히 보여줬다. 그러나, medical 분야와 natural 분야의 domain 차이로 인해 SAM만을 이용해서는 기존의 문제점을 해결할 수 없었다. 또한 SAM은 뚜렷한 경계가 존재하는 image에서 성능이 탁월하나, medical 분야에서는 그 경계가 모호하기 때문에 다소 낮은 성능을 보였다.

 

따라서 MedSAM은 이러한 domain gap을 대규모의 dataset을 통한 학습으로 SOTA model보다 강건하면서 더 높은 성능을 보여준다.

 

Method

Dataset curation and pre-processing

• public으로 이용가능한 dataset을 활용 → kaggle, TCIA 등 appendix에서 확인 가능
• image format 통합
  • 3D dataset(CT, MR) → Nifti format
  • grayscale, RGB dataset → PNG format
• modality별 intensity가 모두 다르기 때문에, 이를 고려한 intensity normalization
  • CT : -2000 ~ 2000, MR : 0 ~ 3000, 그 외 : 0 ~ 255
• input size 1024 x 1024 x 3으로 통일
  • WSI : patch를 overlap 하지 않고 추출, 경계는 0으로 처리
  • 3D(CT, MR) images : 2D slice 후, 1024 x 1024 resize, → 채널별 3번 반복 repeat
  • image : Bi-cubic interpolation
  • mask : nearest-neighbor interpolation → 경계 유지

이러한 과정을 pipeline으로 만들어 후속 단계에 원활하게 적용

 

Network architecture

SAM과 동일하게 ViT 기반의 Encoder, mask decoder 사용했다. 저자는 성능과 비용간 균형을 맞추기 위해 ViT Large, Huge를 사용해 봤으나, 약간의 성능 향상만 발생했다. 따라서 기본 ViT를 사용했다.

• Base ViT
  • 12개의 Transformer layers
  • 각각의 layer는 Multi-head SA, layer norm MLP로 구성

Training protocol & Experimental setting

• 1,570,263 image-mask fair dataset 구축
• Train, val, test : 8 : 1 : 1
  • Internal validation : test dataset
  • external validation : 일반화 성능을 위해 학습에 사용하지 않은 datasets
• WSI는 patch(1024 x 1024 x 3) split 후, overfitting을 막기 위해, parameters를 변경해 가며 training
• 총 3개의 model과 비교실험 진행
  • U-Net, DeepLabV3, SAM
  • 총 10개의 modality에서 training

Loss function

Metrics Reloaded에서 권장하는 Dice Similarity Coefficinet(DSC) & Normalized Surface Distance(NSD)를 사용

• DSC : region-based segmentation metric = dice score
• NSD : boundary-based metric, 두 boundary의 일치도를 평가

 

Result

• MedSAM은 medical image segmentation의 foundation model 급의 성능을 보여줌
• 광범위한 task에서 좋은 성능을 보였으나, Training dataset의 불균형이 존재

Quntitative & qualitative analysis

• MedSAM은 SOTA model인 SAM, U-net, DeepLabV3와 비교
• 가장 bias와 varience가 낮은 model이 MedSAM
• Internal, external을 통한 validation에서도 강건한 성능

The effect of training dataset size

• 굉장히 대량의 dataset으로 학습 진행
• dataset size에 따른 성능 비교실험 → 100k에서 가장 좋은 성능

MedSAM can improve the annotation effciency

• 기존에 annotation은 전문가가 직접 진행 → 시간이 오래 걸림
• 전문가, MedSAM을 이용하여 annotation 작업 시간을 82% 단축, 성능 우수

 

limitation

• modality training set에서 불균형이 존재함
• 혈관 구조같이 분기되는 구조에서 bbox를 효과적으로 찾지 못함