[논문 리뷰] Chest X-ray lung and heart segmentation based on minimal training sets

Balázs Maga|arXiv (Cornell University)|2021. 01. 20.

COVID-19 diagnosis using AI참고 문헌 16인용 수 4

한 줄 요약

이 논문은 최소한의 훈련 데이터로 흉부 X선에서 폐와 심장을 정확하게 분할하기 위해 이중 방향 ConvLSTM과 주의 메커니즘을 통합한 U-Net 기반 아키텍처인 Attention BCDU-Net을 제안한다. JSRT 데이터셋에서 최신 기술 수준의 성능을 기록하며, 98.1%의 Dice 점수와 95.2%의 IoU를 달성했고, 단지 10개의 훈련 샘플로도 97.0%의 Dice와 92.2%의 IoU를 기록하여, mixup 데이터 증강 기법으로 인해 크게 향상되었다.

ABSTRACT

As the COVID-19 pandemic aggravated the excessive workload of doctors globally, the demand for computer aided methods in medical imaging analysis increased even further. Such tools can result in more robust diagnostic pipelines which are less prone to human errors. In our paper, we present a deep neural network to which we refer to as Attention BCDU-Net, and apply it to the task of lung and heart segmentation from chest X-ray (CXR) images, a basic but ardous step in the diagnostic pipeline, for instance for the detection of cardiomegaly. We show that the fine-tuned model exceeds previous state-of-the-art results, reaching $98.1\pm 0.1\%$ Dice score and $95.2\pm 0.1\%$ IoU score on the dataset of Japanese Society of Radiological Technology (JSRT). Besides that, we demonstrate the relative simplicity of the task by attaining surprisingly strong results with training sets of size 10 and 20: in terms of Dice score, $97.0\pm 0.8\%$ and $97.3\pm 0.5$, respectively, while in terms of IoU score, $92.2\pm 1.2\%$ and $93.3\pm 0.4\%$, respectively. To achieve these scores, we capitalize on the mixup augmentation technique, which yields a remarkable gain above $4\%$ IoU score in the size 10 setup.

연구 동기 및 목표

• 제한된 훈련 데이터로도 우수한 성능을 보이는 흉부 X선에서 폐와 심장을 자동으로 분할하는 딥러닝 모델을 개발하는 것.
• 주목적 기능 융합과 주의 메커니즘을 융합한 새로운 네트워크 아키텍처를 통해 기존 최신 기술 수준의 모델을 초월하는 분할 정확도를 향상시키는 것.
• 의료 영상 분할 작업에서 극도로 제한된 데이터 세트에서 mixup 데이터 증강 기법의 효과성을 조사하는 것.
• 특히 훈련 손실 지표는 안정되었음에도 불구하고 검증 지표가 지연되어 향상되는 특이한 훈련 동역학을 분석하는 것.

제안 방법

• 이중 방향 ConvLSTM 레이어를 통합하여 향상된 스킵 커넥션 기반 특징 융합을 구현한 수정된 U-Net 아키텍처인 Attention BCDU-Net을 제안한다.
• 업샘플링 단계에서 네트워크가 관련 특징에 집중할 수 있도록 주의 게이트를 통합하여 국소화 정확도를 향상시킨다.
• 입력 이미지와 마스크 간 선형 보간을 통해 합성 훈련 샘플을 생성하기 위해 하이퍼파라미터 δ를 사용한 mixup 데이터 증강을 적용한다.
• 네트워크 깊이와 손실 함수의 하이퍼파라미터 튜닝을 통해 모델 성능을 최적화하며, 평가 지표로 Dice와 IoU를 사용한다.
• 검증 손실 기반 조기 정지 전략을 적용하지만, 예상치 못한 후기 개선 현상이 관찰된다.
• 10개, 20개, 209개의 훈련 샘플로 구성된 훈련 세트를 사용하여 JSRT 데이터셋에서 모델을 훈련하고 평가한다.

실험 결과

연구 질문

• RQ1 10~20개의 훈련 샘플로만 이루어진 딥러닝 모델이 흉부 X선에서 최신 기술 수준의 분할 성능을 달성할 수 있는가?
• RQ2 훈련 데이터가 극도로 제한된 상황에서 mixup 증강 기법이 분할 성능 향상에 얼마나 효과적인가?
• RQ3 소규모 데이터 설정에서 훈련 지표는 안정되었음에도 불구하고 검증 Dice 점수가 지연되어 향상되는 이유는 무엇인가?
• RQ4 이중 방향 ConvLSTM과 주의 게이트의 통합이 의료 영상 분할에서 특징 융합을 어떻게 향상시키는가?
• RQ5 의료 영상 마스크의 본질적 유사성으로 인해 훈련 세트 크기가 증가함에 따라 증강이 더 이상 필요로 하지 않는가?

주요 결과

• Attention BCDU-Net은 전체 JSRT 데이터셋에서 98.1±0.1%의 Dice 점수와 95.2±0.1%의 IoU를 기록하여 이전 최신 기술 수준의 결과를 초월한다.
• 단지 10개의 훈련 샘플로도 97.0±0.8%의 Dice와 92.2±1.2%의 IoU를 달성하여 극도로 적은 데이터로도 뛰어난 성능을 보인다.
• 20개의 훈련 샘플로는 97.3±0.5%의 Dice와 93.3±0.4%의 IoU를 기록하여 더 작은 데이터 세트보다 일관되게 향상된 성능을 보인다.
• δ=0.2로 설정한 mixup 증강은 10개 샘플 설정에서 증강 없이 비교했을 때 IoU 점수를 4% 이상 향상시켜 소규모 데이터 일반화에서의 핵심적 역할을 입증한다.
• 10개 샘플 설정에서 400 에포크 이후 검증 Dice 점수가 갑작스럽게 상승하는 현상이 관찰되었으며, 이는 안정된 훈련 점수에도 불구하고 예상치 못한 일반화 역학을 시사한다.
• 훈련 세트 크기가 증가함에 따라 mixup의 효과가 감소함을 확인하여, 더 큰 데이터셋에서는 마스크 유사성으로 인해 증강의 필요성이 자연스럽게 줄어든다는 것을 시사한다.