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

Planteamiento del Problema

La obtención de imágenes submarinas es crítica para la ingeniería oceánica y la investigación marina, sin embargo, los datos capturados frecuentemente sufren de baja visibilidad, iluminación no uniforme, atenuación de color y ruido. Estas degradaciones dificultan significativamente las tareas posteriores de visión computacional.

• Métodos Clásicos: Los enfoques tradicionales se dividen en técnicas libres de modelos (que manipulan las intensidades de los píxeles sin supuestos físicos) y métodos basados en modelos físicos (que resuelven problemas inversos basados en el modelo de formación de imagen submarina, típicamente el modelo de Jaffe-McGlamery). Si bien los métodos variacionales ofrecen interpretabilidad, su efectividad depende fuertemente de la selección manual de priors adecuados.
• Métodos de Aprendizaje Profundo: Los enfoques de aprendizaje profundo basados puramente en datos aprenden los priors directamente de los datos, pero a menudo requieren arquitecturas complejas y grandes conjuntos de datos. Los modelos puramente basados en datos pueden carecer de interpretabilidad, mientras que las arquitecturas de aprendizaje profundo estándar (como los Transformers) pueden incurrir en altos costos computacionales para capturar dependencias de largo alcance.
• La Brecha: Existe la necesidad de cerrar la brecha entre la interpretabilidad de la optimización basada en modelos y la flexibilidad del aprendizaje profundo, abordando específicamente el desafío de capturar auto-similitudes de largo alcance en escenas submarinas de manera eficiente.

Metodología

Los autores proponen una red de desplegado profundo basado en modelos (model-based deep unfolding network) que integra el modelado variacional en una arquitectura aprendible. El enfoque consta de tres componentes principales: un baseline variacional, una arquitectura de desplegado profundo y una estrategia de entrenamiento especializada.

1. Baseline Variacional

El método se fundamenta en una formulación variacional basada en una descomposición de desniebla (dehazing). Los autores proponen un modelo de formación de imagen modificado:
$$I = (J + N)t + A(1 - t)$$
donde $I$ es la imagen degradada, $J$ es la imagen restaurada, $N$ es un componente residual multiplicativo (para absorber artefactos restantes), $t$ es el mapa de transmisión y $A$ es la luz de retrodispersión global.

El funcional de energía a minimizar incluye:

• Fidelidad de Datos: Basada en el modelo de descomposición anterior.
• Regularización No Local ($\alpha$-term): Una restricción de gradiente no local para preservar detalles finos y texturas.
• Variación Total ($\beta$-term): Aplicada al mapa de transmisión para asegurar la suavidad espacial preservando las discontinuidades de profundidad.
• Restricción de Residuo ($\lambda$-term): Limita la absorción de componentes residuales.
• Fidelidad de Gradiente No Local ($\mu$-term): Penaliza la distancia entre el gradiente de la imagen restaurada y un campo vectorial específico $V$ (derivado del gradiente amplificado de la entrada degradada). Este término está diseñado para fortalecer la información estructural oculta en imágenes submarinas de bajo contraste.
• Fidelidad de Transmisión ($\rho$-term): Asegura que el mapa de transmisión estimado se mantenga cerca de una estimación inicial (derivada del Dark Channel Prior).

Los autores proporcionan una prueba teórica que establece la existencia de un minimizador para este problema variacional dentro de un espacio admisible definido.

2. Red de Desplegado Profundo (Deep Unfolding Network)

El algoritmo de optimización (una estrategia primal-dual basada en Chambolle–Pock) se "despliega" en una red neuronal profunda con $K$ etapas. En este marco, las actualizaciones iterativas de las variables de optimización ($J, t, N$) corresponden a las capas de la red.

• Priors Aprendibles: Los operadores proximales asociados con el prior de la imagen ($J$) y el mapa de transmisión ($t$) son reemplazados por redes neuronales aprendibles.
• Integración de Mamba: Para capturar eficientemente dependencias de largo alcance y auto-similitudes sin el alto costo computacional de la auto-atención, los autores integran capas Mamba (Modelos de Espacio de Estados).
  • MambaResNet: Reemplaza el operador proximal para la imagen $J$. Utiliza un diseño de múltiples ramas que procesa la imagen de entrada, una versión mejorada por balance de blancos y una versión con ecualización de histograma. Emplea bloques de Escaneo Selectivo 2D (SS2D) para modelar interacciones no locales.
  • ProxNet: Una red residual ligera para el mapa de transmisión $t$.
  • MambaNet: Un módulo que modela el término de fidelidad de gradiente no local, operando en el dominio del gradiente.
• Inicialización: La red se inicializa con la imagen degradada para $J$, cero para $N$, y la estimación del Dark Channel Prior para $t$.

3. Pérdida de Trayectoria Proximal (Proximal Trajectory Loss)

Para asegurar que la red de desplegado se adhiera a la lógica de optimización subyacente, los autores introducen una pérdida de trayectoria proximal. En lugar de supervisar únicamente la salida final, esta pérdida impone consistencia entre las etapas intermedias de la red y las iteraciones de un regularizador de restauración "ideal" (derivado de un problema de optimización simplificado que involucra la verdad de campo o ground truth). Esto evita que las etapas intermedias se desvíen del camino de optimización esperado.

Contribuciones Clave

1. Formulación Variacional: Un nuevo funcional de energía para la mejora submarina que incorpora un componente residual multiplicativo y una restricción de tipo gradiente no local para manejar artefactos y preservar bordes.
2. Análisis Teórico: Una prueba de la existencia de minimizadores para el propuesto problema variacional no convexo.
3. Desplegado Basado en Mamba: Una red de desplegado profundo que reemplaza los operadores proximales con módulos basados en Mamba, permitiendo el modelado eficiente de dependencias de largo alcance en escenas submarinas.
4. Pérdida de Trayectoria Proximal: Una función de pérdida novedosa diseñada para imponer consistencia entre las etapas de despliegue de la red y las iteraciones del algoritmo de optimización ideal, mejorando la interpretabilidad y el rendimiento.

Resultados Experimentales

El método fue evaluado en los conjuntos de datos UIEB y SUIM-E, comparándolo con métodos clásicos (ej. CBAF, BayesianRetinex), métodos de aprendizaje profundo puro (ej. UColor, Phaseformer, Mamba-UIE) y otros métodos de desplegado (UnfoldNet).

• Rendimiento Cuantitativo: El método de desplegado propuesto alcanzó resultados de vanguardia (state-of-the-art) en todas las métricas (PSNR, SSIM, LPIPS) en ambos conjuntos de datos. Por ejemplo, en UIEB, logró un PSNR de 25.16 y un SSIM de 0.9140, superando significativamente al segundo mejor método (Mamba-UIE).
• Calidad Visual: Las comparaciones visuales mostraron que el método propuesto recuperó mejor la precisión del color (reduciendo los tonos rojizos/verdosos comunes en otros métodos) y preservó detalles finos (ej. escamas de peces, texturas del lecho marino) en comparación con los competidores.
• Estudios de Ablación:
  • Desniebla vs. Retinex: La formulación basada en desniebla superó a una alternativa basada en Retinex.
  • Restricción de Gradiente: Eliminar la restricción de gradiente no local resultó en un menor contraste de bordes y calidad perceptual.
  • Pérdida de Trayectoria: La supervisión de las etapas intermedias mediante la pérdida de trayectoria proximal mejoró significativamente las métricas finales en comparación con el uso exclusivo de MSE en la salida final.
  • Eficiencia: Las comparaciones con la Atención de Cabezales Múltiples (MHA) mostraron que la arquitectura basada en Mamba logró menores FLOPs y uso de memoria GPU, particularmente en resoluciones altas, manteniendo al mismo tiempo un mayor rendimiento (throughput).

Significancia y Afirmaciones

El artículo afirma que el marco propuesto logra cerrar con éxito la brecha entre la optimización basada en modelos y el aprendizaje profundo. Al integrar un modelo variacional motivado físicamente con la eficiencia de los Modelos de Espacio de Estados (Mamba), los autores logran:

• Mejor Restauración: Calidad visual y rendimiento cuantitativo superiores en comparación con métodos recientes de vanguardia.
• Interpretabilidad: La estructura de desplegado preserva la lógica del problema de optimización subyacente, potenciada por la garantía teórica de la existencia de la solución.
• Eficiencia: El uso de capas Mamba permite capturar el contexto global (modelado no local) de manera eficiente con complejidad lineal, abordando los cuellos de botella computacionales de los enfoques basados en Transformers en la imagenología submarina.

Los autores señalan que, aunque los resultados son prometedores, el trabajo futuro podría centrarse en mejorar la estimación inicial de la transmisión (actualmente dependiente del Dark Channel Prior) y extender el modelo a la mejora de video y la consistencia temporal.