[논문 리뷰] Boosting Dermatoscopic Lesion Segmentation via Diffusion Models with Visual and Textual Prompts

Shiyi Du, Xiaosong Wang외 6명|arXiv (Cornell University)|2023. 10. 04.

Cutaneous Melanoma Detection and Management인용 수 4

한 줄 요약

이 논문은 시각적 및 텍스트 프롬프트를 활용하여 고해상도의 피부병변 영상 생성을 위한 확산 모델 기반 프레임워크를 제안한다. 이는 병변의 정확한 특성과 함께 높은 정밀도를 달성하여 피부 병변 분할 성능을 크게 향상시킨다. 병변 유형, 형태, 위치 및 특성 등을 프롬프트로 통합함으로써 기존의 GAN 기반 접근법 대비 SSIM에서 9% 향상되고, Dice 계수에서도 5% 이상 향상된다.

ABSTRACT

Image synthesis approaches, e.g., generative adversarial networks, have been popular as a form of data augmentation in medical image analysis tasks. It is primarily beneficial to overcome the shortage of publicly accessible data and associated quality annotations. However, the current techniques often lack control over the detailed contents in generated images, e.g., the type of disease patterns, the location of lesions, and attributes of the diagnosis. In this work, we adapt the latest advance in the generative model, i.e., the diffusion model, with the added control flow using lesion-specific visual and textual prompts for generating dermatoscopic images. We further demonstrate the advantage of our diffusion model-based framework over the classical generation models in both the image quality and boosting the segmentation performance on skin lesions. It can achieve a 9% increase in the SSIM image quality measure and an over 5% increase in Dice coefficients over the prior arts.

연구 동기 및 목표

• 희귀 병변 유형 및 특성에 대해 애초에 부족한 레이블이 부여된 피부병변 영상 문제를 해결하기 위해 다양한 고품질의 합성 병변 영상을 생성함으로써 해결하고자 한다.
• 병변 특화된 시각적 및 텍스트 프롬프트를 통합함으로써 의료 영상에서의 이미지 생성의 제어 가능성을 향상시키고자 한다.
• 제안된 생성 프레임워크를 활용한 데이터 증강을 통해 후속 피부 병변 분할 성능을 향상시키고자 한다.
• 확산 모델이 GAN보다 더 현실적이고 세밀한 피부병변 영상을 정확한 병변 특성과 함께 생성하는 데서 우월함을 입증하고자 한다.
• 확장 가능한 데이터 증강을 위해 자동 병변 형태 및 마스크 생성 모듈을 개발하고자 한다.

제안 방법

• 이미지 생성의 기반으로 Stable Diffusion 모델을 사용하며, 병변의 세분화 마스크(시각적 프롬프트)와 병변 유형 및 특성(텍스트 프롬프트)을 조건으로 삼는다.
• ControlNet을 활용해 확산 과정을 안내함으로써 병변의 형태, 위치, 외관에 대한 정밀한 제어를 가능하게 한다.
• 다양한 크기, 형태, 해부학적 위치를 가진 병변 세분화 마스크를 자동으로 생성하는 모듈을 도입하여 학습 데이터의 다양성을 풍부화시킨다.
• 공개된 피부병변 데이터셋(ISIC)을 사용해 모델을 훈련시키며, ISIC 2017 및 2019의 병변 유형 및 특성 정보를 텍스트 프롬프트로 활용한다.
• 생성된 이미지를 후속 병변 분할 작업의 증강 학습 데이터로 활용하며, Dice 손실 함수와 U-Net 기반 분할 네트워크를 적용한다.
• 동일한 학습 및 평가 프rotocol를 기반으로, 강력한 GAN 기반 기준 모델인 Pix2PixHD와의 비교 및 분석 실험을 수행한다.

실험 결과

연구 질문

• RQ1 시각적 및 텍스트 프롬프트를 활용한 확산 모델이 GAN보다 더 높은 정밀도와 병변 특성에 대한 더 나은 제어를 통해 피부병변 영상을 생성할 수 있는가?
• RQ2 제안된 생성 프레임워크를 활용한 데이터 증강이 피부 병변 분할 성능 향상에 어느 정도 기여하는가?
• RQ3 병변 특화 프롬프트(유형, 형태, 위치, 특성)의 통합이 영상 품질과 후속 작업 정확도에 어떤 영향을 미치는가?
• RQ4 자동 병변 마스크 생성 모듈이 합성 병변 데이터의 다양성과 커버리지 증가에 효과적으로 기여하는가?
• RQ5 제안된 프레임워크는 희귀 또는 비정상적인 경우를 제어적으로 생성이 필요한 다른 의료 영상 작업에 이식 가능한가?

주요 결과

• 제안된 확산 모델 기반 프레임워크는 Pix2PixHD GAN 기준 대비 SSIM에서 9% 상대적 향상을 달성하여 더 뛰어난 영상 품질과 세부 정보 유지 능력을 보여준다.
• 기존 기술 대비 Dice 계수는 5% 이상 향상되었으며, 5,000장의 합성 영상을 사용할 경우 최종 분할 DSC는 0.914를 기록한다.
• 시각적으로 확인된 바, 생성된 영상에서는 질감과 병변 구조의 정확성이 뚜렷이 향상되어 있으며, 그림 4의 확대 영역에서 뚜렷하게 확인된다.
• 1,000장, 3,000장, 5,000장의 다양한 합성 데이터 스케일에서 성능 향상이 일관되게 유지되며, 3,000장 이상에서는 추가적인 이점이 최소화된다.
• 합성 데이터 증강을 적용했을 때, 최신 기술(U-Net, DSC 0.861)과 DCSAU-Net(DSC 0.903)을 모두 초월하는 성능을 기록한다.
• 시각적 및 텍스트 프롬프트를 동시에 사용함으로써 병변 특성에 대한 정밀한 제어가 가능해져 희귀 또는 복잡한 경우에 특화된 데이터 생성이 가능해진다.