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

Probleemstelling

Onderwaterbeeldvorming is cruciaal voor oceaantechniek en marien onderzoek, maar de vastgelegde gegevens lijden vaak onder slechte zichtbaarheid, niet-uniforme verlichting, kleurverzwakking en ruis. Deze degradaties hinderen daaropvolgende computervisie-taken aanzienlijk.

• Klassieke methoden: Traditionele benaderingen zijn verdeeld in modelvrije technieken (het manipuleren van pixelintensiteiten zonder fysieke aannames) en fysische modelgebaseerde methoden (het oplossen van inverse problemen op basis van het onderwater beeldvormingsmodel, doorgaans het Jaffe-McGlamery model). Hoewel variatiemethoden interpreteerbaarheid bieden, hangt hun effectiviteit sterk af van de handmatige selectie van geschikte priors.
• Deep Learning-methoden: Datagestuurde deep learning-benaderingen leren priors direct van data, maar vereisen vaak complexe architecturen en grote datasets. Puur datagestuurde modellen kunnen een gebrek aan interpreteerbaarheid hebben, terwijl standaard deep learning-architecturen (zoals Transformers) hoge computationele kosten kunnen veroorzaken bij het vastleggen van langetermijnafhankelijkheden (long-range dependencies).
• De kloof: Er is een behoefte om de kloof te overbruggen tussen de interpreteerbaarheid van modelgebaseerde optimalisatie en de flexibiliteit van deep learning, specif으로 gericht op het efficiënt aanpakken van de uitdaging om langetermijn-zelfgelijkenissen in onderwater-scènes vast te leggen.

Methodologie

De auteurs stellen een model-gebaseerd deep unfolding netwerk voor dat variatiemodellering integreert in een leerbare architectuur. De aanpak bestaat uit drie hoofdonderdelen: een variatie-baseline, een deep unfolding architectuur en een gespecialiseerde trainingsstrategie.

1. Variatie-baseline

De methode is geworteld in een variatieformulering gebaseerd op een dehazing-decompositie. De auteurs stellen een gewijzigd beeldvormingsmodel voor:
$$I = (J + N)t + A(1 - t)$$
waarbij $I$ het gedegradeerde beeld is, $J$ het herstelde beeld, $N$ een multiplicatieve residuele component (om resterende artefacten op te vangen), $t$ de transmissiekaart is, en $A$ het globale backscattered licht is.

De energiefunctionaal die geminimaliseerd moet worden, bevat:

• Data Fidelity: Gebaseerd op bovenstaande decompositie.
• Nonlocal Regularization ($\alpha$-term): Een niet-lokale gradiëntbeperking om fijne details en texturen te behouden.
• Total Variation ($\beta$-term): Toegepast op de transmissiekaart om ruimtelijke gladheid te garanderen terwijl diepte-discontinuïteiten behouden blijven.
• Residual Constraint ($\lambda$-term): Beperkt de absorptie van residuele componenten.
• Nonlocal Gradient Fidelity ($\mu$-term): Bestraft de afstand tussen de gradiënt van het herstelde beeld en een specifiek vectorveld $V$ (afgeleid van de versterkte gradiënt van de gedegradeerde input). Deze term is ontworpen om de structurele informatie die verborgen zit in laag-contrastrijke onderwaterbeelden te versterken.
• Transmission Fidelity ($\rho$-term): Zorgt ervoor dat de geschatte transmissiekaart dicht bij een initiële schatting blijft (afgeleid van de Dark Channel Prior).

De auteurs leveren een theoretisch bewijs dat de existentie van een minimizer voor dit variatieprobleem vaststelt binnen een gedefinieerde toegestane ruimte.

2. Deep Unfolding Netwerk

Het optimalisatiealgoritme (een primal-dual strategie gebaseerd op Chambolle–Pock) wordt "unfolded" naar een diep neuraal netwerk met $K$ stadia. In dit kader komen de iteratieve updates van de optimalisatievariabelen ($J, t, N$) overeen met netwerklagen.

• Leerbare Priors: De proximale operatoren geassocieerd met de beeldprior ($J$) en de transmissiekaart ($t$) worden vervangen door leerbare neurale netwerken.
• Mamba Integratie: Om efficiënt langetermijn-afhankelijkheden en zelfgelijkenissen vast te leggen zonder de hoge computationele kosten van self-attention, integreren de auteurs Mamba-lagen (State Space Models).
  • MambaResNet: Vervangt de proximale operator voor het beeld $J$. Het gebruikt een multi-branch design dat de inputafbeelding, een witbalans-verbeterde versie en een histogram-geëgaliseerde versie verwerkt. Het maakt gebruik van 2D Selective Scan (SS2D) blokken om niet-lokale interacties te modelleren.
  • ProxNet: Een lichtgewicht residueel netwerk voor de transmissiekaart $t$.
  • MambaNet: Een module die de niet-lokale gradiënt-fidelity term modelleert, werkend in het gradiëntdomein.
• Initialisatie: Het netwerk wordt geïnitialiseerd met het gedegradeerde beeld voor $J$, nul voor $N$, en de Dark Channel Prior schatting voor $t$.

3. Proximal Trajectory Loss

Om ervoor te zorgen dat het unfolding-netwerk voldoet aan de onderliggende optimalisatielogica, introduceren de auteurs een proximal trajectory loss. In plaats van alleen de uiteindelijke output te superviseren, dwingt deze loss consistentie af tussen de tussenliggende stadia van het netwerk en de iteraties van een "ideale" restauratie-regularisator (afgeleid van een vereenvoudigd optimalisatieprobleem met de grondwaarheid/ground truth). Dit voorkomt dat tussenliggende stadia afwijken van het verwachte optimalisatiepad.

Belangrijkste Bijdragen

1. Variatie Formulering: Een nieuwe energiefunctionaal voor onderwaterverbetering die een multiplicatieve residuele component en een niet-lokale gradiënt-type beperking incorporeert om artefacten aan te pakken en randen te behouden.
2. Theoretische Analyse: Een bewijs van de existentie van minimizers voor het voorgestelde niet-convexe variatieprobleem.
3. Mamba-gebaseerde Unfolding: Een deep unfolding netwerk dat proximale operatoren vervangt door Mamba-gebaseerde modules, wat efficiënte modellering van langetermijn-afhankelijkheden in onderwater-scènes mogelijk maakt.
4. Proximal Trajectory Loss: Een nieuwe loss-functie ontworpen om consistentie af te dwingen tussen de unfolding-stadia van het netwerk en de iteraties van het ideale optimalisatiealgoritme, wat zowel de interpreteerbaarheid als de prestaties verbetert.

Experimentele Resultaten

De methode is geëvalueerd op de UIEB en SUIM-E datasets, waarbij het is vergeleken met klassieke methoden (bijv. CBAF, BayesianRetinex), pure deep learning-methoden (bijv. UColor, Phaseformer, Mamba-UIE) en andere unfolding-methoden (UnfoldNet).

• Kwantitatieve Prestaties: De voorgestelde unfolding-methode behaalde state-of-the-art resultaten over alle metrieken (PSNR, SSIM, LPIPS) op beide datasets. Zo behaalde het op UIEB een PSNR van 25.16 en een SSIM van 0.9140, waarmee het de op één na beste methode (Mamba-UIE) met een aanzienlijke marge versloeg.
• Visuele Kwaliteit: Visuele vergelijkingen toonden aan dat de voorgestelde methode de kleuraccuratesse beter herstelde (het verminderen van de rode/groene tinten die bij andere methoden gebruikelijk zijn) en fijne details (bijv. visschubben, zeebodemtexturen) beter conserveerde vergeleken met concurrenten.
• Ablatie-studies:
  • Dehazing vs. Retinex: De dehazing-gebaseerde formulering presteerde beter dan een Retinex-gebaseerde variant.
  • Gradiëntbeperking: Het verwijderen van de niet-lokale gradiëntbeperking resulteerde in lagere randcontrasten en een lagere perceptuele kwaliteit.
  • Trajectory Loss: Het superviseren van tussenliggende stadia via de proximal trajectory loss verbeterde de uiteindelijke metrieken aanzienlijk vergeleken met het enkel gebruiken van MSE op de uiteindelijke output.
  • Efficiëntie: Vergelijkingen met Multi-Head Attention (MHA) toonden aan dat de Mamba-gebaseerde architectuur een lager aantal FLOPs en minder GPU-geheugengebruik bereikte, vooral bij hogere resoluties, terwijl het een hogere throughput behield.

Betekenis en Claims

Het artikel claimt dat het voorgestelde framework er succesvol in is geslaagd de kloof tussen modelgebaseerde optimalisatie en deep learning te overbruggen. Door een fysisch gemotiveerd variatiemodel te integreren met de efficiëntie van State Space Models (Mamba), bereiken de auteurs:

• Verbeterde Restauratie: Superieure visuele kwaliteit en kwantitatieve prestaties vergeleken met recente state-of-the-art methoden.
• Interpreteerbaarheid: De unfolding-structuur behoudt de logica van het onderliggende optimalisatieprobleem, versterkt door het theoretische bewijs van de existentie van de oplossing.
• Efficiëntie: Het gebruik van Mamba-lagen maakt het mogelijk om een globale context (niet-lokale modellering) efficiënt vast te leggen met lineaire complexiteit, wat de computationele knelpunten van Transformer-gebaseerde benaderingen in onderwaterbeeldvorming aanpakt.

De auteurs merken op dat hoewel de resultaten veelbelovend zijn, toekomstig werk zich kan richten op het verbeteren van de initiële transmissieschatting (momenteel afhankelijk van de Dark Channel Prior) en het uitbreiden van het model naar video-verbetering en temporele consistentie.