[논문 리뷰] O-Mamba: O-shape State-Space Model for Underwater Image Enhancement

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

Image Enhancement Techniques인용 수 5

한 줄 요약

O-Mamba는 Spatial Mamba와 Channel Mamba로 구성된 O자형 이중 분기 네트워크를 도입하여 해저 영상 향상을 위한 공간 정보와 채널 간 정보를 모델링하고, 다중 스케일 상호 촉진(MSBMP) 모듈로 다중 스케일 특징을 융합하여 여러 데이터셋에서 SOTA를 달성한다.

ABSTRACT

Underwater image enhancement (UIE) face significant challenges due to complex underwater lighting conditions. Recently, mamba-based methods have achieved promising results in image enhancement tasks. However, these methods commonly rely on Vmamba, which focuses only on spatial information modeling and struggles to deal with the cross-color channel dependency problem in underwater images caused by the differential attenuation of light wavelengths, limiting the effective use of deep networks. In this paper, we propose a novel UIE framework called O-mamba. O-mamba employs an O-shaped dual-branch network to separately model spatial and cross-channel information, utilizing the efficient global receptive field of state-space models optimized for underwater images. To enhance information interaction between the two branches and effectively utilize multi-scale information, we design a Multi-scale Bi-mutual Promotion Module. This branch includes MS-MoE for fusing multi-scale information within branches, Mutual Promotion module for interaction between spatial and channel information across branches, and Cyclic Multi-scale optimization strategy to maximize the use of multi-scale information. Extensive experiments demonstrate that our method achieves state-of-the-art (SOTA) results.The code is available at https://github.com/chenydong/O-Mamba.

연구 동기 및 목표

• 복잡한 조명, 산란, 색상 왜곡으로 인한 해저 영상 저하 문제를 다룬다.
• 공간 정보와 채널 간 정보를 각각 모델링하기 위한 이중 분기 아키텍처를 개발한다.
• MSBMP를 통해 교차 분기 상호작용 및 다중 스케일 정보 융합을 가능케 하여 UIE 성능을 향상시킨다.

제안 방법

• 공간 Mamba 분기(SM)와 채널 Mamba 분기(CM)를 갖는 O자형 이중 분기 네트워크를 제안한다.
• 공간 Mamba 블록과 채널 Mamba 블록을 사용하여 상태 공간 모델링(SSM)로 공간 및 교차 채널 의존성을 포착한다.
• 다중 스케일 모듈(MS-MoE, Mutual Promotion, Cyclic Multi-scale optimization(CMS))으로 구성된 다중 스케일 Bi-mutual Promotion(MSBMP) 모듈을 도입한다.
• 4방향의 공간 방향에 따른 2D 선택적 스캐닝(S-SSM)과 교차 채널 스캐닝 모듈(C-SSM)로 SSM을 확장한다.
• 다양한 해저 조건을 처리하기 위해 Mamba 블록 내 피드포워드 네트워크로 FF-MoE를 사용한다.
• 현재 분기에서의 전문가 블록을 활용하여 각 분기 내에서 다중 스케일 특징을 융합하기 위해 MS-MoE를 구현한다.
• 학습 중 스케일 간 간섭을 완화하기 위해 순환 다중 스케일 손실(CMS-OS)을 적용한다.

실험 결과

연구 질문

• RQ1 이중 분기 Mamba 기반 아키텍처가 UIE를 위해 공간 정보와 채널 간 정보를 효과적으로 모델링할 수 있는가?
• RQ2 MSBMP를 통한 다중 스케일, 교차 분기 상호작용이 단일 분기 또는 비-MSBMP 변종보다 복원 품질을 향상시키는가?
• RQ3 순환 다중 스케일 최적화가 학습 안정성 및 스케일 간 성능에 어떤 영향을 미치는가?
• RQ4 O-Mamba가 여러 해저 벤치마크에서 최첨단(UIE) 방법과 비교하여 어떤 성능 차이를 보이는가?

주요 결과

• O-Mamba는 여러 UIE 데이터셋 및 지표에서 최첨단 성능을 달성한다.
• 특정 제거 실험에서 이중 분기 설계가 단일 분기 변형(O-Mamba-S, O-Mamba-C)보다 우수하다는 것이 확인된다.
• MP 모듈은 분기 간 정보 상호작용을 크게 개선하여 제거 실험에서 PSNR을 29.82에서 30.00으로 올린다.
• MS-MoE는 다중 스케일 정보 융합에 필수적이며 성능 향상에 기여한다.
• 순환 다중 스케일 최적화는 스케일 간 간섭을 줄이고 전반적인 결과를 개선한다.
• 다수의 SOTA UIE 방법과 비교할 때 O-Mamba는 LSUI 및 UIEB와 같은 데이터셋에서 정량적 결과가 우수함(연구에 보고된 바와 같이).