Self-supervised Pretraining of Cell Segmentation Models

1. Problématique

La segmentation d'instances en microscopie est cruciale pour l'analyse morphologique des cellules, la prédiction de la santé cellulaire et la détection de maladies. Cependant, l'entraînement de modèles performants est fortement contraint par la pénurie de jeux de données étiquetés de haute qualité.

Les approches récentes tentent de contourner ce problème en initialisant des modèles avec des poids pré-entraînés sur des modèles de fondation pour images naturelles, tels que le Segment Anything Model (SAM). Toutefois, ces modèles souffrent d'un décalage de domaine (domain shift) significatif :

• Les représentations apprises sur des images naturelles encodent des priors d'« objectité » et de texture qui ne correspondent pas bien aux données de microscopie.
• En microscopie, les contours des objets sont souvent peu distincts, les limites sémantiques sont rares et les formes cellulaires sont extrêmement variées.
• En conséquence, les modèles basés sur SAM peinent à généraliser en dehors de leur distribution d'entraînement, entraînant une dégradation des performances.

2. Méthodologie : DINOCell

Les auteurs proposent DINOCell, un cadre d'apprentissage auto-supervisé conçu pour adapter les représentations visuelles à la microscopie avant l'étape de fine-tuning supervisé. L'hypothèse centrale est que l'apprentissage de représentations visuelles adaptées au domaine est plus limitant que l'apprentissage de la carte de segmentation elle-même.

Le processus s'articule en trois phases :

A. Collecte de Données

• Données d'adaptation de domaine (Non étiquetées) : 130 000 images de cellules non étiquetées provenant de 63 sources publiques (incluant le Broad Institute Cellpainting Gallery, BBBC, IDR). Ces données couvrent plus de 50 types cellulaires et quatre modalités (contraste de phase, champ clair, DIC, fluorescence).
• Données de Fine-Tuning (Étiquetées) : 9 042 images provenant de 15 jeux de données annotés manuellement (incluant LIVECell, BCCD, etc.).

B. Architecture et Prétraitement

• Encodeur : Utilisation d'un transformateur d'images ViT-B initialisé avec les poids de DINOv2.
  • La taille de patch est réduite de 14x14 à 8x8 pour capturer des détails plus fins, nécessitant un redimensionnement des embeddings de position.
  • L'encodeur produit des embeddings de 112x112 (dimension 768).
• Décodeur : Un décodeur léger composé de blocs de convolution transposée et d'upsampling pour restaurer la résolution à 512x512.
• Objectif de prédiction : Inspiré de Cellpose, le modèle ne prédit pas directement un masque, mais des champs de flux vectoriels (flows) et une probabilité cellulaire. Ces flux pointent vers le centre de la cellule, permettant de reconstruire les instances par suivi de gradient.

C. Stratégie d'Entraînement

1. Pré-entraînement auto-supervisé (DINOv2) : Utilisation des poids initiaux pré-entraînés sur 1,2 milliard d'images naturelles.
2. Adaptation de domaine (Domain Adaptation) : Entraînement continu du modèle DINOv2 sur les 130 000 images de microscopie non étiquetées. Cela permet d'aligner les représentations sur les structures biologiques tout en préservant les priors visuels généraux.
3. Fine-tuning supervisé : Entraînement complet du modèle (encodeur dégelé) sur les données étiquetées pour prédire les flux et les probabilités.
   • Technique clé : Ignore Masking. Pour gérer les annotations incomplètes fréquentes dans les jeux de données publics (où certaines cellules sont omises), les pixels non étiquetés sont pondérés à 0,05 dans la fonction de perte. Cela évite que le modèle n'apprenne des faux positifs sur ces zones.

3. Contributions Clés

• DINOCell : Un nouveau cadre de segmentation qui remplace l'initialisation par SAM (priors de tâche) par une initialisation par DINOv2 adaptée au domaine (priors de domaine).
• Adaptation de domaine auto-supervisée : Démonstration qu'un pré-entraînement supplémentaire sur des images de microscopie non étiquetées améliore significativement la capacité de généralisation, même sans étiquettes.
• Gestion des données bruitées : Introduction efficace de l'Ignore Masking pour stabiliser l'entraînement sur des jeux de données comportant des annotations manquantes.
• Évaluation rigoureuse : Comparaison sur le benchmark LIVECell et sur trois jeux de données out-of-distribution (PBL-HEK, PBL-N2A, Glioma C6) pour tester la robustesse en zero-shot.

4. Résultats

Les performances sont évaluées via trois métriques : SEG (précision de segmentation), DET (détection) et MMA (Maximum Matching Accuracy, une métrique développée par les auteurs).

• Sur le benchmark LIVECell (In-distribution) :

  • DINOCell atteint un score SEG de 0,784.
  • Amélioration de 10,42 % par rapport aux modèles basés sur SAM (Cellpose-SAM et SAMCell).
  • Meilleure capacité à séparer les cellules denses et à éviter le regroupement (clumping) ou la sous-segmentation.

• Performance Zero-Shot (Out-of-distribution) :

  • DINOCell démontre une robustesse exceptionnelle sur des données non vues lors de l'entraînement.
  • Sur le jeu de données Glioma C6 (morphologies très variables et conditions d'éclairage changeantes), DINOCell surpasse le meilleur modèle de base (SAMCell-cyto) de 43,84 % en MMA. Les modèles concurrents échouent souvent totalement sur ce jeu de données.
  • Sur PBL-HEK et PBL-N2A, les améliorations sont respectivement de 16,84 % et 8,73 % en MMA.

• Analyse des Représentations (PCA) :

  • La visualisation des composantes principales des embeddings de l'encodeur montre que DINOCell développe des structures spatiales cohérentes et des frontières d'objets nettes dès les premières étapes du fine-tuning, contrairement aux modèles initialisés aléatoirement ou par SAM qui restent bruyants.

5. Signification et Conclusion

Ce travail démontre que pour des domaines spécialisés comme la microscopie cellulaire, l'adaptation de domaine via l'apprentissage auto-supervisé est supérieure à l'utilisation directe de modèles de fondation entraînés sur des images naturelles.

L'approche DINOCell prouve que :

1. Apprendre des représentations visuelles adaptées à la morphologie cellulaire est une étape critique, souvent plus limitante que l'apprentissage de la tâche de segmentation elle-même.
2. L'utilisation de données non étiquetées massives pour l'adaptation de domaine permet d'obtenir des gains de performance significatifs, en particulier pour la généralisation à de nouvelles conditions d'imagerie (zero-shot).
3. Cette méthode offre une solution robuste et efficace pour la quantification cellulaire, réduisant la dépendance aux annotations manuelles coûteuses et améliorant la fiabilité des analyses biologiques automatisées.