You Point, I Learn: Online Adaptation of Interactive Segmentation Models for Handling Distribution Shifts in Medical Imaging

Probleemstelling

Medische beeldsegmentatie is essentieel voor diagnose en behandeling, maar diepe leermodellen kampen vaak met distributieveranderingen (distribution shifts). Dit betekent dat de verdeling van de testdata (bijv. beelden van een andere scanner, een andere kliniek of een andere ziekte) afwijkt van de trainingsdata, wat leidt tot een sterke prestatiedaling.

Hoewel interactieve segmentatiemodellen (waarbij gebruikers klikken gebruiken om voorspellingen te verbeteren) al bestaan, hebben ze tot nu toe geen ingebouwde mechanismen om hun eigen parameters aan te passen op basis van deze gebruikerscorrecties. Bestaande online aanpassingsmethodes zijn vaak beperkt: ze focussen te nauw op de exacte geklikte pixels (via kruis-entropie of focal loss), negeren de omliggende gebieden, en vereisen complexe regularisatoren om overfitting te voorkomen.

Methodologie: OAIMS

De auteurs stellen OAIMS (Online Adaptation for Interactive Medical-image Segmentation) voor. Dit is een raamwerk dat een model continu aanpast aan nieuwe data-distributies tijdens het inferentieproces, gebruikmakend van twee complementaire adaptatiestadia en een nieuwe verliesfunctie.

1. Kerncomponent: Click-Centered Gaussian (CCG) Loss

Een centrale innovatie is de CCG Loss. In tegenstelling tot traditionele methodes die alleen straffen toepassen op de exacte pixel van een klik, straft deze loss de voorspellingen in het omringende gebied van de klik af, gewogen door een Gaussische kernel.

• Doel: Het model dwingen om sterk te reageren op gebruikersinvoer in de directe omgeving van de klik.
• Mechanisme: De loss wordt alleen toegepast op pixels die dezelfde klasse hebben als de klik (bijv. alleen voorgrond-pixels bij een voorgrond-klik). Dit voorkomt dat het model verwart wordt door pixels van een andere klasse in de buurt, wat overfitting en slechte generalisatie bij distributieveranderingen zou kunnen veroorzaken.

2. Twee Adaptatiestadia

Het raamwerk combineert twee soorten online learning:

• Mid-Interaction Adaptation (MI):

  • Wanneer: Na elke individuele klik van de gebruiker.
  • Proces: Het model voert een voorspelling uit, de gebruiker corrigeert met een klik, en het model past zijn gewichten direct aan voordat de volgende iteratie begint.
  • Pseudo-ground-truth: De voorspelling na de klik (maar voor de update) wordt gebruikt als tijdelijke "pseudo-ground-truth" om de update te sturen, in combinatie met de CCG Loss.
  • Voordeel: Verbeterd de kwaliteit van de huidige afbeelding en de volgende afbeeldingen in de sequentie.

• Post-Interaction Adaptation (PI):

  • Wanneer: Nadat de gebruiker klaar is met het corrigeren van een volledige afbeelding (na $T$ klikken).
  • Aanneming: De definitieve, door de gebruiker gecorrigeerde masker ($P_{final}$) is van voldoende kwaliteit om als pseudo-ground-truth te dienen voor het trainen van het model op de nieuwe distributie.
  • Proces (Twee fasen):
    1. Fase 1: Fijnafstelling (fine-tuning) met de initiële lokalisatie-klik als input, gebruikmakend van Dice-Focal loss tegen $P_{final}$.
    2. Fase 2: Fijnafstelling met kunstmatige correctie-klikken. Omdat de originele klikken al gebruikt zijn, genereert het systeem nieuwe klikken op locaties waar de output van Fase 1 afwijkt van $P_{final}$ (valse positieven/negatieven). Deze nieuwe klikken worden gebruikt met de CCG Loss en Dice-Focal loss om het model verder te optimaliseren.

Belangrijkste Bijdragen

1. OAIMS Framework: Een praktisch en effectief raamwerk voor online adaptatie van interactieve segmentatiemodellen zonder complexe regularisatoren.
2. CCG Loss: Een nieuwe verliesfunctie die de focus legt op de omgeving van gebruikersklikken, wat cruciaal is voor het leren van nieuwe distributies.
3. Gebruik van Pseudo-Ground-Truth: Het bewijs dat door de gebruiker gecorrigeerde maskers effectief kunnen dienen als trainingsdata voor online adaptatie, zelfs zonder menselijke annotatie van de volledige afbeelding.
4. Combinatie van MI en PI: Het tonen aan dat het combineren van adaptatie tijdens en na de interactie superieure resultaten oplevert dan alleen adaptatie na afloop.

Resultaten

De methode is getest op 5 fundus-databases en 4 hersen-MRI-databases met diverse distributieveranderingen (verschillende scanners, modaliteiten en pathologieën).

• Prestatieverbetering: OAIMS presteert consistent beter dan state-of-the-art methoden (zoals IA+SA en TSCA) en niet-adaptieve basismodellen.
  • Op hersen-MRI-data (bijv. BRATS) werden verbeteringen in de Dice-score van meer dan 10% waargenomen ten opzichte van eerdere methoden.
  • Zelfs met slechts Post-Interaction adaptatie (zonder MI) presteert het model beter dan methoden die meer backpropagaties vereisen.
• Robuustheid:
  • Weinig klikken: De methode werkt effectief zelfs met een beperkt budget aan klikken (3 of 5 per afbeelding).
  • Foutieve invoer: Het model is robuust tegen foutieve gebruikersklikken (tot 40% verkeerde klassen) en slechte initiële pseudo-ground-truth (bij extreme domeinverschillen).
  • Overfitting: In tegenstelling tot concurrenten (zoals TSCA) vertoont OAIMS geen tekenen van overfitting, zelfs niet bij een extreem hoog aantal klikken (50 per afbeelding).
• Efficiëntie: De computertijd voor adaptatie is verwaarloosbaar (0,05s - 0,09s op een GPU, <0,5s op een CPU), wat het geschikt maakt voor real-time klinische workflows.

Significantie

Dit paper biedt een oplossing voor een van de grootste uitdagingen in medische AI: het aanpassen van modellen aan nieuwe, onbekende data-distributies zonder opnieuw te hoeven trainen. Door gebruikersinvoer (klikken) direct om te zetten in modelverbetering via een slimme online leerstrategie, maakt OAIMs modellen adaptief en toekomstbestendig.

De methode elimineert de noodzaak om meerdere modellen te beheren voor verschillende ziektes of scanners. In plaats daarvan kan één model continu leren van de klinische praktijk, wat de drempel voor de implementatie van interactieve AI in de zorg verlaagt en de nauwkeurigheid van diagnose en behandeling direct verbetert. De open-source code en de brede validatie op verschillende medische modaliteiten onderstrepen de praktische toepasbaarheid van de aanpak.