[논문 리뷰] UNet#: A UNet-like Redesigning Skip Connections for Medical Image Segmentation

Ledan Qian, Xiaofang Zhou외 2명|arXiv (Cornell University)|2022. 05. 24.

Radiomics and Machine Learning in Medical Imaging인용 수 9

한 줄 요약

이 논문은 다중 척도 특징 간의 특징 집합을 향상시키기 위해 밀집 및 전체 스케일 스케일 연결을 조합한 새로운 UNet-유사 아키텍처인 UNet#을 제안한다. 이는 장기 및 병변 경계 분할 성능을 향상시킨다. 모델 정제를 위한 딥 서프비전과 오진 억제를 위한 분류 지도 모듈을 통합함으로써, UNet#은 CT, MRI 및 EM 데이터셋에서 다양한 2D 및 3D 의료 영상 분할 벤치마크에서 최신 기술(SOTA) 성능을 달성한다.

ABSTRACT

As an essential prerequisite for developing a medical intelligent assistant system, medical image segmentation has received extensive research and concentration from the neural network community. A series of UNet-like networks with encoder-decoder architecture has achieved extraordinary success, in which UNet2+ and UNet3+ redesign skip connections, respectively proposing dense skip connection and full-scale skip connection and dramatically improving compared with UNet in medical image segmentation. However, UNet2+ lacks sufficient information explored from the full scale, which will affect the learning of organs' location and boundary. Although UNet3+ can obtain the full-scale aggregation feature map, owing to the small number of neurons in the structure, it does not satisfy the segmentation of tiny objects when the number of samples is small. This paper proposes a novel network structure combining dense skip connections and full-scale skip connections, named UNet-sharp (UNet\#) for its shape similar to symbol \#. The proposed UNet\# can aggregate feature maps of different scales in the decoder sub-network and capture fine-grained details and coarse-grained semantics from the full scale, which benefits learning the exact location and accurately segmenting the boundary of organs or lesions. We perform deep supervision for model pruning to speed up testing and make it possible for the model to run on mobile devices; furthermore, designing two classification-guided modules to reduce false positives achieves more accurate segmentation results. Various experiments of semantic segmentation and instance segmentation on different modalities (EM, CT, MRI) and dimensions (2D, 3D) datasets, including the nuclei, brain tumor, liver, and lung, demonstrate that the proposed method outperforms state-of-the-art models.

연구 동기 및 목표

• UNet2+가 전체 스케일의 맥락 정보를 포착하는 데에 한계가 있고, UNet3+가 제한된 훈련 데이터에서 소형 물체를 처리하는 데에 어려움을 겪는 문제를 해결한다.
• 밀집 및 전체 스케일 스케일 연결을 모두 활용하여 의미 수준 간의 특징 융합을 향상시키는 통합 아키텍처를 설계한다.
• 딥 서프비전 기반의 모델 정제를 통해 효율적인 추론을 가능하게 하여 모바일 장치에의 배포를 가능하게 한다.
• 결정 경계를 정밀하게 조정하기 위해 분류 지도 모듈을 사용하여 분할에서의 오진 억제를 개선한다.

제안 방법

• 다중 척도 특징 집합을 위해 밀집 스케일 연결과 전체 스케일 스케일 연결을 조합한 #-형 구조를 가진 새로운 UNet# 아키텍처를 도입한다.
• 최종 출력 및 네 개의 중간 레이어를 포함한 여덟 개의 디코더 브랜치에 걸쳐 딥 서프비전을 구현하여 계층적 표현 학습과 모델 정제를 가능하게 한다.
• 딥 서프비전 브랜치에 대해 크기 기반 정제를 적용하여 모델 파라미터를 감소시키면서도 성능을 유지한다.
• 과다 분할 및 오진 억제를 위해 두 개의 분류 지도 모듈(CGM)을 서로 다른 단계에 통합하여 결정 경계를 정밀하게 조정한다.
• 다양한 해상도에서의 정답을 활용한 다중 척도 서프비전을 통해 디코더의 모든 수준에서 특징 학습을 안내한다.
• 분할 헤드와 분류 헤드 양쪽 모두에서 교차 엔트로피 손실과 DICE 손실을 사용하여 국소화 및 클래스 식별 최적화를 동시에 수행한다.

실험 결과

연구 질문

• RQ1 밀집 연결과 전체 스케일 연결을 조합한 하이브리드 스케일 연결 설계가 의료 영상 분할에서 특징 표현을 향상시킬 수 있는가?
• RQ2 여러 디코더 브랜치에 걸친 딥 서프비전이 성능 저하 없이 효과적인 모델 정제를 가능하게 하는가?
• RQ3 분류 지도 모듈이 장기 및 병변 분할에서 오진을 감소시키고 경계 정확도를 향상시킬 수 있는가?
• RQ4 UNet#은 다양한 의료 영상 모odalities와 차원에서 최신 기술 모델과 비교해 어떻게 성능을 내는가?

주요 결과

• BraTs19 데이터셋에서 UNet#은 DICE 스코어 95.60%와 IoU 91.68%를 기록하여 비교된 모든 SOTA 모델을 압도했다.
• 분류 지도 모듈(CGM)을 적용한 UNet#은 BraTs19에서 DICE 스코어 96.14%와 IoU 92.08%를 기록하여 기본 모델 대비 각각 0.54%와 0.4% 향상되었다.
• Lits17 데이터셋에서 UNet#은 DICE 97.29%와 IoU 94.79%를 기록하여 TransUNet 및 AttnUNet을 포함한 모든 베이스라인을 뛰어넘었다.
• CGM 모듈은 Lits17에서 DICE 0.95% 향상 및 IoU 0.52% 향상으로 성능을 추가로 향상시켜 오진 억제에 효과적임을 확인했다.
• 딥 서프비전 기반 정제로 최대 97배(9.71M에서 0.1M로) 파라미터를 감소시켰지만 강력한 추론 정확도를 유지했다.
• 정제된 L¹ 모델(0.1M 파라미터)은 DSB2018에서 IoU 88.5%를 기록하여 모바일 배포 가능성의 타당성을 입증했다.