[논문 리뷰] Learning A Physical-aware Diffusion Model Based on Transformer for Underwater Image Enhancement

Chen Zhao, Chenyu Dong외 2명|arXiv (Cornell University)|2024. 03. 03.

Underwater Acoustics Research인용 수 10

한 줄 요약

PA-Diff는 물속 영상 향상을 위한 물리 가이드 확산 프레임워크를 도입하여, Physics Prior Generation 분기와 Physics-aware Diffusion Transformer를 결합해 수중 이미징 물리 및 확산 모델링을 활용합니다.

ABSTRACT

Underwater visuals undergo various complex degradations, inevitably influencing the efficiency of underwater vision tasks. Recently, diffusion models were employed to underwater image enhancement (UIE) tasks, and gained SOTA performance. However, these methods fail to consider the physical properties and underwater imaging mechanisms in the diffusion process, limiting information completion capacity of diffusion models. In this paper, we introduce a novel UIE framework, named PA-Diff, designed to exploiting the knowledge of physics to guide the diffusion process. PA-Diff consists of Physics Prior Generation (PPG) Branch, Implicit Neural Reconstruction (INR) Branch, and Physics-aware Diffusion Transformer (PDT) Branch. Our designed PPG branch aims to produce the prior knowledge of physics. With utilizing the physics prior knowledge to guide the diffusion process, PDT branch can obtain underwater-aware ability and model the complex distribution in real-world underwater scenes. INR Branch can learn robust feature representations from diverse underwater image via implicit neural representation, which reduces the difficulty of restoration for PDT branch. Extensive experiments prove that our method achieves best performance on UIE tasks.

연구 동기 및 목표

• 수중 물리적 영상 메커니즘을 확산 모델에 통합하여 UIE 향상 동기를 제시한다.
• 물리 프리어와 확산 기반 복원을 결합하는 이중 분기 아키텍처를 제안한다.
• 물리 인식 트랜스포머를 통한 장거리 의존성 모델링 및 영역 인식 향상을 구현한다.
• 실세계 수중 데이터셋에서 최첨단 성능을 입증하고 기여를 입증하는 차등 가능한 평가를 수행한다.

제안 방법

• 두 분기 PA-Diff 아키텍처: Physics Prior Generation (PPG) 분기와 Physics-aware Diffusion Transformer (PDT) 분기.
• PPG는 수정된 Koschmieder light scanning model을 통해 transmission maps와 global background light를 생성하며, 물리 프리어가 재구성 손실을 통해 확산을 가이드한다.
• PDT는 수중 입력과 물리 프리어에 조건부로 확산 과정(DDPM 프레임워크)을 사용하며, Physics-aware Self-Attention (PA-SA) 및 Multi-Scale Dynamic FFN (MS-FFN)을 포함한다.
• PA-SA는 전송(prior)과 광(prior)을 이용해 특징을 보정하여 장거리 확산 모델링을 위한 transmission-aware 및 light-aware 특징을 생성한다.
• MS-FFN은 다중 스케일 깊이별 컨볼루션을 사용해 수용영역을 확장하고 지역 특징을 집계한다.
• 훈련은 조건부 확산 과정에서 추정된 노이즈를 실제 노이즈와 일치시키는 확산 손실을 최적화한다.

실험 결과

연구 질문

• RQ1 물리적 수중 이미징 지식을 확산 기반 UIE 프레임워크에 어떻게 통합할 수 있는가?
• RQ2 전달 맵과 배경광(priors) 정보를 포함시키면 수중 이미지의 확산 기반 복원 성능이 향상되는가?
• RQ3 물리 인식 어텐션이 포함된 트랜스포머 기반 확산 모듈이 수중 장면의 장거리 의존성을 더 잘 모델링하는가?
• RQ4 각 물리 유도 컴포넌트(전달 맵, 배경광, MS-FFN)의 UIE 성능에 대한 기여도는 무엇인가?

주요 결과

Dataset	FID ↓	LPIPS ↓	PSNR ↑	SSIM ↑
UIEBD (Ours)	28.76	0.1324	21.14	0.8620
LSUI (Ours)	22.15	0.0923	25.89	0.9354

• PA-Diff가 UIE 벤치마크(UIEBD 및 LSUI)에서 최첨단 주관적 및 충실도 지표를 달성한다.
• 추적 연구에서 전달 프리어(T)가 배경광(B)만 사용하는 것보다 더 큰 이득을 제공하고, MS-FFN 및 PA-SA와 함께 두 프리어를 완전히 통합하면 최상의 성능을 달성한다.
• LSUI에서 제안된 방법이 여러 베이스라인 대비 PSA 개선을 달성하고 FID가 낮고 SSIM이 더 높다.
• 발췌에 따르면 생성된 물리 프리어가 유익하나 정확한 프리어의 정확도가 최종 이득에 영향을 미쳐 물리 프리어 추정의 개선 여지가 있음을 시사한다.