Prompt-SID: Learning Structural Representation Prompt via Latent Diffusion for Single-Image Denoising

Dit paper introduceert Prompt-SID, een zelftoezichtend framework voor het verwijderen van ruis uit een enkele afbeelding dat structurele details behoudt door middel van een prompt-leringsbenadering met latent diffusion en een transformer-architectuur.

De kern

Stel je voor dat je een prachtige foto hebt gemaakt, maar deze is bevroren met sneeuwvlokken (ruis). Je wilt de sneeuw verwijderen om de foto weer scherp te maken, maar je hebt geen "zuivere" foto om als voorbeeld te gebruiken. Dit is het grote probleem bij het verwijderen van ruis uit foto's zonder dat je een schone versie hebt.

De meeste bestaande methoden doen alsof ze de foto in kleine stukjes snijden, waarbij ze sommige pixels (de kleintjes waaruit een foto bestaat) weglaten om te raden wat eronder zit. Het probleem is dat hierdoor belangrijke details, zoals de randen van een gebouw of de textuur van een bloem, verloren gaan. Het is alsof je een puzzel probeert op te lossen, maar je gooit de helft van de stukjes weg en hoopt dat je het plaatje toch kunt reconstrueren.

Prompt-SID is een nieuwe, slimme manier om dit probleem op te lossen. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Gids" in plaats van het "Gokken"

Stel je voor dat je een schilderij moet restaureren, maar je hebt alleen een versleten, vervaagde kopie. In plaats van blindelings te raden, maakt Prompt-SID eerst een sneltekening (een "prompt") van de originele structuur.

• De Analogie: Denk aan een architect die een blauwdruk maakt van een oud huis voordat hij begint met de renovatie. Hij kijkt naar de ruwe, beschadigde muren en maakt een schets van hoe de muren moeten zijn. Deze schets gebruikt hij als een leidraad (de "prompt") om de echte muren weer op te bouwen.
• In de tech-taal: Het systeem gebruikt een speciale "tijdmachine" (een diffusiemodel) om een compacte, gestructureerde versie van de foto te genereren. Dit is geen willekeurige gok, maar een slimme samenvatting van hoe de foto eruit zou moeten zien, zonder dat er pixels verloren gaan.

2. De "Tolken" die de instructies vertalen

Zodra de "sneltekening" (de prompt) klaar is, moet deze worden vertaald naar de echte foto.

• De Analogie: Stel je voor dat de sneltekening in een vreemde taal is geschreven. Je hebt een tolk nodig die deze instructies vertaalt naar de taal van de verf en de penseelstreken. Prompt-SID heeft een speciale "Tolk-module" (de Structural Attention Module) die precies weet welke details belangrijk zijn (zoals de rand van een oog of een blad) en welke details ruis zijn (de sneeuwvlokken).
• Hoe het werkt: Deze module kijkt naar de sneltekening en zegt tegen het systeem: "Kijk hier, dit is een scherpe rand, houd dit vast!" en "Hier is alleen ruis, verwijder dit." Zo wordt de foto scherp gemaakt zonder dat de fijne details vervagen.

3. De "Repetitie" op verschillende schalen

Een ander groot probleem is dat de training vaak gebeurt op kleine, verkleinde foto's, maar de echte foto's zijn groot. Het is alsof je een orkest traint op een fluitje, maar dan een symfonie moet spelen op een groot podium. De muziek klinkt anders.

• De Analogie: Prompt-SID gebruikt een slimme repetitiemethode. Het oefent eerst op de kleine versie van de foto, maar daarna "repeteert" het ook op de grote, originele versie. Het systeem kijkt naar het resultaat op de grote foto en zegt: "Ah, hier klopt iets niet, ik moet mijn strategie aanpassen."
• Het resultaat: Hierdoor leert het systeem niet alleen om kleine stukjes te repareren, maar ook om de hele grote foto perfect te maken, zonder dat de details "verdwijnen" door de schaalvergroting.

Waarom is dit zo speciaal?

Vroeger moesten systemen vaak kiezen tussen "goed ruis verwijderen" en "alle details behouden". Vaak werd de foto dan wel schoon, maar vervaagd (zoals een wazige foto).

Prompt-SID slaat twee vliegen in één klap:

1. Het mist geen details: Omdat het de hele foto "leest" via de sneltekening, worden geen pixels weggegooid.
2. Het is slim: Het gebruikt de kracht van moderne AI (diffusiemodellen, zoals die ook worden gebruikt om nieuwe kunst te maken) om de structuur van de foto te begrijpen, in plaats van alleen te rekenen.

Kortom: Prompt-SID is als een super-slimme restaurateur die niet alleen de vuile plekken verwijdert, maar ook een mentale blauwdruk heeft van hoe het schilderij eruit moet zien. Hierdoor komt de foto er niet alleen schoon uit, maar ook haarscherp en met alle fijne details intact, zelfs als er geen perfecte voorbeeldfoto beschikbaar is.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Prompt-SID: Learning Structural Representation Prompt via Latent Diffusion for Single-Image Denoising" in het Nederlands.

Probleemstelling

Bestaande methoden voor beelddenoising (ruisverwijdering) kampen met aanzienlijke beperkingen:

• Supervised methods: Vereisen grote datasets met gepaarde ruisvrije en ruizige afbeeldingen, wat duur en tijdrovend is om te verzamelen.
• Self-supervised/Unsupervised methods: Bestaande benaderingen zoals Blind-Spot Networks (bijv. Noise2Void) of sub-image sampling (bijv. Noise2Noise) leiden vaak tot verlies van pixelinformatie en beschadiging van gedetailleerde structurele informatie.
  • Bij blind-spot netwerken is het centrale pixel van het convolutiekern onzichtbaar.
  • Bij downsampling van sub-afbeeldingen gaan pixels verloren en treedt er semantische degradatie op, wat de kwaliteit van het herstel beperkt.

Het doel van dit onderzoek is een self-supervised framework te ontwikkelen dat deze structurele schade en semantische degradatie overwint zonder afhankelijk te zijn van gelabelde data.

Methodologie: Prompt-SID

Het voorgestelde framework, Prompt-SID, is een self-supervised systeem dat gebaseerd is op prompt-learning en latent diffusion. De kernarchitectuur bestaat uit drie hoofdcomponenten:

1. Ruimtelijke Redundantie Sampling Strategie (Spatial Redundancy Sampling)
Om het verlies van pixels te minimaliseren die gebruikelijk is bij downsampling, gebruikt het systeem een strategie waarbij de originele ruizige afbeelding $x$ wordt opgedeeld in blokken van 4 pixels. Uit elk blok worden willekeurig 3 aangrenzende pixels geselecteerd om drie sub-afbeeldingen ($m_1, m_2, m_3$) te genereren die elk 1/4e van de oorspronkelijke grootte hebben.

• $m_1(x)$ dient als invoer voor het denoising-netwerk.
• $m_2(x)$ en $m_3(x)$ fungeren als doelen voor supervisie (reconstructieverlies).

2. Structuur Representatie Generatie Diffusie (RG-Diff)
Dit is een innovatieve module die een latent diffusion model gebruikt om structurele prompts te genereren.

• Encodering: Een Pixel Structure Encoder (PSE) comprimeert zowel de gedownsamplede afbeelding ($m_1$) als de originele afbeelding ($x$) naar een impliciete vectorruimte, resulterend in $c_{sub}$ en $c_{org}(0)$.
• Diffusieproces: Een forward diffusion proces voegt ruis toe aan $c_{org}(0)$.
• Reverse Diffusie met Conditie: Tijdens het denoisingproces (reverse diffusion) wordt $c_{sub}$ gebruikt als conditionele input. Het model leert de beschadigde structuurrepresentatie van de gedownsamplede afbeelding te gebruiken om de intacte structuurrepresentatie van de originele afbeelding ($\hat{c}_{org}(0)$) te reconstrueren.
• Doel: Dit proces genereert een "prompt" die de semantische en structurele informatie van de originele schaal vastlegt, zonder directe identity mapping (het model leert niet simpelweg de ruis te kopiëren).

3. Structuur Attention Module (SAM) en SPIformer
Het denoising-netwerk is gebaseerd op een Vision Transformer (SPIformer).

• De gegenereerde structuurrepresentatie $\hat{c}_{org}(0)$ wordt geïntegreerd in het netwerk via de Structural Attention Module (SAM).
• SAM berekent kanaal-attentie-weights en combineert deze met de structuur-prompt om de feature maps te sturen. Dit zorgt ervoor dat het netwerk focus legt op kanalen met rijke structurele details en ruis onderdrukt.

4. Schaal Replay Mechanisme (Scale Replay)
Om de kloof (domain gap) tussen gedownsamplede trainingdata en originele resolutie afbeeldingen tijdens inferentie te overbruggen:

• Tijdens elke training-iteratie wordt er, na het verwerken van de gedownsamplede afbeeldingen, een extra inferentiestap uitgevoerd op de originele schaal afbeelding $x$.
• De output wordt gedownsampled en vergeleken met de doelen, maar zonder backpropagatie van de gradiënten van deze stap.
• Dit fungeert als een regularisatie die het model dwingt om generaliseerbaar te zijn voor originele schalen, zonder identity mapping te veroorzaken.

Belangrijkste Bijdragen

1. Prompt-Learning voor Self-Supervised Denoising: Het introduceren van een pipeline die structurele representaties uit originele afbeeldingen haalt om het herstelproces van gedownsamplede inputs te sturen, wat semantische degradatie voorkomt.
2. Overbrugging van Schaalverschillen: Een speciale tak voor het verwerken van originele resolutie (via Scale Replay) die indirect bijdraagt aan de optimalisatie en identity mapping voorkomt.
3. Toepassing van Diffusiemodellen: Pionierswerk door latent diffusion te gebruiken voor het genereren van structurele prompts in de self-supervised setting, in plaats van alleen voor generatie.
4. Superieure Prestaties: Het bereiken van state-of-the-art resultaten op synthetische, real-world en fluorescentie-imaging datasets met een relatief klein aantal parameters (6M).

Resultaten

Het paper presenteert uitgebreide experimenten op drie soorten datasets:

• Synthetische Data (Gaussian & Poisson ruis): Op datasets zoals Kodak, BSD300 en Set14 overtreft Prompt-SID bestaande self-supervised methoden (zoals N2V, B2U, NBR2NBR) consistent met 0.21 tot 0.34 dB in PSNR. Het presteert zelfs beter dan sommige supervised baselines.
• Real-world Data (SIDD): Op de SIDD benchmark (raw-RGB) behaalt het een PSNR van 51.02 dB, wat een verbetering is van 0.55 dB ten opzichte van de NBR2NBR architectuur en 0.23 dB ten opzichte van B2U. Het behoudt beter randdetails en vermijdt kleuronevenwichtigheden.
• Fluorescentie Imaging: Op 3D neuronale datasets presteert het beter dan andere self-supervised methoden en komt het in de buurt van supervised methoden, zelfs bij hoge scansnelheden (30 Hz).

Ablatie Studies bevestigen dat elk onderdeel essentieel is:

• Zonder RG-Diff (structuur prompt) neemt de kwaliteit van semantische details af.
• Zonder de schaal replay-mechanisme worden de beelden waziger door het ontbreken van hogere resolutie informatie tijdens training.

Betekenis en Conclusie

Prompt-SID markeert een belangrijke stap in de evolutie van self-supervised beelddenoising. Door de combinatie van latent diffusion voor het genereren van structurele prompts en een scale replay mechanisme voor domeinadaptatie, lost het fundamentele problemen op van eerdere methoden: het verlies van pixelinformatie en de beschadiging van structurele details.

De methode demonstreert dat diffusion modellen niet alleen nuttig zijn voor generatieve taken, maar ook effectief kunnen worden ingezet als conditionele mechanismen in restauratietaken. Dit leidt tot een robuust, lichtgewicht model dat zonder gelabelde data uitstekend presteert in diverse complexe scenario's, van natuurlijke foto's tot wetenschappelijke fluorescentie-imaging.