[논문 리뷰] WaterMamba: Visual State Space Model for Underwater Image Enhancement

Meisheng Guan, Haiyong Xu외 5명|arXiv (Cornell University)|2024. 05. 14.

Image Enhancement Techniques인용 수 6

한 줄 요약

WaterMamba는 SCOSS 블록을 갖춘 선형 복잡도 시각 상태 공간 모델을 도입하여 수중 영상 향상을 효율적으로 수행하고, 여러 데이터셋에서 더 적은 파라미터와 더 낮은 FLOPs로 최첨단 PSNR/SSIM을 달성합니다.

ABSTRACT

Underwater imaging often suffers from low quality due to factors affecting light propagation and absorption in water. To improve image quality, some underwater image enhancement (UIE) methods based on convolutional neural networks (CNN) and Transformer have been proposed. However, CNN-based UIE methods are limited in modeling long-range dependencies, and Transformer-based methods involve a large number of parameters and complex self-attention mechanisms, posing efficiency challenges. Considering computational complexity and severe underwater image degradation, a state space model (SSM) with linear computational complexity for UIE, named WaterMamba, is proposed. We propose spatial-channel omnidirectional selective scan (SCOSS) blocks comprising spatial-channel coordinate omnidirectional selective scan (SCCOSS) modules and a multi-scale feedforward network (MSFFN). The SCOSS block models pixel and channel information flow, addressing dependencies. The MSFFN facilitates information flow adjustment and promotes synchronized operations within SCCOSS modules. Extensive experiments showcase WaterMamba's cutting-edge performance with reduced parameters and computational resources, outperforming state-of-the-art methods on various datasets, validating its effectiveness and generalizability. The code will be released on GitHub after acceptance.

연구 동기 및 목표

• 비비정합적 저하 및 색 채널 손실이 있는 수중 영상에 대한 UIE 동기 부여.
• 상태 공간 모델링을 활용해 장거리 의존성을 포착하는 경량의 효율적인 UIE 모델 개발.
• SCOSS 블록을 설계해 공간 정보와 채널 정보를 공동으로 모델링하여 강건한 향상을 달성.
• U-네트 백본에 SCOSS 블록을 통합해 고품질의 수중 이미지를 재구성.
• 다양한 벤치마크 데이터셋에서 SOTA 메서드 대비 효율성과 효과를 입증하기 위해 평가

제안 방법

• 이미지 향상을 위한 선형 복잡도를 가지는 구조화된 상태 공간 모델(SSM)을 채용한 WaterMamba를 도입.
• 스 SCCOSS(공간)와 COOSS(채널) 모듈 plus MSFFN로 구성된 SCOSS 블록 및 다중 스케일 특징 처리를 제시.
• SCCOSS에서 2D 선택적 스캔을 네 방향 시퀀스로 구현해 장거리 공간 의존성을 모델링.
• COOSS에서 공간 풀링과 Mamba 기반 채널 처리를 결합한 좌표 기반 채널 어텐션을 적용.
• 패턴 붕괴를 방지하기 위해 다중 스케일 특징을 융합하고 잔차 학습을 마지막에 수행하는 MSFFN을 사용.
• Skip 연결과 잔차 출력을 포함한 DR+I 형태의 U-네트 like 인코더-디코더에 WaterMamba를 구성

실험 결과

연구 질문

• RQ1 WaterMamba가 선형 복잡도 SSM 기반 아키텍처로 최첨단 UIE 성능을 달성할 수 있는가?
• RQ2 SCOSS 블록이 비균일 저하 하에서 수중 영상의 공간 및 채널 의존성을 효과적으로 모델링하는가?
• RQ3 WaterMamba가 PSNR, SSIM, UIQM, UCIQE 및 계산 효율성 측면에서 CNN 기반 및 Transformer 기반 UIE 방법과 비교해 어떤 차이가 있는가?

주요 결과

• WaterMamba는 R90에서 PSNR/SSIM이 24.715/0.931, U100에서 21.992/0.843로 달성(표 1의 UIEB 유도 데이터셋).
• UIEB 유사 평가에서 WaterMamba는 경쟁력 있는 UIQM 및 UCIQE 점수를 달성(UIQM 0.555; SQUID 기반 메트릭에서 UCIQE 2.767).
• WaterMamba는 3.69M 매개변수와 7.53G FLOPs를 사용하므로 Restormer와 같은 Transformer 기반 기본 비교대상 대비 복잡도를 크게 줄임.
• 다양한 수중 장면(산호초 및 해안지형)에서 더 선명한 디테일과 더 자연스러운 색상 복원을 보여주는 정성적 결과.
• WaterMamba는 PSNR/SSIM에서 다수의 데이터셋(UIEB, UCIOD, UCCS, SQUID)에서 10개의 SOTA 방법을 능가하며 계산 비용이 더 낮음