[논문 리뷰] CT organ segmentation using GPU data augmentation, unsupervised labels and IOU loss

Blaine Rister, Darvin Yi외 3명|arXiv (Cornell University)|2018. 11. 27.

Radiomics and Machine Learning in Medical Imaging참고 문헌 18인용 수 9

한 줄 요약

이 논문은 GPU 가속 데이터 증강, 폐 및 뼈에 대한 비지도 형태학적 레이블링, 그리고 교차율(Intersection-over-Union, IOU) 손실 함수를 사용하여 높은 정확도를 달성하는 경량 3D 완전 컨volution 네트워크를 제안한다. 이 방법은 특히 불균형 세그멘테이션 작업에서 가중치가 부여된 교차 엔트로피 손실보다 뛰어나며, 5개의 기관에 대해 78–92%의 Dice 점수를 기록하고, 일반 소비자용 GPU에서 평균 5.98초의 추론 시간을 기록한다.

ABSTRACT

Fully-convolutional neural networks have achieved superior performance in a variety of image segmentation tasks. However, their training requires laborious manual annotation of large datasets, as well as acceleration by parallel processors with high-bandwidth memory, such as GPUs. We show that simple models can achieve competitive accuracy for organ segmentation on CT images when trained with extensive data augmentation, which leverages existing graphics hardware to quickly apply geometric and photometric transformations to 3D image data. On 3 mm^3 CT volumes, our GPU implementation is 2.6-8X faster than a widely-used CPU version, including communication overhead. We also show how to automatically generate training labels using rudimentary morphological operations, which are efficiently computed by 3D Fourier transforms. We combined fully-automatic labels for the lungs and bone with semi-automatic ones for the liver, kidneys and bladder, to create a dataset of 130 labeled CT scans. To achieve the best results from data augmentation, our model uses the intersection-over-union (IOU) loss function, a close relative of the Dice loss. We discuss its mathematical properties and explain why it outperforms the usual weighted cross-entropy loss for unbalanced segmentation tasks. We conclude that there is no unique IOU loss function, as the naive one belongs to a broad family of functions with the same essential properties. When combining data augmentation with the IOU loss, our model achieves a Dice score of 78-92% for each organ. The trained model, code and dataset will be made publicly available, to further medical imaging research.

연구 동기 및 목표

• 폐 및 뼈와 같은 기관에 대해 수동 레이블링과 비지도 형태학적 레이블링을 조합하여 의료 영상 데이터의 부족 문제를 해결한다.
• 복잡한 아키텍처 대신 단순하고 메모리 효율적인 모델을 사용하여 3D 의료 영상 세그멘테이션의 계산 요구량을 줄인다.
• 표준 교차 엔트로피 손실 대신 IOU 기반 손실 함수를 도입하여 불균형 데이터 세트에서의 세그멘테이션 정확도를 향상시킨다. 이는 의료 영상에서 흔한 불균형 문제를 해결하기 위함이다.
• GPU 텍스처 샘플링 하드웨어를 활용하여 데이터 증강을 가속화하여 학습 중 실시간 기하학적 및 광학적 변환을 가능하게 한다.
• 미래의 다기관 CT 세그멘테이션 연구를 지원하기 위해 공개 가능한 데이터셋과 훈련된 모델을 구축한다.

제안 방법

• CUDA와 텍스처 샘플링을 활용해 저지연 처리가 가능한 기하학적 및 광학적 변환을 사용한 GPU 가속 3D 데이터 증강을 구현한다.
• 3D 푸리에 변환으로 가속화된 형태학적 연산을 통해 폐와 뼈에 대한 비지도 훈련 레이블을 생성하여 수동 레이블링의 노력을 감소시킨다.
• 간, 신장, 방광에 대한 반자동 수동 레이블링과 비지도 레이블을 융합하여 총 130개의 스캔으로 구성된 데이터셋을 구성한다.
• 메모리 사용량을 줄이고 학습 효율성을 향상시키기 위해 최소한의 파라미터를 가진 경량 3D 완전 컨volution 네트워크(FCN)를 사용한다.
• 교차율(Intersection-over-Union, IOU) 손실 함수를 도입하였으며, 수학적으로 분석하여 불균형 세그멘테이션 작업에서 균형 잡힌 오차 페널티를 제공함을 입증하였다.
• IOU 손실 함수가 동일한 최적화 성질을 가진 함수의 일반적인 가족에 属하며, 실질적으로 가중치가 부여된 교차 엔트로피보다 뛰어나다는 것을 입증하였다.

실험 결과

연구 질문

• RQ1 복잡한 아키텍처 없이도 광범위한 데이터 증강과 IOU 손실 함수를 사용하여 훈련된 단순한 3D FCN이 경쟁력 있는 기관 세그멘테이션 성능을 달성할 수 있는가?
• RQ2 GPU 기반 데이터 증강은 3D CT 세그멘테이션에서 CPU 기반 방법에 비해 속도와 학습 효율성 측면에서 어떻게 비교되는가?
• RQ3 비지도 형태학적 연산이 폐나 뼈와 같은 기관에 대해 신뢰할 수 있는 훈련 레이블을 얼마나 잘 생성할 수 있으며, 수동 레이블링과 비교해 볼 때 어떤가?
• RQ4 IOU 손실 함수는 불균형 의료 영상 세그멘테이션 작업에서 특히 작은 또는 비정형 모양을 가진 기관에 대해 가중치가 부여된 교차 엔트로피보다 항상 뛰어나게 작용하는가?
• RQ5 단일의 일반 목적 모델이 높은 정확도와 낮은 계산 비용으로 동시에 여러 기관을 효과적으로 세그멘테이션할 수 있는가?

주요 결과

• IOU 손실 함수는 모든 기관에서 가중치가 부여된 교차 엔트로피 손실보다 뚜렷이 뛰어나며, 특히 소수 기관인 방광과 같은 구조에서의 가짜 양성 결과를 크게 감소시켰다.
• 데이터 증강은 모델의 일반화 능력을 향상시켰으며, IOU 손실과 조합했을 경우 Dice 점수를 최대 10%까지 향상시켰다. 반면 교차 엔트로피와 조합했을 경우 성능이 저하되었다.
• 130개 스캔 데이터셋을 사용하여 3mm³ 등방성 해상도로 간, 신장, 방광, 폐, 뼈 등 5개 기관에서 78–92%의 Dice 점수를 기록하였다.
• LiTS 간 세그멘테이션 챌린지 테스트 세트에서 최고 성능을 낸 모델은 평균 Dice 점수 90.5%를 기록하여 100개 이상의 제출 중 49위를 기록하였다.
• 모델은 일반 소비자용 하드웨어에서 3D CT 볼륨을 평균 5.98초 만에 처리하였으며, GPU 메모리는 2429 MB로 제한되어 있어 효율성을 입증하였다.
• 저자들은 IOU 손실 함수가 유일한 함수가 아니며, 동일한 최적화 성질을 가진 함수의 가족에 属하며, 불균형 세그멘테이션에서 균형 잡힌 오차 페널티를 제공함을 입증하였다.