[논문 리뷰] PSGR: Pixel-wise Sparse Graph Reasoning for COVID-19 Pneumonia Segmentation in CT Images

Haozhe Jia, Haoteng Tang외 5명|arXiv (Cornell University)|2021. 08. 09.

COVID-19 diagnosis using AI참고 문헌 43인용 수 4

한 줄 요약

이 논문은 CT 영상에서 코로나19 폐렴을 정확하게 분할하기 위해 깊이 학습 모델을 향상시키는 픽셀 단위 희소 그래프 추론 모듈인 PSGR을 제안한다. 특징 맵에서 희소로 연결된 그래프를 구성하고 장거리 의존성 추론을 수행함으로써, 기준 모델 대비 분할 성능을 향상시키며, COVID19-CT-100 데이터셋에서 DSC 84.83%를 달성한다.

ABSTRACT

Automated and accurate segmentation of the infected regions in computed tomography (CT) images is critical for the prediction of the pathological stage and treatment response of COVID-19. Several deep convolutional neural networks (DCNNs) have been designed for this task, whose performance, however, tends to be suppressed by their limited local receptive fields and insufficient global reasoning ability. In this paper, we propose a pixel-wise sparse graph reasoning (PSGR) module and insert it into a segmentation network to enhance the modeling of long-range dependencies for COVID-19 infected region segmentation in CT images. In the PSGR module, a graph is first constructed by projecting each pixel on a node based on the features produced by the segmentation backbone, and then converted into a sparsely-connected graph by keeping only K strongest connections to each uncertain pixel. The long-range information reasoning is performed on the sparsely-connected graph to generate enhanced features. The advantages of this module are two-fold: (1) the pixel-wise mapping strategy not only avoids imprecise pixel-to-node projections but also preserves the inherent information of each pixel for global reasoning; and (2) the sparsely-connected graph construction results in effective information retrieval and reduction of the noise propagation. The proposed solution has been evaluated against four widely-used segmentation models on three public datasets. The results show that the segmentation model equipped with our PSGR module can effectively segment COVID-19 infected regions in CT images, outperforming all other competing models.

연구 동기 및 목표

• 기존 깊이 학습 모델의 국소 수용장치가 제한되고 전반적인 추론 능력이 약한 점을 감안해, 비정상적이고 산만하며 작은 코로나19 감염 영역을 분할하는 데 도전한다.
• 자기 주의(Self-attention) 및 비국소(non-local) 기법의 한계를 극복한다. 이 기법들은 부과적인 상관관계와 노이즈를 유발하지만, 전반적인 맥락을 충분히 활용하지 못한다.
• 픽셀 수준의 정보를 유지하고 효과적인 장거리 특징 전파를 가능하게 하는 그래프 기반 추론 모듈을 개발한다.
• 다양한 병리적 패턴과 영상 프로토콜을 가진 여러 공개 데이터셋을 통해 제안된 방법의 유효성을 입증한다.

제안 방법

• 각 특징 맵의 픽셀을 직접 노드에 매핑하여 공간적 및 의미적 통일성을 유지하는 픽셀 단위 매핑을 통해 그래프를 구성한다.
• 분할 백본에서 생성된 특징 임bedding을 기반으로 모든 픽셀 특징 간의 쌍별 유사도 점수를 계산하여 초기 그래프 간선를 형성한다.
• 불확실한 픽셀 각각에 대해 K-가장 가까운 이웃 전략을 적용하여 유일한 K개의 강력한 연결만 유지함으로써 노이즈와 계산 부담을 줄인 희소로 연결된 그래프를 생성한다.
• 그래프 신경망을 사용하여 희소 그래프를 통해 메시지 전달 및 특징 집약을 수행함으로써 장거리 맥락 정보를 반영한 특징을 개선한다.
• 기존의 U-Net 및 U²-Net 아키텍처에 PSGR 모듈을 플러그인 방식으로 통합하여 아키텍처의 전반적인 개편 없이도 전역 추론 능력을 향상시킨다.
• 학습 가능한 불확실성 임계값 기반 메커니즘을 사용하여 불확실한 픽셀을 식별하고 우선순위를 정함으로써, 작은 또는 모호한 병변에 대한 민감도를 향상시킨다.

실험 결과

연구 질문

• RQ1 클러스터 기반 또는 영역 기반의 노드 매핑 대비 픽셀 단위 그래프 구축 전략이 의료 영상 분할에서 전역 맥락 모델링의 정밀도를 향상시키는가?
• RQ2 K-강력한 간선를 통해 그래프 노드를 희소로 연결함으로써 노이즈 전파를 줄이고 효과적인 장거리 의존성 학습을 유지할 수 있는가?
• RQ3 PSGR 모듈이 CT 영상에서 작고 산만하며 비정형적인 형태의 코로나19 폐렴 영역 분할 성능에 얼마나 기여하는가?
• RQ4 불확실한 픽셀의 선택 비율이 그래프 추론 과정의 강건성과 성능에 미치는 영향은 어떠한가?
• RQ5 PSGR 모듈은 다양한 분할 백본에 효과적으로 통합되고 다양한 공개 데이터셋에서 일반화 가능한가?

주요 결과

• U-2-Net+PSGR 모델은 COVID19-CT-100 데이터셋에서 최고의 DSC 84.83%를 기록하여 모든 기준 모델을 능가했다.
• U-2-Net+PSGR의 HD(Hausdorff Distance)는 42.85 mm로, 특히 작은 병변이나 분리된 병변에서 진짜값과의 경계 일치도가 뛰어나다는 것을 나타냈다.
• 제거 실험에서 최적의 불확실한 픽셀 선택 비율 $ R_u = 0.005 $ 가 분할 성능을 최대화했으며, 이 비율을 초과할 경우 노이즈 누적로 인해 성능 저하가 발생했다.
• 접근 행렬의 시각화 결과, PSGR이 먼 거리에 있지만 관련 있는 감염된 픽셀을 효과적으로 연결함으로써 장거리 의존성을 성공적으로 포착했음을 확인했다. 이는 미세하거나 고립된 영역에서도 마찬가지였다.
• PSGR 모듈은 모든 세 개의 공개 데이터셋에서 일관되게 분할 정확도를 향상시켰으며, 다양한 영상 프로토콜과 질환 표현 방식에 걸쳐 강건성과 일반화 능력을 입증했다.
• 대규모 외부 데이터를 사전 학습에 사용한 Inf-Net과 비교했을 때, U-2-Net+PSGR는 추가 데이터 없이도 유사한 성능을 달성했으며, 이는 제안된 모듈의 효율성과 효과성을 시사한다.