Self-supervised Pretraining of Cell Segmentation Models

Probleemstelling

De kwantificering van celmorfologie in microscopie-afbeeldingen is cruciaal voor onderzoek en klinische toepassingen, zoals het voorspellen van celgezondheid en het onderscheiden van kankercellen. De standaardmethode hiervoor is instance segmentation (het identificeren van alle pixels die bij elke individuele cel horen).

• Beperkingen: De vooruitgang wordt gehinderd door het gebrek aan hoogwaardige, gelabelde datasets. Handmatige annotatie is tijdrovend, vereist expertise en leidt tot inconsistenties tussen annotatoren.
• Huidige aanpak: Moderne deep-learning-modellen gebruiken vaak vooraf getrainde gewichten van foundation modellen die op natuurlijke afbeeldingen zijn getraind, zoals het Segment Anything Model (SAM).
• Het kernprobleem: Er bestaat een aanzienlijke domain shift (domeinverschuiving) tussen natuurlijke afbeeldingen en microscopie-afbeeldingen. In microscopie zijn objectranden minder scherp, zijn er weinig semantische grenzen en zijn celvormen extreem gevarieerd. Modellen die op natuurlijke afbeeldingen zijn getraind, coderen vaak "object-gevoel" en textuurpriors die slecht overeenkomen met biologische data, wat leidt tot slechte generalisatie en prestaties.

Methodologie: DINOCell

De auteurs stellen DINOCell voor, een zelftoezichtend (self-supervised) framework dat de representaties van DINOv2 aanpast aan het domein van microscopie voordat er supervisie plaatsvindt. Het proces bestaat uit drie fasen:

1. Zelftoezichtend Pretraining (DINOv2):

   • Het model start met de gewichten van DINOv2 (een Vision Transformer, ViT-B), getraind op een grote dataset van natuurlijke afbeeldingen. Dit levert sterke, generieke visuele representaties op.

2. Domeinadaptatie (Domain Adaptation):

   • Dit is de kerninnovatie. In plaats van direct te finetunen voor segmentatie, wordt het DINOv2-encoder verder getraind op een dataset van 130.000 ongelabelde celafbeeldingen (verzameld uit diverse openbare bronnen zoals Broad Institute, BBBC, IDR).
   • Deze training gebruikt de DINOv2-architectuur (student-teacher framework met self-distillation) om de visuele representaties aan te passen aan de structurele kenmerken van celmorfologie, zonder dat er handmatige labels nodig zijn.
   • De dataset bevat diverse modaliteiten (fasecontrast, helderveld, DIC, fluorescentie) en celtypen.

3. Supervisie Finetuning:

   • Het aangepaste model wordt vervolgens gefinetuned op een gelabelde dataset van 9.042 afbeeldingen met handgetekende segmentaties.
   • Architectuur: De encoder (ViT-B) produceert patch-embeddings die worden verwerkt door een lichte decoder (transposed convoluties).
   • Doel: In plaats van directe maskers te voorspellen (zoals U-Net), voorspelt het model object proxies (specifiek flow fields zoals in Cellpose: vectorvelden die naar het celcentrum wijzen en een cel-kanskaart). Deze velden worden later omgezet in instance maskers via gradient-tracking.
   • Technische details: Er wordt gebruikgemaakt van "Ignore Masking" om omgang te hebben met datasets waar annotatoren gebieden hebben overgeslagen (missende segmentaties), waardoor overfitting op onnauwkeurige data wordt voorkomen.

Belangrijkste Bijdragen

• DINOCell Framework: Een nieuwe aanpak die domeinadaptatie via zelftoezicht toepast op microscopie-afbeeldingen voordat supervisie plaatsvindt, in plaats van te vertrouwen op task-priors van natuurlijke afbeeldingen (zoals bij SAM).
• Grote Ongelabelde Dataset: Het samenstellen van een divers, ongelabeld corpus van 130k celafbeeldingen voor domeinadaptatie.
• Robuustheid: Het aantonen dat het leren van domein-specifieke visuele representaties (via DINO) effectiever is dan het leren van de segmentatietaak zelf vanuit willekeurige initialisatie of natuurlijke afbeeldingen.
• Nieuwe Evaluatiemetric: Introductie van MMA (Maximum Matching Accuracy) als aanvullende metric naast de standaard SEG en DET, om de overlap tussen voorspellingen en ground truth optimaal te matchen.

Resultaten

De prestaties zijn geëvalueerd op de LIVECell-benchmark (in-distribution) en drie out-of-distribution datasets (PBL-HEK, PBL-N2A, Glioma C6) voor zero-shot prestaties.

• In-Distribution (LIVECell):

  • DINOCell behaalde een SEG-score van 0.784.
  • Dit is een verbetering van 10,42% ten opzichte van toonaangevende SAM-gebaseerde modellen (zoals SAMCell en Cellpose-SAM).
  • DINOCell toont minder neiging tot het samenvoegen van dicht op elkaar liggende cellen (clumping) en mist minder cellen dan concurrenten.

• Zero-Shot Generalisatie:

  • DINOCell presteert aanzienlijk beter op onbekende datasets. Op de Glioma C6-dataset (extreem variabel in vorm en belichting) verbeterde DINOCell de MMA-score met 43,84% ten opzichte van de beste baseline (SAMCell-cyto).
  • Op andere zero-shot datasets (PBL-HEK, PBL-N2A) werden verbeteringen van 16-20% geobserveerd.

• Ablatie Studies:

  • Pretraining met DINOv2 alleen (zonder domeinadaptatie) leverde al een enorme verbetering op (+29,51% op testset, +52-285% op zero-shot).
  • Domeinadaptatie gaf een extra, zij het kleiner, voordeel voor generalisatie naar zeer ongebruikelijke celtypen (zoals Glioma).
  • Visualisatie van PCA (Principal Component Analysis) van de encoder-embeddings toont aan dat DINOCell veel scherpere en beter gescheiden objectrepresentaties leert dan modellen zonder pretraining of met alleen SAM-weights.

Betekenis en Conclusie

Het paper demonstreert dat voor microscopie-taken, waar de domeinverschuiving met natuurlijke afbeeldingen extreem groot is, het initialiseren van modellen met domein-geadaptede self-supervised pretraining superieur is aan het gebruik van task-priors van foundation modellen zoals SAM.

DINOCell bewijst dat het leren van domein-specifieke visuele priors (via zelftoezicht op ongelabelde data) een cruciale stap is om robuuste segmentatiemodellen te bouwen die goed generaliseren naar nieuwe, ongeziene microscopie-omstandigheden. Dit opent de deur voor nauwkeurigere en snellere analyse van celcultuur in onderzoek en klinische settings, zelfs met beperkte gelabelde data.