[논문 리뷰] UT-Net: Combining U-Net and Transformer for Joint Optic Disc and Cup Segmentation and Glaucoma Detection

Rukhshanda Hussain, Hritam Basak|arXiv (Cornell University)|2023. 03. 08.

Retinal Imaging and Analysis인용 수 6

한 줄 요약

UT-Net는 망막 후면도 영상에서 시신경판과 시신경구를 동시에 분할하고 정확한 녹내장 진단을 가능하게 하기 위해 U-Net과 비전 트랜스포머 아키텍처를 융합한 혁신적인 하이브리드 딥러닝 프레임워크를 제안한다. 다중 헤드 컨텍스트 추론 주의, 주의 가중 양방향 융합, 향상된 혼합 손실을 통합함으로써, UT-Net은 세 가지 공개 데이터셋에서 최신 기술 수준(SOTA)의 성능을 달성하며, 특히 시신경판 대비 비율(CDR) 추정 정확도를 크게 향상시킨다.

ABSTRACT

Glaucoma is a chronic visual disease that may cause permanent irreversible blindness. Measurement of the cup-to-disc ratio (CDR) plays a pivotal role in the detection of glaucoma in its early stage, preventing visual disparities. Therefore, accurate and automatic segmentation of optic disc (OD) and optic cup (OC) from retinal fundus images is a fundamental requirement. Existing CNN-based segmentation frameworks resort to building deep encoders with aggressive downsampling layers, which suffer from a general limitation on modeling explicit long-range dependency. To this end, in this paper, we propose a new segmentation pipeline, called UT-Net, availing the advantages of U-Net and transformer both in its encoding layer, followed by an attention-gated bilinear fusion scheme. In addition to this, we incorporate Multi-Head Contextual attention to enhance the regular self-attention used in traditional vision transformers. Thus low-level features along with global dependencies are captured in a shallow manner. Besides, we extract context information at multiple encoding layers for better exploration of receptive fields, and to aid the model to learn deep hierarchical representations. Finally, an enhanced mixing loss is proposed to tightly supervise the overall learning process. The proposed model has been implemented for joint OD and OC segmentation on three publicly available datasets: DRISHTI-GS, RIM-ONE R3, and REFUGE. Additionally, to validate our proposal, we have performed exhaustive experimentation on Glaucoma detection from all three datasets by measuring the Cup to Disc Ratio (CDR) value. Experimental results demonstrate the superiority of UT-Net as compared to the state-of-the-art methods.

연구 동기 및 목표

• 깊은 인코더에서의 급격한 다운샘플링으로 인해 CNN 기반 모델이 장거리 상관관계를 포착하는 데 한계를 가짐을 해결하기 위해.
• 초기 녹내장 진단을 위해 망막 후면도 영상에서 시신경판과 시신경구 분할 정확도를 향상시키기 위해.
• 시신경판 대비 비율(CDR)을 핵심 생물학적 지표로 사용하여 분할과 녹내장 분류를 동시에 수행하는 통합 파ip라인 개발을 위해.
• 자기주의 주의 메커니즘을 통한 전역적 맥락 통합과 U-Net 스위프트 연결을 통한 국소 세부 정보 유지로 특징 표현을 향상시키기 위해.
• 분할 및 분류 작업의 공동 학습을 강력하게 감독하기 위한 새로운 손실 함수 도입을 위해.

제안 방법

• CNN의 국소 인덕티브 바이어스와 자기주의 주의를 통한 전역 맥락 모델링을 조합하기 위해 U-Net 유사 인코더-디코더 아키텍처와 비전 트랜스포머 기반 인코더를 통합한다.
• 특징 맵에서 장거리 상관관계를 더 효과적으로 모델링하기 위해 표준 자기주의 주의를 개선하기 위해 다중 헤드 컨텍스트 주의를 적용한다.
• 다양한 인코더 스테이지의 특징을 동적으로 가중하고 융합하기 위해 주의 가중 양방향 융합을 적용하여 표현 품질을 향상시킨다.
• 다양한 인코딩 계층에서 맥락 정보를 추출하여 깊은 계층적 표현을 구축하고 수신 영역 커버리지 향상에 기여한다.
• 이진 교차 엔트로피 손실과 딱지 손실을 클래스 불균형 및 공간 가중치와 결합한 향상된 혼합 손실을 제안하여 분할 감독을 향상시킨다.
• 다중 스케일 특징 융합과 스위프트 연결을 활용하여 디코딩 과정에서 세밀한 세부 정보를 유지함으로써 정확한 시신경판/구 경계 검출에 기여한다.

실험 결과

연구 질문

• RQ1 U-Net과 비전 트랜스포머를 융합한 하이브리드 아키텍처가 순수한 CNN 또는 ViT 모델 대비 시신경판 및 시신경구 분할 정확도를 향상시킬 수 있는가?
• RQ2 다중 헤드 컨텍스트 주의 통합이 망막 영상 분할에서 장거리 특징 모델링을 향상시키는가?
• RQ3 주의 가중 양방향 융합 메커니즘이 다중 수준 특징을 융합하여 분할 및 녹내장 진단 성능을 향상시키는 데 얼마나 효과적인가?
• RQ4 향상된 혼합 손실이 불균형한 망막 데이터셋에서 학습 안정성과 분할 성능을 얼마나 향상시키는가?
• RQ5 공동 분할 및 분류 파이프라인은 CDR 추정을 통한 녹내장 진단 벤치마크 데이터셋에서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• UT-Net는 DRISHTI-GS 데이터셋에서 기존 최신 기술 수준의 방법들보다 시신경판 및 시신경구 딱지 점수 모두에서 뛰어난 분할 성능를 달성한다.
• RIM-ONE R3 데이터셋에서 UT-Net는 시신경판의 평균 딱지 점수 0.948과 시신경구의 평균 딱지 점수 0.872를 기록하여 다양한 영상 품질에서도 뛰어난 강건성을 입증한다.
• REFUGE 데이터셋에서 UT-Net는 CDR 예측 정확도 92.3%를 달성하여 이전 방법들보다 뚜렷이 뛰어난 녹내장 진단 성능를 보였다.
• 제거 실험 결과, 다중 헤드 컨텍스트 주의, 주의 가중 융합, 향상된 혼합 손실 각각이 최종 성능에 기여하는 데 중요한 역할을 한다는 것이 확인되었다.
• 다양한 해상도, 대비, 병변 유병률을 가진 데이터셋 간에도 모델이 잘 일반화되어 있어 강력한 도메인 적응 능력을 보였다.
• 공동 학습 프레임워크는 분할 및 녹내장 분류의 엔드 투 엔드 최적화를 가능하게 하여 추론 시간을 단축하고 CDR 추정 일관성을 향상시켰다.