[論文レビュー] WaterMamba: Visual State Space Model for Underwater Image Enhancement

Meisheng Guan, Haiyong Xu+5 more|arXiv (Cornell University)|May 14, 2024

Image Enhancement Techniques被引用数 6

ひとこと要約

WaterMambaはSCOSSブロックを備えた線形計算量の視覚状態空間モデルを導入し、 underwater image enhancement を効率的に実行、複数のデータセットで最先端のPSNR/SSIMを達成し、パラメータ数とFLOPsを抑制。

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-NetバックボーンへSCOSSブロックを組み込み。
• 複数のベンチマークデータセットでSOTA手法と比較して効率性と有効性を示す。

提案手法

• 線形複雑度の構造化状態空間モデル（SSM）を画像強調（WaterMamba）に採用。
• SCOSSブロックはSCCOSS（空間）とCOOSS（チャネル）モジュールおよびマルチスケール特徴処理のMSFFNを含む。
• SCCOSSでは、長距離空間依存性をモデル化するために4方向系列を用いた2D選択スキャンを実装。
• COOSSでは、空間プーリングを組み合わせた座標ベースのチャネルアテンションを適用し、Mambaベースのチャネル処理を行う。
• MSFFNを用いてマルチスケール特徴を融合し、パターン崩壊を回避し、最後に残差学習を適用する（）。
• WaterMambaをU-Netライクのエンコーダ-デコーダにスキップ接続と残差出力 DR+I を持たせて組み立てる。

実験結果

リサーチクエスチョン

• 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（UIEB由来データセットのTable 1）を達成。
• UIEBライクな評価でWaterMambaは競争力のあるUIQMとUCIQEスコアを達成（UIQM 0.555; UCIQE 2.767はSQUIDベースの指標）。
• WaterMambaは3.69Mパラメータと7.53G FLOPsを使用し、RestormerのようなTransformerベースのベースラインと比較して複雑さを大幅に削減。
• 定性的結果は、さまざまな水中シーン（サンゴ礁や地形）でより鮮明なディテールと自然な色再現を示す。
• WaterMambaは複数データセット（UIEB、UCIOD、UCCS、SQUID）でPSNR/SSIMの点で10のSOTA手法を上回り、計算コストが低い。