[논문 리뷰] ConDSeg: A General Medical Image Segmentation Framework via Contrast-Driven Feature Enhancement

Mengqi Lei, Haochen Wu외 2명|arXiv (Cornell University)|2024. 12. 11.

Medical Imaging and Analysis인용 수 4

한 줄 요약

ConDSeg는 일致성 강화 사전 훈련 전략과 불확실성 감소를 위한 의미 정보 분리 모듈을 통해 저대비도 및 연약한 경계 조건에 대한 강건성을 향상시키는 일반적인 의료 영상 분할 프레임워크이다. 대비 기반 특징 집합과 크기 인지 디코더를 사용하여 전경-배경 구분과 다중 척도 국소화를 향상시키며, 다섯 가지 다양한 의료 영상 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Medical image segmentation plays an important role in clinical decision making, treatment planning, and disease tracking. However, it still faces two major challenges. On the one hand, there is often a ``soft boundary'' between foreground and background in medical images, with poor illumination and low contrast further reducing the distinguishability of foreground and background within the image. On the other hand, co-occurrence phenomena are widespread in medical images, and learning these features is misleading to the model's judgment. To address these challenges, we propose a general framework called Contrast-Driven Medical Image Segmentation (ConDSeg). First, we develop a contrastive training strategy called Consistency Reinforcement. It is designed to improve the encoder's robustness in various illumination and contrast scenarios, enabling the model to extract high-quality features even in adverse environments. Second, we introduce a Semantic Information Decoupling module, which is able to decouple features from the encoder into foreground, background, and uncertainty regions, gradually acquiring the ability to reduce uncertainty during training. The Contrast-Driven Feature Aggregation module then contrasts the foreground and background features to guide multi-level feature fusion and key feature enhancement, further distinguishing the entities to be segmented. We also propose a Size-Aware Decoder to solve the scale singularity of the decoder. It accurately locate entities of different sizes in the image, thus avoiding erroneous learning of co-occurrence features. Extensive experiments on five medical image datasets across three scenarios demonstrate the state-of-the-art performance of our method, proving its advanced nature and general applicability to various medical image segmentation scenarios. Our released code is available at \url{https://github.com/Mengqi-Lei/ConDSeg}.

연구 동기 및 목표

• 의료 영상에서 모호한 경계와 저대비도로 인해 발생하는 분할 정확도 저하 문제를 해결한다.
• 유사한 크기의 병변 간에 발생하는 공존 편향을 극복한다. 이는 모델이 유사한 크기의 병변 간에 인위적인 상관관계를 학습하는 것을 의미한다.
• 특징 추출 과정에서 조명 및 대비 변화에 대한 모델의 강건성을 향상시킨다.
• 크기 인지 디코더 아키텍처를 통해 다양한 크기의 구조를 정밀하게 국소화할 수 있도록 한다.
• 다양한 의료 영상 모odalities와 분할 작업에 적용 가능한 일반화 가능한 프레임워크를 개발한다.

제안 방법

• 일관성 강화(Consistency Reinforcement, CR)를 도입하여 원본 입력과 강력하게 증강된 입력에서의 예측 간 일관성을 극대화하는 이중 단계 사전 훈련 전략을 제안함으로써 인코더의 강건성을 향상시킨다.
• 의미 정보 분리(Semantic Information Decoupling, SID) 모듈을 제안하여 인코더 특징을 전경, 배경, 불확실성 영역으로 분리하고, 불확실성 감소와 마스크 정밀도 향상을 위한 손실 함수를 적용한다.
• 대비 기반 특징 집합(Contrast-Driven Feature Aggregation, CDFA) 모듈을 설계하여 전경과 배경 특징 간의 대비를 대조함으로써 다수준 특징 융합을 이끌어내고, 분할 성능 향상을 위한 핵심 특징을 강화한다.
• 다양한 크기 범위에 대해 별도의 헤드를 사용하는 크기 인지 디코더(Size-Aware Decoder, SA-Decoder)를 구현하여 다양한 크기의 객체에 대한 국소화 정확도를 향상시킨다.
• 이중 단계 훈련 파이프라인을 적용한다: 첫 번째 단계에서 CR을 사용해 인코더를 사전 훈련하고, 두 번째 단계에서 전체 네트워크를 미세조정함으로써 수렴 속도와 성능을 향상시킨다.
• SID 및 CDFA 모듈을 최적화하기 위해 교차 엔트로피 손실과 Dice 손실을 조합하고, 상호 보완적 제약 조건을 적용한다.

실험 결과

연구 질문

• RQ1 대비 기반 사전 훈련 전략은 의료 영상에서 저대비도 및 조명 변화에 대한 인코더의 강건성을 향상시킬 수 있는가?
• RQ2 전경, 배경, 불확실성 영역으로 특징을 분리함으로써 모델의 불확실성 감소와 분할 정확도 향상이 가능한가?
• RQ3 전경과 배경 특징 간의 대비 정보를 활용하면 다수준 특징 융합과 분할 성능 향상이 가능한가?
• RQ4 크기 인지 디코더 아키텍처는 척도 특이성 문제를 완화하고 다양한 크기의 해부학적 구조의 국소화를 향상시킬 수 있는가?
• RQ5 모든 제안된 구성 요소를 통합하면 다양한 의료 영상 분할 벤치마크에서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• ConDSeg는 Kvasir-SEG, GlaS 등 다섯 가지 의료 영상 분할 데이터셋에서 최신 기술 수준의 성능을 기록하였으며, Kvasir-SEG 데이터셋에서 평균 mIoU는 84.6%이고 mDSC는 90.5%였다.
• 사전 훈련 없이 초기 학습을 수행한 경우 대비, 일관성 강화 전략을 통해 mIoU가 7.3%p 향상되었다.
• 첫 번째 단계에서 CR을 적용한 이중 단계 훈련 전략은 단일 단계 훈련보다 빠른 수렴 속도와 더 높은 성능을 달성했다.
• 기본 모델 대비 SID 및 CDFA 모듈을 추가함으로써 GlaS 데이터셋에서 mIoU는 5.0%p, mDSC는 2.9%p 향상되었다.
• 완전한 ConDSeg 프레임워크(사용자 정의 SA-Decoder 포함)는 GlaS 데이터셋에서 최고의 mIoU(85.1%)와 mDSC(91.6%)를 기록하여 모든 추상화된 변형보다 뛰어난 성능을 보였다.
• 추상화 연구를 통해 CR, SID, CDFA, SA-Decoder 각 구성 요소가 성능 향상에 점진적으로 기여하며, 전체 모델이 최적의 성능을 달성하는 것으로 확인되었다.