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

Problem Statement (Enunciato del problema)

L'imaging subacqueo è fondamentale per l'ingegneria oceanica e la ricerca marina, tuttavia i dati catturati soffrono frequentemente di scarsa visibilità, illuminazione non uniforme, attenuazione del colore e rumore. Queste degradazioni ostacolano significativamente le successive attività di computer vision.

• Metodi Classici: Gli approcci tradizionali si dividono in tecniche model-free (che manipolano le intensità dei pixel senza assunzioni fisiche) e metodi basati su modelli fisici (che risolvono problemi inversi basati sul modello di formazione dell'immagine subacquea, tipicamente il modello di Jaffe-McGlamery). Sebbene i metodi variazionali offrano interpretabilità, la loro efficacia dipende fortemente dalla selezione manuale di prior adeguati.
• Metodi di Deep Learning: Gli approcci basati sui dati del deep learning apprendono i prior direttamente dai dati, ma spesso richiedono architetture complesse e grandi dataset. I modelli puramente data-driven possono mancare di interpretabilità, mentre le standard architetture di deep learning (come i Transformer) possono comportare elevati costi computazionali per la cattura di dipendenze a lungo raggio.
• Il Gap: Esiste la necessità di colmare il divario tra l'interpretabilità dell'ottimizzazione basata su modelli e la flessibilità del deep learning, affrontando specificamente la sfida della cattura efficiente delle auto-similarità a lungo raggio nelle scene subacquee.

Methodology (Metodologia)

Gli autori propongono una rete di deep unfolding basata su modello che integra la modellazione variazionale in un'architettura apprendibile. L'approccio consiste in tre componenti principali: un baseline variazionale, un'architettura di deep unfolding e una strategia di addestramento specializzata.

1. Variational Baseline (Baseline variazionale)

Il metodo si fonda su una formulazione variazionale basata su una decomposizione di de-hazing. Gli autori propongono un modello di formazione dell'immagine modificato:
$$I = (J + N)t + A(1 - t)$$
dove $I$ è l'immagine degradata, $J$ è l'immagine restaurata, $N$ è una componente residua moltiplicativa (per assorbire gli artefatti rimanenti), $t$ è la mappa di trasmissione e $A$ è la luce di backscatter globale.

Il funzionale di energia da minimizzare include:

• Data Fidelity (Fedeltà ai dati): Basata sul modello di decomposizione sopra citato.
• Nonlocal Regularization ($\alpha$-term) (Regolarizzazione non locale): Un vincolo di gradiente non locale per preservare i dettagli fini e la texture.
• Total Variation ($\beta$-term) (Variazione totale): Applicata alla mappa di trasmissione per garantire la regolarità spaziale preservando al contempo le discontinuità di profondità.
• Residual Constraint ($\lambda$-term) (Vincolo residuo): Limita l'assorbimento delle componenti residue.
• Nonlocal Gradient Fidelity ($\mu$-term) (Fedeltà del gradiente non locale): Penalizza la distanza tra il gradiente dell'immagine restaurata e un campo vettoriale specifico $V$ (derivato dal gradiente amplificato dell'input degradato). Questo termine è progettato per rafforzare le informazioni strutturali nascoste nelle immagini subacquee a basso contrasto.
• Transmission Fidelity ($\rho$-term) (Fedeltà della trasmissione): Assicura che la mappa di trasmissione stimata rimanga vicina a una stima iniziale (derivata dal Dark Channel Prior).

Gli autori forniscono una prova teorica che stabilisce l'esistenza di un minimizzatore per questo problema variazionale all'interno di uno spazio ammissibile definito.

2. Deep Unfolding Network (Rete di deep unfolding)

L'algoritmo di ottimizzazione (una strategia primal-dual basata su Chambolle–Pock) viene "unfolded" (srotolato) in una rete neurale profonda con $K$ stadi. In questo framework, gli aggiornamenti iterativi delle variabili di ottimizzazione ($J, t, N$) corrispondono agli strati della rete.

• Prior Apprendibili: Gli operatori prossimali associati al prior dell'immagine ($J$) e alla mappa di trasmissione ($t$) sono sostituiti da reti neurali apprendibili.
• Integrazione di Mamba: Per catturare efficientemente le dipendenze a lungo raggio e le auto-similarità senza l'alto costo computazionale dell'auto-attenzione, gli autori integrano i layer Mamba (State Space Models).
  • MambaResNet: Sostituisce l'operatore prossimale per l'immagine $J$. Utilizza un design multi-branch che elabora l'immagine di input, una versione con bilanciamento del bianco migliorato e una versione con equalizzazione dell'istogramma. Impiega blocchi 2D Selective Scan (SS2D) per modellare le interazioni non locali.
  • ProxNet: Una rete residua leggera per la mappa di trasmissione $t$.
  • MambaNet: Un modulo che modella il termine di fedeltà del gradiente non locale, operando nel dominio del gradiente.
• Inizializzazione: La rete viene inizializzata con l'immagine degradata per $J$, con zero per $N$ e con la stima del Dark Channel Prior per $t$.

3. Proximal Trajectory Loss (Perdita di traiettoria prossimale)

Per garantire che la rete di unfolding aderisca alla logica di ottimizzazione sottostante, gli autori introducono una perdita di traiettoria prossimale. Inve al di là della supervisione del solo output finale, questa perdita impone la coerenza tra gli stadi intermedi della rete e le iterazioni di un "ideale" regolarizzatore di restauro (derivato da un problema di ottimizzazione semplificato che coinvolge la ground truth). Ciò impedisce agli stadi intermedi di deviare dal percorso di ottimizzazione previsto.

Key Contributions (Contributi chiave)

1. Formulazione Variazionale: Un nuovo funzionale di energia per il miglioramento subacqueo che incorpora una componente residua moltiplicativa e un vincolo di tipo gradiente non locale per gestire gli artefatti e preservare i bordi.
2. Analisi Teorica: Una prova dell'esistenza di minimizzatori per il proposto problema variazionale non convesso.
3. Mamba-Based Unfolding: Una rete di deep unfolding che sostituisce gli operatori prossimali con moduli basati su Mamba, consentendo la modellazione efficiente delle dipendenze a lungo raggio nelle scene subacquee.
4. Proximal Trajectory Loss: Una nuova funzione di perdita progettata per imporre la coerenza tra gli stadi di unfolding della rete e le iterazioni dell'algoritmo di ottimizzazione ideale, migliorando interpretabilità e prestazioni.

Experimental Results (Risultati sperimentali)

Il metodo è stato valutato sui dataset UIEB e SUIM-E, confrontandolo con i metodi classici (es. CBAF, BayesianRetinex), i metodi di puro deep learning (es. UColor, Phaseformer, Mamba-UIE) e altri metodi di unfolding (UnfoldNet).

• Prestazioni Quantitative: Il metodo di unfolding proposto ha ottenuto risultati allo stato dell'arte in tutte le metriche (PSNR, SSIM, LPIPS) su entrambi i dataset. Ad esempio, su UIEB, ha raggiunto un PSNR di 25.16 e un SSIM di 0.9140, superando significativamente il secondo miglior metodo (Mamba-UIE).
• Qualità Visiva: I confronti visivi hanno mostrato che il metodo proposto recupera meglio l'accuratezza del colore (riducendo le tonalità rossastre/verdastre comuni in altri metodi) e preserva i dettagli fini (es. scaglie di pesce, texture del fondale) rispetto ai competitor.
• Studi di Ablazione:
  • Dehazing vs. Retinex: La formulazione basata sul de-hazing ha superato un'alternativa basata su Retinex.
  • Gradiente di Vincolo: Rimuovere il vincolo di gradiente non locale ha comportato un minor contrasto dei bordi e una minore qualità percettiva.
  • Traiettoria Loss: La supervisione degli stadi intermedi tramite la perdita di traiettoria prossimale ha migliorato significativamente le metriche finali rispetto all'uso del solo MSE sull'output finale.
  • Efficienza: I confronti con la Multi-Head Attention (MHA) hanno mostrato che l'architettura basata su Mamba ha raggiunto un numero inferiore di FLOPs e un minor uso di memoria GPU, specialmente ad alte risoluzioni, mantenendo al contempo un throughput più elevato.

Significance and Claims (Significato e affermazioni)

Il documento afferma che il framework proposto riesce a colmare il divario tra l'ottimizzazione basata su modelli e il deep learning. Integrando un modello variazionale motivato fisicamente con l'efficienza degli State Space Models (Mamba), gli autori ottengono:

• Restauro Migliorato: Qualità visiva e prestazioni quantitative superiori rispetto ai recenti metodi allo stato dell'arte.
• Interpretabilità: La struttura di unfolding preserva la logica del problema di ottimizzazione sottostante, potenziata dalla garanzia teorica dell'esistenza della soluzione.
• Efficienza: L'uso di layer Mamba consente di catturare efficientemente il contesto globale (modellazione non locale) con complessità lineare, affrontando i colli di bottiglia computazionali degli approcci basati su Transformer nell'imaging subacqueo.

Gli autori notano che, sebbene i risultati siano promettenti, il lavoro futuro potrebbe concentrarsi sul miglioramento della stima iniziale della trasmissione (attualmente basata sul Dark Channel Prior) ed estendere il modello al miglioramento video e alla coerenza temporale.