[논문 리뷰] Deep Q-Network-Driven Catheter Segmentation in 3D US by Hybrid Constrained Semi-Supervised Learning and Dual-UNet

Hongxu Yang, Caifeng Shan외 2명|arXiv (Cornell University)|2020. 06. 25.

Advanced Neural Network Applications참고 문헌 16인용 수 6

한 줄 요약

이 논문은 3D 초음파 카테터 분할을 위한 이단계 딥러닝 프레임워크를 제안하며, 깊이 Q-네트워크(DQN) 기반의 군집화 국소화와 준감독형 듀얼-UNet를 조합한다. 비라벨 데이터에 대해 불확실성 및 문맥 손실과 같은 하이브리드 제약 조건을 활용함으로써, 주어진 라벨 수를 크게 줄이며 최신 기술 수준의 성능을 달성하며, 라벨링 부담을 줄이고 정확도 및 효율성을 기존의 감독 및 준감독 접근 방식보다 향상시킨다.

ABSTRACT

Catheter segmentation in 3D ultrasound is important for computer-assisted cardiac intervention. However, a large amount of labeled images are required to train a successful deep convolutional neural network (CNN) to segment the catheter, which is expensive and time-consuming. In this paper, we propose a novel catheter segmentation approach, which requests fewer annotations than the supervised learning method, but nevertheless achieves better performance. Our scheme considers a deep Q learning as the pre-localization step, which avoids voxel-level annotation and which can efficiently localize the target catheter. With the detected catheter, patch-based Dual-UNet is applied to segment the catheter in 3D volumetric data. To train the Dual-UNet with limited labeled images and leverage information of unlabeled images, we propose a novel semi-supervised scheme, which exploits unlabeled images based on hybrid constraints from predictions. Experiments show the proposed scheme achieves a higher performance than state-of-the-art semi-supervised methods, while it demonstrates that our method is able to learn from large-scale unlabeled images.

연구 동기 및 목표

• 3D 초음파 카테터 분할에서 감독형 딥러닝의 높은 라벨링 비용을 줄이기 위해.
• 제한된 라벨 데이터로도 풍부한 비라벨 3D US 이미지를 효과적으로 활용하여 분할 성능을 향상시키기 위해.
• 강화학습을 통한 사전 국소화를 통합하여 더 효율적이고 정확한 분할 파이프라인을 개발하기 위해.
• 기존 준감독 학습 방법의 한계, 예를 들어 파라미터 상관관계 및 불안정한 불확실성 추정을 극복하기 위해.
• 전체 볼륨 또는 철저한 패치 기반 분할 대비 계산 비용을 줄여 실시간 임상 적용을 가능하게 하기 위해.

제안 방법

• 딥러닝 Q-네트워크(DQN)는 볼륨 내에서 정책을 학습하여 볼륨을 탐색함으로써 3D US에서 카테터의 군집화 국소화를 수행한다. 이 과정에서 병렬 볼륨 수준의 라벨링이 필요하지 않다.
• DQN 에이전트는 $128^3$ 볼륨에서 $55^3$ 관측 공간을 사용하며, 카테터 중심을 찾기 위한 이산적인 3D 이동으로 행동을 정의한다.
• 패치 기반 듀얼-UNet는 DQN에 의해 감지된 영역 주변에서 카테터를 분할한다. 이때 Compact-UNet을 백본으로 사용하고 드롭아웃을 통해 불확실성 추정을 수행한다.
• 듀얼-UNet는 감독 손실, 내부 및 상호 네트워크 불확실성 제약 조건, 그리고 문맥 제약 손실을 조합한 새로운 준감독 학습 기반으로 훈련된다.
• 불확실성은 T=8번의 전방향 전파를 통해 몬테카를로 드롭아웃을 사용하여 추정되며, 비라벨 이미지에서 신뢰할 수 있는 예측 결과를 토대로 가짜 라벨이 생성된다.
• 비라벨 패치 간의 의미적 문맥 일관성을 강제함으로써 특징 학습을 향상시키기 위해 문맥 제약 손실이 도입된다.

실험 결과

연구 질문

• RQ1 딥러닝 기반 강화학습은 국소화 정확도를 희생시키지 않고 3D 초음파 카테터 분할에서 라벨링 작업을 줄일 수 있는가?
• RQ2 제한된 라벨 데이터로도 비라벨 3D US 데이터를 효과적으로 활용하여 분할 성능을 향상시킬 수 있는가?
• RQ3 하이브리드 제약 조건 기반 준감독 학습 기반은 기존의 불확실성 인식 및 적대적 SSL 방법보다 3D 카테터 분할에서 더 우수한 성능을 낼 수 있는가?
• RQ4 군집화 국소화와 패치 기반 분할의 통합은 전체 볼륨 또는 철저한 패치 기반 방법 대비 추론 효율성을 향상시킬 수 있는가?
• RQ5 불확실성 추정 및 문맥 모델링은 준감독 3D 분할의 강인성과 정확도에 어떤 영향을 미치는가?

주요 결과

• 제안된 방법은 외부에서 채취한 3D 초음파 데이터셋에서 최신 기술 수준의 준감독 방법보다 평균 딜레이 스코어(Dice Score)가 4% 높게 달성되었다.
• 라벨이 18장 뿐이고 비라벨이 42장일 때도, 동일한 백본 네트워크를 사용한 감독 기반 베이스라인을 초월하는 성능을 보였다.
• DQN 기반 사전 국소화는 $128^3$ 볼륨에서 평균 ± 표준편차로 4.3 ± 2.6 볼륨에서 국소화 정확도를 달성하였으며, $160^3$ 경우(4.7 ± 4.8 볼륨)보다 뛰어났다.
• 이중 단계 프레임워크는 3D 볼륨을 약 1.2초 내에 처리하며, 철저한 패치 기반 방법이 요구하는 10초 이상보다 크게 빠르게 작동한다.
• 절단 실험을 통해 하이브리드 손실 제약 조건의 각 구성 요소—불확실성 및 문맥—이 성능 향상에 기여하며, 전체 모델이 가장 우수한 성능을 기록하였다.
• $π$-모델 기반 베이스라인은 복잡한 3D 초음파 이미지에서 불안정한 불확실성 추정으로 인해 신뢰할 수 없는 가짜 라벨을 생성하여 제안된 방법에 비해 성능이 열 劣했다.