Restore, Assess, Repeat: A Unified Framework for Iterative Image Restoration

Probleemstelling

Beeldherstel (Image Restoration - IR) richt zich op het herstellen van hoge-kwaliteit beelden uit inputs die zijn aangetast door factoren zoals wazigheid, slecht weer, of weinig licht. Bestaande methoden kampen echter met twee belangrijke beperkingen:

1. Beperkte generalisatie: Veel modellen zijn getraind op specifieke, bekende degradaties en falen bij onbekende of samengestelde (composiete) degradaties.
2. Inefficiëntie en loskoppeling:
   • Agent-based modellen gebruiken een agent die iteratief tools selecteert op basis van een Image Quality Assessment (IQA). Dit proces is echter traag, complex en vereist dat beelden tussen de IQA- en IR-stappen worden gedecodeerd en opnieuw gecodeerd, wat informatieverlies en latentie veroorzaakt.
   • All-in-one modellen zijn efficiënter maar presteren vaak slechter bij complexe scenario's omdat ze vaak afhankelijk zijn van vooraf gedefinieerde degradatielabels (bijv. via CLIP) die niet goed generaliseren naar vrije, onbekende degradaties.

Methodologie: RAR (Restore, Assess, Repeat)

De auteurs stellen RAR voor, een unificerend framework dat Image Quality Assessment (IQA) en Image Restoration (IR) integreert in één end-to-end trainbaar model dat volledig in de latent space (latente ruimte) opereert.

1. Kernarchitectuur

Het framework bestaat uit twee nauw met elkaar verbonden modules:

• Latent Quality Assessment (LQA): Gebaseerd op DepictQA (een Vision-Language Model), maar aangepast om te opereren in de latent space van het herstelmodel in plaats van in pixelruimte.
  • Image Adapter ($A_I$): Brengt de input van het herstelmodel (latent) in overeenstemming met de IQA-module.
  • Conditioning Adapter ($A_Q$): Projecteert de output van de IQA (vrije tekst over degradaties) direct naar de conditioning embeddings van het herstelmodel, zonder dat er tekst hoeft te worden gedecodeerd en opnieuw gecodeerd.
• Unified Image Restoration (UIR): Een generatief model (gebaseerd op Stable Diffusion 3.5 met Flow Matching) dat het beeld herstelt op basis van de feedback van de LQA.

2. Het Iteratieve Proces

In tegenstelling tot eerdere methoden die tekst gebruiken als tussenstap, werkt RAR volledig in de latent space:

1. Restore: Het UIR-model neemt een gedegradeerde latent ($z_{deg}$) en een conditionering ($Q_{deg}$) als input en genereert een verbeterde latent.
2. Assess: De LQA-module analyseert de nieuwe latent om de kwaliteit te verifiëren en nieuwe degradaties te identificeren.
3. Repeat: Als de kwaliteit niet optimaal is, wordt de conditionering dynamisch bijgewerkt op basis van de nieuwe LQA-voorspelling en wordt de cyclus herhaald.

3. Noise-Free Flow Matching

Een cruciaal technisch detail is het gebruik van Flow Matching zonder toegevoegde ruis (noise-free).

• Traditionele diffusion-modellen voegen bij elke stap ruis toe, wat de latent space "vervuilt" en de IQA-module onnauwkeurige metingen laat doen.
• RAR leert een directe mapping van de gedegradeerde verdeling ($\rho_{deg}$) naar de hoge-kwaliteit verdeling ($\rho_{hq}$). Dit zorgt ervoor dat de LQA-module de tussentijdse latents kan beoordelen zonder dat deze vervormd zijn door ruis, wat essentieel is voor stabiele iteratieve feedback.

4. Stopcriterium

Het proces stopt niet op een vast aantal iteraties, maar dynamisch. De LQA-module vergelijkt de kwaliteit van het beeld voor en na een reeks stappen. Als er geen significante verbetering wordt waargenomen, stopt het proces (STOP). Dit voorkomt over-restauratie.

Belangrijkste Bijdragen

1. Unificatie van IQA en IR: De eerste methode die een beschrijvende IQA (vrije tekst) en een generatief IR-model volledig integreert in één end-to-end trainbaar model dat in de latent space werkt.
2. Eliminatie van Latentie en Informatieverlies: Door de tekst-decoding en her-encoding te elimineren en direct te werken met embeddings in de latent space, wordt de latentie drastisch verlaagd en informatiebehoud gemaximaliseerd.
3. Dynamische Adaptatie: Het model past het herstelproces dynamisch aan op basis van iteratieve feedback, waardoor het effectief is bij onbekende en samengestelde degradaties.
4. State-of-the-Art Prestaties: Het framework bereikt nieuwe records op zowel enkele, samengestelde als onbekende degradaties.

Resultaten

De auteurs hebben RAR getest op diverse benchmarks (SIDD, GoPro, LOL, MiO100, UDC, EUVP) en vergeleken met state-of-the-art methoden zoals AutoDIR, AgenticIR, en PromptIR.

• Samengestelde Degradaties (Composite): RAR presteert aanzienlijk beter dan concurrenten. Op de meest uitdagende set (3 degradaties tegelijk) behaalde het een verbetering van +2.71 dB (PSNR) tot +72.4 punten (MANIQA) ten opzichte van de beste bestaande methoden.
• Onbekende Degradaties: Op de UDC-dataset (onbekende degradaties) overtreft RAR alle baselines, met name op perceptuele metrics zoals MUSIQ en CLIP-IQA.
• Efficiëntie: RAR is 11.27x sneller dan de state-of-the-art AgenticIR. Dit komt doordat het geen complexe agentische loop met meerdere tools en decodestappen vereist, maar een enkel, geoptimaliseerd model gebruikt.
• Kwaliteit: Visuele vergelijkingen tonen aan dat RAR rijkere details behoudt en realistischer beelden produceert dan methoden die vaak texturen verliezen of artefacten introduceren.

Betekenis en Impact

Dit paper markeert een verschuiving in het veld van beeldherstel van gescheiden, statische pipelines naar adaptieve, iteratieve systemen die in de latent space opereren.

• Het lost het probleem op van "onbekende" degradaties door een model te gebruiken dat niet afhankelijk is van vooraf gedefinieerde labels, maar de inhoud van het beeld zelf analyseert en beschrijft.
• Het bewijst dat het integreren van assessment en herstel in één model niet alleen de prestaties verbetert, maar ook de rekentijd drastisch verlaagt, waardoor het toepasbaar wordt voor real-time toepassingen.
• De methode biedt een nieuwe standaard voor hoe generatieve AI-modellen kunnen worden gebruikt voor complexe, iteratieve taakoplossing zonder de efficiëntie te verliezen.