[논문 리뷰] Stack-U-Net: Refinement Network for Image Segmentation on the Example of Optic Disc and Cup

Artem Sevastopolsky, Stepan Drapak외 4명|arXiv (Cornell University)|2018. 04. 30.

Retinal Imaging and Analysis인용 수 25

한 줄 요약

이 논문은 망막 영상에서 시신경두 및 시신경구멍 분할을 위한 U-Net 기반의 개선 네트워크의 계층적 조합인 Stack-U-Net을 제안한다. 이는 더 큰 데이터셋이 필요 없이도 단일 U-Net 및 최신 기법들보다 분할 정확도를 크게 향상시킨다. 모델은 스킵 커넥션과 입력 이미지의 맥락을 활용해 예측을 반복적으로 개선하며, 공개 및 비공개 데이터셋에서 SOTA 성능을 달성한다. 시신경두의 경우 최대 DICE 스코어 0.97, 시신경구멍의 경우 0.85를 기록한다.

ABSTRACT

In this work, we propose a special cascade network for image segmentation, which is based on the U-Net networks as building blocks and the idea of the iterative refinement. The model was mainly applied to achieve higher recognition quality for the task of finding borders of the optic disc and cup, which are relevant to the presence of glaucoma. Compared to a single U-Net and the state-of-the-art methods for the investigated tasks, very high segmentation quality has been achieved without a need for increasing the volume of datasets. Our experiments include comparison with the best-known methods on publicly available databases DRIONS-DB, RIM-ONE v.3, DRISHTI-GS, and evaluation on a private data set collected in collaboration with University of California San Francisco Medical School. The analysis of the architecture details is presented, and it is argued that the model can be employed for a broad scope of image segmentation problems of similar nature.

연구 동기 및 목표

• 망막 망막도 영상에서 시신경두 및 시신경구멍 분할 정확도를 향상시켜 조기 녹내람 검출을 가능하게 하기 위해.
• 의료 영상 분할에서 훈련 데이터가 제한적이고 정밀도 요구 수준이 높은 문제를 해결하기 위해.
• 데이터셋 크기를 늘리지 않고도 분할 품질을 향상시킬 수 있는 확장성 있고 종단 간 개선 아키텍처를 개발하기 위해.
• 공개 벤치마크 및 UCSF 의과대학에서 확보한 대규모 비공개 데이터셋에서 모델 성능을 평가하기 위해.

제안 방법

• 모델은 U-Net 블록의 계층적 조합을 사용하며, 각 후속 네트워크는 이전 네트워크의 출력을 개선하면서도 원본 입력 이미지를 맥락으로 수신한다.
• 스택 내 각 기본 네트워크는 입력 이미지에서 첫 번째 레이어로의 스킵 커넥션을 포함하여 공간적 및 의미적 맥락을 유지한다.
• 일부 변형에서는 잔차 블록(ResU-Net)을 사용하여 기울기 흐름과 특징 학습을 향상시킨다.
• 개선 과정은 반복적이다: 예측은 계층별로 개선되며, 각 블록은 점차 더 세밀한 특징을 학습한다.
• 이중 교차 엔트로피와 DICE 손실을 사용해 엔드 투 엔드로 훈련되며, 옵timization은 Adam을 사용한다.
• 블록 수는 성능과 계산 비용의 균형을 맞추기 위해 조정되며, 15개 블록이 최적의 성능을 보였다.

실험 결과

연구 질문

• RQ1 U-Net 기반 네트워크의 계층적 조합은 시신경두 및 시신경구멍과 같이 작고 임상적으로 중요한 구조의 분할 정확도를 향상시킬 수 있는가?
• RQ2 입력 이미지 맥락을 활용한 반복적 개선은 작은 의료 영상 데이터셋에서 과적합을 줄이고 강건성을 향상시키는가?
• RQ3 스택된 U-Net 블록의 수가 시신경두 및 시신경구멍 분할의 성능 및 일반화 능력에 어떤 영향을 미치는가?
• RQ4 제안된 Stack-U-Net 아키텍처는 추가 훈련 데이터가 없이도 단일 U-Net 및 최신 기법들을 능가할 수 있는가?
• RQ5 각 블록의 첫 번째 레이어에 입력 이미지에서의 스킵 커넥션은 분할 정확도에 어떤 영향을 미치는가?

주요 결과

• 15개의 ResU-Net 블록을 사용한 Stack-U-Net은 DRISHTI-GS 데이터셋에서 시신경두 분할에 DICE 스코어 0.97을 기록하며, 모든 기준 모델을 능가했다.
• RIM-ONE v.3 데이터셋에서 모델은 시신경구멍 분할에 DICE 스코어 0.95를 달성하여 이전 최신 기법을 초월했다.
• 비공개 UCSF-DB 데이터셋에서 Stack-U-Net은 시신경구멍 분할에 DICE 스코어 0.85를 기록했으며, 인간 전문가의 평균 인간 대 인간 DICE 스코어 0.66를 초월했다.
• 15개 블록을 가진 모델가 최고의 성능을 보였으며, 15개를 초과해 블록 수를 늘여도 성능 향상이 없었고, 오히려 성능 저하가 발생했다.
• 각 블록의 첫 번째 레이어에 입력 이미지에서의 스킵 커넥션을 도입함으로써 성능이 소폭이지만 일관되게 향상되었으며, 유일한 구성에서는 스킵 커넥션의 부재가 약간 더 좋은 성능을 보였다.
• 시각적 분석 결과, 최고 성능을 낸 모델이 단일 U-Net이나 이전 기법보다 도전적인 케이스(예: 저대비 또는 흐린 영역)를 더 강건하게 처리하는 것으로 나타났다.