Diffusion Autoencoder for Unsupervised Artifact Restoration in Handheld Fundus Images

Probleemstelling

Handheld funduscamera's (op basis van smartphones) hebben de oogheelkundige diagnostiek toegankelijker, efficiënter en goedkoper gemaakt. Echter, beelden die met deze apparaten worden gemaakt, lijden vaak onder ernstige artefacten zoals flitsreflecties, over- of onderbelichting en bewegingsonscherpte. Deze kwaliteitsproblemen verstoren de beeldstructuur en belemmeren betrouwbare downstream-analyses, zoals ziektegradering (bijv. diabetische retinopathie) en vaatsegmentatie.

Bestaande generatieve modellen voor beeldherstel (zoals GAN's en eerdere diffusion-modellen) zijn vaak afhankelijk van gepaarde supervisie (paren van beschadigde en schone beelden) of gedefinieerde maskerstructuren. Dit maakt ze minder flexibel voor de ongestructureerde en variabele degradaties die typisch zijn voor handheld-opnames. Bovendien kunnen conventionele diffusion-processen zonder extra context de semantische informatie (zoals bloedvatstructuren) verliezen die gedeeltelijk door artefacten wordt bedekt.

Methodologie

De auteurs stellen een ongesuperviseerde Diffusion Autoencoder (DiffAE) voor die specifiek is ontworpen om artefacten te herstellen in handheld fundusbeelden, zonder dat er gepaarde trainingsdata nodig is.

1. Architectuur (DiffAE):

   • Het model combineert een context-encoder ($Enc_\phi$) met een denoising UNet ($f_\theta$).
   • Training: Het model wordt uitsluitend getraind op hoogwaardige, schone fundusbeelden van tafeltop-systemen (EyePACS dataset). De encoder leert een compacte latente vector ($z$) te genereren die globale retinale structuren en semantische prioriteiten vastlegt. De UNet leert de reverse diffusion-proces om het originele beeld te reconstrueren op basis van deze latentie.
   • Inferentie (Herstel): Tijdens het herstelproces van een handheld-beeld met artefacten, wordt een artefact-masker ($M$) gebruikt. De reconstructie wordt geleid door de contextuele representatie $z$ (afgeleid van het schone beeldconcept) en het masker. De formule voor de inpainting op tijdstap $t$ is:
     $$x^{inpaint}_t = M \odot x_t + (1 - M) \odot x^{noisy}_0$$
     Hierbij wordt de beschadigde regio vervangen door de generatie van het model, terwijl de gezonde delen behouden blijven.

2. Ongepaarde aanpak:

   • Het systeem vereist geen paren van "beschadigd vs. schoon" beelden voor training. Het leert de verdeling van gezonde retinale beelden en past deze kennis toe om beschadigde gebieden in handheld-beelden in te vullen (inpainting).

Belangrijkste Bijdragen

• Diffusion Auto-encoding Framework: Introductie van een framework dat een context-encoder integreert in het diffusion-proces om semantisch betekenisvolle representaties te leren voor het herstellen van artefacten.
• Ongesuperviseerde Inpainting: Ontwikkeling van een mechanisme dat alleen getraind wordt op hoogwaardige tafeltop-beelden, maar effectief inferentie uitvoert op artefact-bevattende handheld-beelden.
• Validatie op Downstream Taken: Het bewijzen dat de herstelde beelden niet alleen visueel beter zijn, maar ook leiden tot een hogere nauwkeurigheid in diagnostische taken (zoals ziektegradering en vaatsegmentatie) op onzichtbare datasets.

Resultaten

De prestaties zijn geëvalueerd op zowel synthetische datasets (met gepaarde ground truth) als een echte testset van handheld-beelden.

• Kwantitatieve Evaluatie:

  • Beeldkwaliteit: Het model behaalde een PSNR van 38,64 dB en een SSIM van 0,97, wat hoger is dan de beste baselines (DeepFillV2, EdgeConnect, etc.).
  • Vaatsegmentatie: De Dice-score voor vasculaire segmentatie bedroeg 0,891, wat aangeeft dat de bloedvatstructuren beter behouden blijven dan bij andere methoden.
  • Kwaliteitsbeoordeling: Op de echte handheld-testset werd 35,8% van de beelden geclassificeerd als "Good" (goed), vergeleken met 31,6% voor de beste concurrent.
  • Diagnostische Nauwkeurigheid: De nauwkeurigheid van de classificatie van diabetische retinopathie (DR) steeg naar 81,17% op de herstelde beelden, een significante verbetering ten opzichte van de niet-gerestaureerde baseline (77,11%).

• Kwalitatieve Evaluatie:

  • Visuele vergelijkingen tonen aan dat het model natuurlijker fundus-achtige patronen herstelt, terwijl andere modellen (zoals GAN's) vaak onnatuurlijke gele vlekken of vage structuren introduceren. Het model slaagt erin artefacten van verschillende maten en structuren (zoals flitsreflecties) te verwijderen zonder de onderliggende anatomie te verstoren.

• Ablatiestudie:

  • Het gebruik van een geïnterpoleerde context-vector (gemengd tussen de artefact-beeld latentie en een schone latentie) resulteerde in een nog hogere PSNR (41,46 dB), wat bevestigt dat de context-encoder cruciaal is voor de kwaliteit van het herstel.

Betekenis en Conclusie

Dit werk toont aan dat diffusion auto-encoders een krachtige oplossing bieden voor het ongesuperviseerd herstellen van medische beelden. Door de afhankelijkheid van gepaarde trainingsdata te elimineren en gebruik te maken van semantische context, kan het model robuust omgaan met de complexe en variabele artefacten die voorkomen in handheld fundusopnames.

De belangrijkste implicatie is dat deze technologie de bruikbaarheid van goedkope, mobiele oogscreeningsapparatuur aanzienlijk kan verbeteren. Door de beeldkwaliteit te herstellen, worden downstream AI-taken (zoals ziekte-detectie) betrouwbaarder, wat essentieel is voor het uitbreiden van oogheelkundige zorg naar gebieden met beperkte middelen. De auteurs suggereren dat toekomstig werk zich kan richten op het integreren van AI-gedreven lokalisatie van artefacten om het maskeringsproces verder te automatiseren.