[논문 리뷰] Generative Image Translation for Data Augmentation of Bone Lesion Pathology

Anant Gupta, Srivas Venkatesh외 2명|arXiv (Cornell University)|2019. 02. 06.

AI in cancer detection참고 문헌 18인용 수 30

한 줄 요약

이 논문은 희귀 뼈병변 분류에서의 데이터 부족과 클래스 불균형 문제를 해결하기 위해 정상 X선 영상에 현실적인 뼈 병변을 합성하는 순환일致적 생성 적대적 네트워크(CycleGAN)를 제안한다. 병변이 없는 패치를 병변 유사 패치로 변환하고 전체 영상에 다시 융합함으로써 분류기 성능을 크게 향상시켰으며, 햄머스에서 훈련한 모델을 전이 학습하여 대퇴부에서 최대 15%의 AUC 향상을 달성했다.

ABSTRACT

Insufficient training data and severe class imbalance are often limiting factors when developing machine learning models for the classification of rare diseases. In this work, we address the problem of classifying bone lesions from X-ray images by increasing the small number of positive samples in the training set. We propose a generative data augmentation approach based on a cycle-consistent generative adversarial network that synthesizes bone lesions on images without pathology. We pose the generative task as an image-patch translation problem that we optimize specifically for distinct bones (humerus, tibia, femur). In experimental results, we confirm that the described method mitigates the class imbalance problem in the binary classification task of bone lesion detection. We show that the augmented training sets enable the training of superior classifiers achieving better performance on a held-out test set. Additionally, we demonstrate the feasibility of transfer learning and apply a generative model that was trained on one body part to another.

연구 동기 및 목표

• 희귀 뼈병변을 X선 영상에서 분류할 때 발생하는 심각한 클래스 불균형과 제한된 훈련 데이터 문제를 해결한다.
• 정상 해부학적 구조에 현실적이고 병리학적으로 특화된 병변을 합성하는 데이터 증강 전략을 개발한다.
• 희귀 또는 샘플 수가 부족한 다른 해부학적 부위(예: 경골, 대퇴부)에 대해, 한 부위(예: 어깨 뼈)에서 훈련한 생성 모델을 활용해 전이 학습을 가능하게 한다.
• 분류기 신뢰도 기반의 가짜 레이블링을 통해 고품질의 합성 병변을 자동으로 필터링하여 후속 모델 성능을 향상시킨다.

제안 방법

• 정상 패치와 병변이 있는 패치 간에 공유된 잠재 공간을 가진 CycleGAN 프레임워크를 사용하여 비지도 이미지-이미지 번역 과제로 병변 생성을 설정한다.
• 비병변 패치를 공유된 잠재 표현으로 매핑하고, 이를 통해 병변이 있는 패치를 재구성할 수 있도록 인코더와 생성기 모델을 훈련한다.
• 합성된 병변 패치를 원본 전체 X선 영상에 자연스럽게 융합하기 위해 알파 블렌딩을 적용하여 해부학적 현실감을 유지한다.
• 기본 분류기를 사용하여 생성된 이미지에 대해 가짜 레이블링을 수행하고, 신뢰도 점수가 임계치(t) 이상인 경우에만 증강된 훈련 세트에 포함시킨다.
• 검증 세트에서 임계치(t)를 최적화하여 합성 샘플의 품질과 양 사이의 균형을 맞춘다.
• 어깨 뼈에서 훈련한 어깨 뼈 전용 생성 모델을 적용하여 경골과 대퇴부에 병변을 생성함으로써, 데이터 부족으로 전용 모델 훈련이 불가능한 경우에도 전이 학습을 가능하게 한다.

실험 결과

연구 질문

• RQ1 생성적 이미지 번역 기법이 희귀 병변 분류에서 데이터 부족 문제를 완화하기 위해 정상 X선 영상에 현실적인 뼈 병변을 효과적으로 합성할 수 있는가?
• RQ2 합성된 병변으로 증강된 데이터를 사용할 경우, 후속 분류기 성능은 어떻게 향상되는가?
• RQ3 어떤 정도로 한 해부학적 부위(예: 어깨 뼈)에서 훈련한 생성 모델이 다른 부위(예: 경골, 대퇴부)에 대해 제한적 또는 노이즈가 많은 데이터로 병변을 생성하는 데 전이 가능할 수 있는가?
• RQ4 분류기 성능을 최대화하면서 저품질 샘플을 도입하지 않기 위해, 합성 이미지에 대한 가짜 레이블링의 최적 임계치는 무엇인가?

주요 결과

• 가짜 레이블링에 대해 임계치 t=0.9를 사용했을 때, 햄머스 테스트 세트에서 AUC가 통계적으로 유의미하게 5% 향상되어 0.876에서 0.924로 증가했다.
• 경골의 경우, 경골 전용 생성 모델을 사용했을 때 약 2%의 미미하지만 일관된 AUC 향상이 있었고, 어깨 뼈 기반 기본 분류기와 전이 학습을 결합했을 때는 8%의 향상이 이루어졌다.
• 대퇴부에서는 전용 생성 모델 훈련이 데이터 다양성과 부족으로 인해 불가능했지만, 어깨 뼈 모델에서 전이 학습을 통해 기존 기본 분류기 대비 15%의 AUC 향상을 달성했다.
• 성능은 가짜 레이블링 임계치(t)에 민감했으며, 햄머스의 최적 성능은 t=0.9, 대퇴부는 t=0.95에서 달성되어, 적절한 초모수 튜닝의 중요성을 시사했다.
• 제거 분석 결과, 과도하거나 부족한 합성 데이터 모두 성능 저하를 초래함을 확인하여, 신뢰도 기반 필터링을 통한 품질 제어의 필요성을 강조했다.
• 결과적으로 전이 학습이 제한된 훈련 데이터를 가진 해부학적 부위에 대해서도 효과적인 데이터 증강을 가능하게 하여, 자원이 부족한 의료 영상 환경에서 생성 모델의 활용도를 넓혔다.