[논문 리뷰] Localise to segment: crop to improve organ at risk segmentation accuracy

Abraham George Smith, Denis Kutnár외 3명|arXiv (Cornell University)|2023. 04. 10.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 방사선 치료 계획에서 위험 기관 분할 정확도를 향상시키기 위해 국소화 이후 기관 분할을 수행하는 두 단계의 딥러닝 접근법을 제안한다. 분할 이전에 관련 해부학적 구조에 집중하기 위해 이미지를 자르는 방식으로, 클래스 불균형을 줄이고 학습 속도를 높이며, 특히 비장, 췐물선, 심장와 같은 작은 기관에서 DICE 스코어를 향상시킨다.

ABSTRACT

Increased organ at risk segmentation accuracy is required to reduce cost and complications for patients receiving radiotherapy treatment. Some deep learning methods for the segmentation of organs at risk use a two stage process where a localisation network first crops an image to the relevant region and then a locally specialised network segments the cropped organ of interest. We investigate the accuracy improvements brought about by such a localisation stage by comparing to a single-stage baseline network trained on full resolution images. We find that localisation approaches can improve both training time and stability and a two stage process involving both a localisation and organ segmentation network provides a significant increase in segmentation accuracy for the spleen, pancreas and heart from the Medical Segmentation Decathlon dataset. We also observe increased benefits of localisation for smaller organs. Source code that recreates the main results is available at \href{https://github.com/Abe404/localise_to_segment}{this https URL}.

연구 동기 및 목표

• 의료 영상 데이터셋의 클래스 불균형 영향을 줄임으로써 방사선 치료에서 위험 기관의 분할 정확도를 향상시키는 것.
• 단일 단계의 전체 영상 분할 대비 두 단계의 국소화-분할 파이프라인의 성능 향상 여부를 조사하는 것.
• 다양한 기관 크기에서 국소화가 학습 효율성, 수렴 속도, 모델 안정성에 미치는 영향을 평가하는 것.
• 작은 기관에서 클래스 불균형이 더 심하기 때문에 국소화가 더 큰 이점을 제공하는지 평가하는 것.
• 기본 진단 기반 국소화가 전체 해상도 영상에서의 엔드 투 엔드 학습보다 일관된 이점을 제공하는지 확인하는 것.

제안 방법

• 국소화 및 분할 네트워크 모두에 3D U-Net 아키텍처를 사용하였으며, 구현을 단순화하기 위해 공통된 모델 구조를 활용하였다.
• 학습 중 입력 및 출력 공간 차원을 일관되게 유지하기 위해 컨볼루션 레이어에 제로 패딩을 적용하였다.
• 작은 배치 크기(1 GPU당 2개)로도 학습을 안정화하기 위해 배치 정규화 대신 그룹 정규화를 사용하였다.
• 특히 작은 기관에서 중요한 클래스 불균형 문제를 해결하기 위해 DICE 손실과 교차 엔트로피 손실을 병합하였다.
• 64×256×256 크기의 자른 패치를 사용하여 최대한의 맥락을 유지하면서도 64GB GPU 메모리 제약에 맞추었다.
• 과적합을 방지하고 학습 효율성을 향상시키기 위해 조기 정지 및 모델 체크포인팅을 구현하였다.

Figure 1: Illustration showing the three different methods compared, including the baseline, two stage involving both a localisation network and organ segmentation network and the organ segmentation network that uses the ground truth to localise.

실험 결과

연구 질문

• RQ1 두 단계의 국소화-분할 파이프라인은 단일 단계의 전체 영상 학습 대비 위험 기관 분할 정확도를 향상시키는가?
• RQ2 국소화가 다양한 기관 크기에서 학습 시간, 수렴 속도, 모델 안정성에 미치는 영향은 어떠한가?
• RQ3 클래스 불균형이 더 심한 작은 기관에서 국소화의 이점이 더 두드러지는가?
• RQ4 기본 진단 기반 경계 상자로 국소화하는 방식이 전체 영상에서의 엔드 투 엔드 학습보다 일관되게 슈퍼어리어를 제공하는가?
• RQ5 국소화가 의료 영상 분할에서 클래스 불균형의 부정적 영향을 어느 정도 완화시키는가?

주요 결과

• 두 단계의 국소화 접근법은 비장, 췐물선, 심장의 DICE 스코어를 유의미하게 향상시켰으며, 특히 작은 기관에서 가장 큰 성과를 보였다.
• 국소화된 패치로 학습함으로써 전체 영상 패치 학습 대비 학습 시간을 단축시키고 수렴 속도를 향상시켰다. 이는 두 개의 네트워크를 학습하는 데에도 불구하고 성립한다.
• 기본 단일 단계 모델은 불안정한 학습 동역학을 보였으며, DICE 스코어의 변동성이 있었는데, 이는 심각한 클래스 불균형과 패치 내에서 전경 존재의 일관성 부족 때문일 가능성이 높다.
• 심장(좌심방)은 가장 느린 수렴 속도를 보였으며, 이는 작은 데이터셋과 높은 클래스 불균형 때문이지만, 국소화 덕분에 성능 향상이 이루어졌다.
• 가장 큰 학습 데이터셋을 가진 췐물선도 수렴률이 85%에 그쳐, 데이터셋 크기 자체만으로는 클래스 불균형 문제를 해결할 수 없음을 시사한다.
• 기본 진단 기반 국소화가 일관된 개선 효과를 보였으며, 이는 임상 워크플로우에서 수동 또는 상호작용식 국소화가 고정밀도 분할에 유용할 수 있음을 시사한다.