Cell Instance Segmentation via Multi-Task Image-to-Image Schrödinger Bridge

Probleemstelling

Bestaande methoden voor cel-instantiesegmentatie (het onderscheiden van individuele cellen in dicht op elkaar gepakte microscopische beelden) vertrouwen doorgaans op een tweestapsproces:

1. Deterministische voorspelling: Een CNN genereert tussenliggende representaties (zoals waarschijnlijkheidskaarten, afstandkaarten of vectorvelden).
2. Post-processing: Algorithmische stappen (zoals watershed-transformaties, clustering of morfologische operaties) worden toegepast om deze representaties om te zetten in definitieve maskers.

Deze aanpak heeft twee belangrijke beperkingen:

• Afhankelijkheid van heuristieken: De post-processing is vaak dataset-specifiek en vereist handmatige hyperparameter-tuning, wat de robuustheid en generaliseerbaarheid beperkt.
• Gebrek aan globale structuur: Deterministische modellen leggen geen expliciete beperkingen op aan de globale structuur van de maskers. Dit leidt soms tot morfologisch onwaarschijnlijke output, zoals samengevoegde cellen of spookfragmenten, die later moeten worden gecorrigeerd.

Het paper stelt voor dat instantiesegmentatie beter kan worden gemodelleerd als een distributiegebaseerd probleem, waarbij het model leert om morfologisch plausibele maskers te genereren vanuit een geleerde kansverdeling, in plaats van één enkele deterministische voorspelling.

Methodologie

De auteurs stellen een nieuw kader voor: een Multi-Task Image-to-Image Schrödinger Bridge (SB) framework.

1. Schrödinger Bridge (SB) Formulering:

• In plaats van een deterministische mapping, modelleert de SB een stochastisch transportproces tussen de verdeling van de inputbeelden en de verdeling van de geldige outputmaskers.
• Dit is een generatief model dat de conditional verdeling van geldige instantiemaskers leert gegeven een inputbeeld. Hierdoor worden morfologisch onwaarschijnlijke vormen (zoals samengevoegde cellen) van nature onderdrukt.
• Het behoudt fijne ruimtelijke correspondentie tussen input en output, wat essentieel is voor gedetailleerde celgrenzen.

2. Multi-Task Learning met Reverse Distance Map:

• Om de scheiding van dicht op elkaar gepakte cellen te verbeteren, introduceert het model een reverse distance map als extra leerdoel.
• Reverse Distance Map: In tegenstelling tot conventionele afstandkaarten (die de afstand tot de dichtstbijzijnde grens meten), wordt hier de afstand omgekeerd genormaliseerd: pixels dichter bij de celgrens krijgen hogere waarden.
• Het model leert tijdens training gezamenlijk het instantiemasker en deze reverse distance map te voorspellen. Dit dwingt het model om randinformatie expliciet te leren, wat de scheiding van cellen verbetert zonder dat post-processing nodig is.

3. Deterministische Inferentie:

• Hoewel het trainingsproces stochastisch is (een diffusieproces), vereist de inferentie een enkele, stabiele voorspelling.
• Tijdens inferentie wordt het stochastische ruisinjectie uitgeschakeld. Het model volgt deterministisch de meest waarschijnlijke trajectorie (het gemiddelde van de posterior-verdeling) van de input naar de output.
• De uiteindelijke celmaskers worden verkregen door het gemiddelde te nemen van de outputkanalen en een vaste drempelwaarde toe te passen. De reverse distance map wordt tijdens inferentie verworpen.

Belangrijkste Bijdragen

1. Formulering als Generatief Probleem: De eerste toepassing van een Schrödinger Bridge voor cel-instantiesegmentatie, die het probleem herschrijft als een distributiegebaseerde image-to-image generatie, waardoor afhankelijkheid van heuristische post-processing wordt verwijderd.
2. Structuurbehoudende Generatie: Een framework dat de ruimtelijke correspondentie behoudt en morfologisch plausibele celvormen genereert door te leren uit de onderliggende data-distributie.
3. Boundary-Aware Multi-Task Strategie: De introductie van een reverse distance map als auxiliary supervision binnen het SB-framework, wat leidt tot betere scheiding van individuele cellen.
4. Robuustheid zonder SAM of Post-processing: Het bewijs dat het model concurrerende prestaties levert zonder gebruik te maken van vooringevormde modellen zoals SAM (Segment Anything Model) of complexe post-processing pipelines.

Resultaten

De methode is getest op twee datasets: PanNuke (groot scala aan weefsels, veel data) en MoNuSeg (beperkte trainingsdata).

• PanNuke (Grote dataset):

  • De voorgestelde methode behaalde de hoogste bPQ (Binary Panoptic Quality), F1-score en Precision van alle vergeleken methoden.
  • Het presteerde beter dan state-of-the-art methoden zoals CellViT (dat gebruikmaakt van SAM-pretraining en post-processing), ondanks dat de voorgestelde methode noch SAM noch post-processing gebruikt.
  • Een ablatiestudie toonde aan dat de multi-task aanpak (masker + reverse distance map) superieur is aan single-task varianten.

• MoNuSeg (Beperkte dataset):

  • Het model toonde sterke robuustheid bij weinig trainingsdata.
  • Het behaalde een bPQ vergelijkbaar met CellViT met post-processing, en een hogere F1-score dan CellViT zonder post-processing.
  • Dit suggereert dat de distributiegebaseerde aanpak minder afhankelijk is van enorme hoeveelheden gelabelde data dan methoden die volledig vertrouwen op pretraining.

• Morfologische Analyse:

  • Analyse van de verdeling van celgrootte en circulariteit toonde aan dat het SB-model een outputdistributie genereert die dichter bij de ground truth ligt dan CellViT. CellViT produceerde vaker atypische vormen (in lage-dichtheidsgebieden), terwijl het SB-model deze onderdrukte.

Betekenis en Conclusie

Dit werk markeert een verschuiving in cel-instantiesegmentatie van deterministische, post-processing-afhankelijke pipelines naar distributiegebaseerde generatieve modellen.

• Eliminatie van Post-processing: De methode demonstreert dat complexe, dataset-specifieke post-processing overbodig kan worden gemaakt door het probleem correct te formuleren als een generatief transportprobleem.
• Morfologische Consistentie: Door te leren vanuit de verdeling van geldige maskers, produceert het model van nature morfologisch plausibele resultaten, wat de betrouwbaarheid in de biomedische analyse verhoogt.
• Toekomstperspectief: De resultaten suggereren dat Schrödinger Bridges een krachtig kader bieden voor fijnkorrelige beeldsegmentatie taken waar ruimtelijke precisie en structurele plausibiliteit cruciaal zijn, zelfs onder beperkte data-omstandigheden.