[논문 리뷰] Contrastive Pretraining for Echocardiography Segmentation with Limited Data

M. A. Saeed, Rand Muhtaseb외 1명|arXiv (Cornell University)|2022. 01. 16.

Cardiac Valve Diseases and Treatments인용 수 5

한 줄 요약

이 논문은 제한된 애너테이션 데이터에서 심장초음파 영상의 좌심실 분할을 향상시키기 위해 SimCLR와 BYOL을 사용한 자기지도 대비 미리학습 방법을 제안한다. 대규모 레이블이 없는 초음파 영상에서 미리학습하고 ImageNet 또는 EchoNet-Dynamic 데이터의 5%만을 사용해 피니튜닝함으로써, EchoNet-Dynamic에서 Dice 점수 0.9252를 달성하여 완전히 지도 학습된 최신 기술을 능가한다.

ABSTRACT

Contrastive learning has proven useful in many applications where access to labelled data is limited. The lack of annotated data is particularly problematic in medical image segmentation as it is difficult to have clinical experts manually annotate large volumes of data such as cardiac structures in ultrasound images of the heart. In this paper, We propose a self supervised contrastive learning method to segment the left ventricle from echocardiography when limited annotated images exist. Furthermore, we study the effect of contrastive pretraining on two well-known segmentation networks, UNet and DeepLabV3. Our results show that contrastive pretraining helps improve the performance on left ventricle segmentation, particularly when annotated data is scarce. We show how to achieve comparable results to state-of-the-art fully supervised algorithms when we train our models in a self-supervised fashion followed by fine-tuning on just 5\% of the data. We show that our solution outperforms what is currently published on a large public dataset (EchoNet-Dynamic) achieving a Dice score of 0.9252. We also compare the performance of our solution on another smaller dataset (CAMUS) to demonstrate the generalizability of our proposed solution. The code is available at (https://github.com/BioMedIA-MBZUAI/contrastive-echo).

연구 동기 및 목표

• 임상 환경에서 딥러닝 도입을 저해하는 좌심실 분할을 위한 제한된 애너테이션 데이터 문제를 해결한다.
• 레이블 데이터가 부족할 경우 자기지도 대비 학습이 분할 성능을 향상시킬 수 있는지 조사한다.
• 두 가지 기존 네트워크인 UNet과 DeepLabV3를 사용하여 SimCLR 및 BYOL 대비 학습 프레임워크의 심장초음파 영상 분할 성능에 대한 효과를 평가한다.
• EchoNet-Dynamic와 더 작은 CAMUS 데이터셋 모두에서 테스트하여 모델의 일반화 능력을 입증한다.
• 미리학습 데이터 크기와 백본 초기화 방식이 최종 분할 성능에 미치는 영향을 탐색한다.

제안 방법

• EchoNet-Dynamic 및 CAMUS의 대규모 레이블이 없는 심장초음파 영상에서 SimCLR 및 BYOL을 사용해 특징 추출기의 미리학습을 수행한다.
• 랜덤 크롭, 색상 왜곡, 가우시안 블러와 같은 데이터 증강 기법을 적용하여 동일한 이미지의 증강된 쌍(양성 쌍)과 다른 이미지의 증강된 쌍(음성 쌍)을 생성한다.
• UNet 및 DeepLabV3 아키텍처를 사용해 최종 분할 작업에서 레이블 데이터의 5%만을 사용해 미리학습된 백본을 피니튜닝한다.
• 대비 손실을 사용하여 동일한 이미지의 증강된 시각 간의 일치를 최대화하고, 다른 이미지의 증강된 시각 간의 일치를 최소화한다.
• 다양한 데이터 제약 조건(100% 및 5% 레이블 데이터)에서 EchoNet-Dynamic 및 CAMUS 데이터셋에서 Dice 점수를 사용해 성능을 평가한다.
• 미리학습 전략을 비교한다: ImageNet, 무작위 초기화, EchoNet-Dynamic 또는 CAMUS에서의 대비 미리학습.

실험 결과

연구 질문

• RQ1 레이블 데이터가 제한적으로 존재할 경우 자기지도 대비 미리학습이 좌심실 분할 성능을 향상시킬 수 있는가?
• RQ2 레이블 데이터를 최소로 사용해 피니튜닝할 경우 SimCLR가 BYOL보다 심장초음파 영상 분할에서 성능이 뛰어나게 되는가?
• RQ3 대규모이고 다양한 데이터셋인 EchoNet-Dynamic에서의 미리학습과 더 작은 도메인 특화 데이터셋인 CAMUS에서의 미리학습 간 성능 차이는 어떻게 되는가?
• RQ4 레이블 데이터의 5%만을 사용해 피니튜닝한 모델이 완전히 지도 학습된 최신 기술 모델과 유사한 성능을 달성할 수 있는가?
• RQ5 분할 백본 선택(UNet 대비 DeepLabV3)이 대비 미리학습에서 얻는 이점에 영향을 미치는가?

주요 결과

• EchoNet-Dynamic에서 SimCLR를 사용한 대비 미리학습은 EchoNet-Dynamic 데이터셋에서 Dice 점수 0.9252를 달성하여 완전히 지도 학습된 최신 기술을 능가했다.
• SimCLR 미리학습 후 레이블 데이터의 5%만을 사용해 피니튜닝한 결과, 완전히 지도 학습된 모델과 유사한 성능을 보였으며, 높은 데이터 효율성을 입증했다.
• EchoNet-Dynamic에서 SimCLR 미리학습은 CAMUS에서의 성능을 향상시켜 Dice 점수 0.9311을 달성했으며, ImageNet 및 무작위 초기화를 초월했다.
• BYOL 미리학습은 특히 5% 데이터에서 피니튜닝할 경우 성능이 열 劣하였으며, Dice 점수 0.7256에 머물러 안정성과 일반화 능력이 떨어지는 것으로 나타났다.
• CAMUS(400개 샘플)와 같이 작은 데이터셋에서의 미리학습은 SimCLR 및 BYOL 모두 효과적이지 않았으며, 대규모 미리학습 데이터의 필요성을 강조했다.
• 모든 설정에서 DeepLabV3가 UNet을 능가했으며, 특히 SimCLR 미리학습과 조합했을 경우 ASPP 모듈 덕분에 다중 척도 특징 학습 능력이 향상되어 성능이 뛰어났다.