FunduSegmenter: Leveraging the RETFound Foundation Model for Joint Optic Disc and Optic Cup Segmentation in Retinal Fundus Images

Probleemstelling

De nauwkeurige segmentatie van de anatomische structuren van het netvlies, met name de optische schijf (OD) en de optische beker (OC), is cruciaal voor het onderzoek naar biomarkers voor systemische ziekten (zoals diabetes en glaucoom). Hoewel overtoezicht op CNN-modellen (Convolutional Neural Networks) sterke prestaties heeft geleverd, lijden ze onder drie belangrijke beperkingen:

1. Afhankelijkheid van gelabelde data: Het animeren van medische beelden is tijdrovend en vereist experts, waardoor veel bestaande datasets ongelabeld blijven.
2. Slechte generalisatie: Supervised CNN-modellen presteren vaak slecht op "out-of-distribution" data (bijv. beelden van andere apparaten of klinische centra).
3. Beperkte data-augmentatie: Bestaande methoden gebruiken vaak eenvoudige transformaties of synthetische beelden, wat door clinici minder betrouwbaar wordt gevonden dan training op echte beelden.

De auteurs stellen dat Foundation Models (FMs), zoals RETFound (een model getraind op een enorme verzameling ongelabeld fundus- en OCT-materiaal via self-supervised learning), deze problemen kunnen oplossen door robuuste, algemene representaties te leren. Echter, er was tot nu toe geen succesvolle aanpassing van RETFound voor de specifieke taak van segmentatie van OD en OC.

Methodologie: FunduSegmenter

De auteurs stellen FunduSegmenter voor, een model dat de RETFound-encoder combineert met een aangepaste decoder en een reeks nieuwe modules om de overgang van classificatie naar segmentatie te faciliteren. De architectuur omvat de volgende kerncomponenten:

• Encoder (RETFound): De auteurs gebruiken de vooraf getrainde ViT-large (Vision Transformer) encoder van RETFound. De gewichten worden gefixeerd (frozen) om de rijke, algemene representaties te behouden.
• Decoder (Segmenter): In plaats van de standaard UPerNet of SETR decoders die vaak bij FMs worden gebruikt, kiezen ze voor de decoder van Segmenter, die beter presteert bij natuurlijke beelden.
• Nieuwe Modules:
  1. Pre-adapter: Een lichtgewicht module (geïnspireerd op MobileNetV2) die invoerbeelden van verschillende maten accepteert en deze verwerkt voordat ze de encoder bereiken. Dit is superieur aan de standaard "position interpolation" methode van RETFound.
  2. Post-adapter: Een module die de segmentatiemaps stapsgewijs (progressief) omhoog schaalt (upsampling) naar de originele beeldgrootte. Dit lost het probleem op van onduidelijke grenzen tussen OD en OC, wat bij directe bilineaire upsampling vaak tot fouten leidt.
  3. Skip Connections met CBAM: Feature maps uit de middelste lagen van de RETFound-encoder worden via skip-connections naar de decoder geleid, verrijkt met Convolutional Block Attention Modules (CBAM). Dit benut multi-schaal informatie, wat essentieel is voor het segmenteren van de kleinere OC.
  4. ViT Block Adapter: Een adapter geplaatst tussen de normalisatielaag (LN) en de Multi-Head Attention (MHA) blokken van de transformer. Dit stelt het model in staat om zich aan te passen aan de segmentatietaken zonder de algemene representatieve kracht van de gefixeerde encoder te verliezen.
• Verliesfunctie: Een combinatie van Dice Loss en Cross Entropy Loss.

Experimenteel Ontwerp

Het model werd geëvalueerd op een eigen dataset (GoDARTS) en vier publieke datasets (IDRiD, Drishti-GS, RIM-ONE-r3, REFUGE). Er werden drie soorten experimenten uitgevoerd:

1. Interne verificatie: Trainen en testen binnen dezelfde dataset.
2. Externe verificatie: Trainen op één dataset en testen op alle andere datasets.
3. Domein-generalisatie: Trainen op drie datasets en testen op de vierde (ongezien) dataset, vergelijkbaar met de benchmark DoFE.

De prestaties werden vergeleken met zes state-of-the-art baselines, waaronder nnU-Net, DUNet, TransUNet, en specifieke domein-generalisatiemodellen zoals DoFE, RAM-DSIR en TVConv.

Belangrijkste Resultaten

FunduSegmenter overtrof in bijna alle scenario's de bestaande baselines:

• Interne Verificatie: Het model bereikte een gemiddelde Dice Similarity Coefficient (DSC) van 90,51%. Dit is een significante verbetering ten opzichte van de baselines (bijv. nnU-Net: 82,91%, DUNet: 89,17%, TransUNet: 87,91%).
• Externe Verificatie: Het model presteerde consistent ongeveer 3% beter dan de beste baselines, zelfs wanneer het werd getraind op kleine datasets (zoals Drishti-GS met slechts 50 beelden).
• Domein-Generalisatie: Het model toonde sterke stabiliteit op onbekende data. Hoewel het iets achterbleef bij TVConv (een model dat specifiek is ontworpen voor generalisatie), presteerde het beter dan DoFE en RAM-DSIR.
• Stabiliteit: In tegenstelling tot nnU-Net, dat instabiel werd op onbekende data (bijv. slechts 17,96% DSC op GoDARTS), behield FunduSegmenter hoge prestaties (~92% voor OD) ongeacht de bron van de data.
• Ablatie Studies: De studies bevestigden dat elke nieuwe module (Pre-adapter, Post-adapter, Skip connections, ViT adapter) bijdroeg aan betere prestaties en een aanzienlijk snellere convergentie, vooral op kleine datasets waar baselines vaak vastliepen.

Bijdragen en Betekenis

De belangrijkste bijdragen van dit werk zijn:

1. Eerste toepassing van RETFound voor segmentatie: Dit is het eerste werk dat RETFound succesvol adapteert voor de segmentatie van anatomische structuren in fundusbeelden.
2. Superieure Generalisatie: Het bewijst dat foundation models, getraind op ongelabelde data, superieure generalisatie-eigenschappen hebben voor medische segmentatietaken, zelfs met beperkte gelabelde data.
3. Generieke Architectuur: De voorgestelde modules (Pre-adapter, Post-adapter, ViT adapter) zijn generiek en kunnen worden toegepast op andere foundation models met een ViT-encoder die via self-supervised learning zijn getraind.
4. Translatieve Relevantie: De stabiliteit van het model maakt het een betrouwbare basis voor geautomatiseerde klinische taken, variërend van het bepalen van retinale coördinaten tot het ontdekken van nieuwe biomarkers.

Conclusie: FunduSegmenter demonstreert dat het benutten van de "latent general representations" van foundation models een krachtige route is om de beperkingen van traditionele supervised deep learning in de oogheelkunde te overwinnen, met name wat betreft generalisatie over verschillende domeinen en datasets. De code en getrainde gewichten zijn openbaar beschikbaar.