Variational Deep Unfolding with Mamba-Based Nonlocal Modeling for Underwater Image Enhancement

Definição do Problema

O imageamento subaquático é crítico para a engenharia oceânica e a pesquisa marinha, mas os dados capturados frequentemente sofrem com baixa visibilidade, iluminação não uniforme, atenuação de cor e ruído. Essas degradações dificultam significativamente as tarefas subsequentes de visão computacional.

• Métodos Clássicos: As abordagens tradicionais dividem-se em técnicas sem modelo (manipulando intensidades de pixels sem suposições físicas) e métodos baseados em modelos físicos (resolvendo problemas inversos baseados no modelo de formação de imagem subaquática, tipicamente o modelo de Jaffe-McGlamery). Embora os métodos variacionais ofereçam interpretabilidade, sua eficácia depende fortemente da seleção manual de priors adequados.
• Métodos de Aprendizado Profundo (Deep Learning): Abordagens baseadas em dados aprendem priors diretamente dos dados, mas frequentemente exigem arquiteturas complexas e grandes conjuntos de dados. Modelos puramente baseados em dados podem carecer de interpretabilidade, enquanto arquiteturas padrão de aprendizado profundo (como Transformers) podem incorrer em altos custos computacionais para capturar dependências de longo alcance.
• A Lacuna: Há uma necessidade de preencher a lacuna entre a interpretabilidade da otimização baseada em modelos e a flexibilidade do aprendizado profundo, abordando especificamente o desafio de capturar autossimilaridades de longo alcance em cenas subaquáticas de forma eficiente.

Metodologia

Os autores propõem uma rede de unfolding profundo baseado em modelo que integra a modelagem variacional em uma arquitetura aprendível. A abordagem consiste em três componentes principais: um baseline variacional, uma arquitetura de unfolding profundo e uma estratégia de treinamento especializada.

1. Baseline Variacional

O método fundamenta-se em uma formulação variacional baseada em uma decomposição de desfoque (dehazing). Os autores propõem um modelo de formação de imagem modificado:
$$I = (J + N)t + A(1 - t)$$
onde $I$ é a imagem degradada, $J$ é a imagem restaurada, $N$ é um componente residual multiplicativo (para absorver artefatos remanescentes), $t$ é o mapa de transmissão e $A$ é a luz de retroespalhamento global.

O funcional de energia a ser minimizado inclui:

• Fidelidade de Dados: Baseado no modelo de decomposição acima.
• Regularização Não Local ($\alpha$-term): Uma restrição de gradiente não local para preservar detalhes finos e texturas.
• Variação Total ($\beta$-term): Aplicada ao mapa de transmissão para garantir suavidade espacial enquanto preserva descontinuidades de profundidade.
• Restrição de Resíduo ($\lambda$-term): Limita a absorção de componentes residuais.
• Fidelidade de Gradiente Não Local ($\mu$-term): Penaliza a distância entre o gradiente da imagem restaurada e um campo vetorial específico $V$ (derivado do gradiente amplificado da entrada degradada). Este termo é projetado para fortalecer a informação estrutural oculta em imagens subaquáticas de baixo contraste.
• Fidelidade de Transmissão ($\rho$-term): Garante que o mapa de transmissão estimado permaneça próximo a uma estimativa inicial (derivada do Dark Channel Prior).

Os autores fornecem uma prova teórica estabelecendo a existência de um minimizador para este problema variacional dentro de um espaço admissível definido.

2. Rede de Unfolding Profundo (Deep Unfolding Network)

O algoritmo de otimização (uma estratégia primal-dual baseada em Chambolle–Pock) é "desdobrado" (unfolded) em uma rede neural profunda com $K$ estágios. Neste framework, as atualizações iterativas das variáveis de otimização ($J, t, N$) correspondem às camadas da rede.

• Priors Aprendíveis: Os operadores proximais associados ao prior da imagem ($J$) e ao mapa de transmissão ($t$) são substituídos por redes neurais aprendíveis.
• Integração Mamba: Para capturar eficientemente dependências de longo alcance e autossimilaridades sem o alto custo computacional da autoatenção, os autores integram camadas Mamba (Modelos de Espaço de Estados).
  • MambaResNet: Substitui o operador proximal para a imagem $J$. Utiliza um design de múltiplos ramos processando a imagem de entrada, uma versão com balanço de branco aprimorado e uma versão com equalização de histograma. Emprega blocos de Varredura Seletiva 2D (SS2D) para modelar interações não locais.
  • ProxNet: Uma rede residual leve para o mapa de transmissão $t$.
  • MambaNet: Um módulo que modela o termo de fidelidade de gradiente não local, operando no domínio do gradiente.
• Inicialização: A rede é inicializada com a imagem degradada para $J$, zero para $N$, e a estimativa do Dark Channel Prior para $t$.

3. Perda de Trajetória Proximal (Proximal Trajectory Loss)

Para garantir que a rede de unfolding adira à lógica de otimização subjacente, os autores introduzem uma perda de trajetória proximal. Em vez de supervisionar apenas a saída final, esta perda impõe consistência entre os estágios intermediários da rede e as iterações de um regularizador de restauração "ideal" (derivado de um problema de otimização simplificado envolvendo a verdade fundamental/ground truth). Isso evita que os estágios intermediários desviem do caminho de otimização esperado.

Principais Contribuições

1. Formulação Variacional: Um novo funcional de energia para o realce subaquático incorporando um componente residual multiplicativo e uma restrição do tipo gradiente não local para lidar com artefatos e preservar bordas.
2. Análise Teórica: Uma prova da existência de minimizadores para o proposto problema variacional não convexo.
3. Unfolding Baseado em Mamba: Uma rede de unfolding profundo que substitui operadores proximais por módulos baseados em Mamba, permitindo a modelagem eficiente de dependências de longo alcance em cenas subaquáticas.
4. Perda de Trajetória Proxicional: Uma função de perda inovadora projetada para impor consistência entre os estágios de unfolding da rede e as iterações do algoritmo de otimização ideal, melhorando a interpretabilidade e o desempenho.

Resultados Experimentais

O método foi avaliado nos conjuntos de dados UIEB e SUIM-E, comparando-o com métodos clássicos (ex: CBAF, BayesianRetinex), métodos de aprendizado profundo puro (ex: UColor, Phaseformer, Mamba-UIE) e outros métodos de unfolding (UnfoldNet).

• Desempenho Quantitativo: O método de unfolding proposto alcançou resultados de estado da arte em todas as métricas (PSNR, SSIM, LPIPS) em ambos os conjuntos de dados. Por exemplo, no UIEB, alcançou um PSNR de 25.16 e SSIM de 0.9140, superando o segundo melhor método (Mamba-UIE) por uma margem significativa.
• Qualidade Visual: Comparações visuais mostraram que o método proposto recuperou melhor a precisão das cores (reduzindo tons avermelhados/esverdeados comuns em outros métodos) e preservou detalhes finos (ex: escamas de peixes, texturas do leito marinho) em comparação aos concorrentes.
• Estudos de Ablação:
  • Dehazing vs. Retinex: A formulação baseada em dehazing superou uma alternativa baseada em Retinex.
  • Restrição de Gradiente: A remoção da restrição de gradiente não local resultou em menor contraste de borda e qualidade perceptual.
  • Perda de Trajetória: A supervisão dos estágios intermediários via perda de trajetória proximal melhorou significativamente as métricas finais em comparação ao uso apenas de MSE na saída final.
  • Eficiência: Comparações com Multi-Head Attention (MHA) mostraram que a arquitetura baseada em Mamba alcançou menores FLOPs e uso de memória GPU, particularmente em resoluções mais altas, mantendo um maior throughput.

Significância e Alegações

O artigo afirma que o framework proposto consegue preencher a lacuna entre a otimização baseada em modelos e o aprendizado profundo. Ao integrar um modelo variacional motivado fisicamente com a eficiência dos Modelos de Espaço de Estados (Mamba), os autores alcançam:

• Restauração Melhorada: Qualidade visual superior e desempenho quantitativo em comparação com métodos recentes de estado da arte.
• Interpretabilidade: A estrutura de unfolding preserva a lógica do problema de otimização subjacente, potencializada pela garantia teórica de existência de solução.
• Eficiência: O uso de camadas Mamba permite a captura eficiente de contexto global (modelagem não local) com complexidade linear, abordando os gargalos computacionais de abordagens baseadas em Transformers na imagem subaquática.

Os autores observam que, embora os resultados sejam promissores, trabalhos futuros podem focar em melhorar a estimativa inicial de transmissão (atualmente dependente do Dark Channel Prior) e estender o modelo para realce de vídeo e consistência temporal.