DiffuSAM: Diffusion-Based Prompt-Free SAM2 for Few-Shot and Source-Free Medical Image Segmentation

1. Énoncé du problème

La segmentation d'images médicales est cruciale pour le diagnostic clinique, mais elle fait face à des obstacles majeurs avec les modèles de base les plus avancés actuels, tels que le Segment Anything Model (SAM) et son successeur SAM2.

• Écart de domaine : SAM et SAM2 sont entraînés sur des images naturelles et peinent à généraliser aux données médicales en raison des différences de texture, de contraste et de structures anatomiques.
• Dépendance aux invites : Ces modèles nécessitent des invites fournies par l'utilisateur (points, boîtes ou texte) pour fonctionner, ce qui est peu pratique pour les flux de travail cliniques automatisés.
• Pénurie de données : Le fine-tuning supervisé nécessite de grands jeux de données annotés par des experts, dont la création est coûteuse et longue.
• Limites des adaptations existantes : Les tentatives précédentes d'adapter SAM/SAM2 aux données médicales nécessitent souvent un fine-tuning lourd en termes de calcul ou reposent encore sur des invites. De plus, les variantes sans invites existantes n'exploitent pas les capacités de mémoire volumétrique de SAM2.

2. Méthodologie : DiffuSAM

Les auteurs proposent DiffuSAM, un cadre qui adapte SAM2 pour la segmentation médicale sans invites en utilisant un modèle de diffusion latent conditionnel. L'idée centrale est de synthétiser des "embeddings de mémoire" que SAM2 peut utiliser pour générer des masques, remplaçant efficacement le besoin d'invites utilisateur.

Architecture de base

1. Backbone figé : L'encodeur d'images et l'encodeur de mémoire de SAM2 restent figés pendant l'entraînement. Le modèle n'apprend qu'un prior de diffusion léger.
2. Entraînement du prior de diffusion :
   • Entrée : Le modèle prend l'embedding d'image ($z_{img}$) provenant de l'encodeur d'images figé de SAM2 et l'étiquette de classe ($l$).
   • Cible : Il apprend à générer l'embedding de mémoire correspondant ($z_{mem}$) que SAM2 produirait normalement si une invite lui était fournie.
   • Processus : Un modèle de diffusion basé sur UNet est entraîné selon un calendrier de Modèle Probabiliste de Diffusion Débruitant (DDPM). Il apprend à débruiter un embedding de mémoire bruité conditionné par l'embedding d'image et l'étiquette.
   • Fonction de perte : L'objectif d'entraînement combine une perte de prior (reconstruction de l'embedding de mémoire) et une perte de segmentation (assurant que l'embedding généré produit des masques précis lorsqu'il est transmis au décodeur de masques de SAM2).
3. Inférence (Génération sans invites) :
   • Lors de l'inférence, le modèle de diffusion commence avec du bruit pur et l'embedding d'image.
   • Il débruite itérativement pour synthétiser un embedding de mémoire ($\hat{z}_{mem}$).
   • Cet embedding synthétisé est injecté dans l'attention de mémoire et le décodeur de masques de SAM2 pour produire le masque de segmentation final sans aucune entrée utilisateur.

Cohérence volumétrique (Conditionnement 3D)

Pour répondre à la nature 3D des volumes médicaux (CT/IRM), DiffuSAM impose une cohérence spatiale entre les tranches :

• Lors de la génération de l'embedding de mémoire pour une tranche spécifique $i$, le modèle se conditionne sur les prédictions de segmentation des tranches adjacentes ($adj$).
• Les prédictions des tranches adjacentes sont encodées via l'encodeur de mémoire et fusionnées avec les données de la tranche courante.
• Cela assure la continuité anatomique et prévient les incohérences entre les tranches dans le volume 3D.

Adaptation de domaine non supervisée sans source (SF-UDA)

Le cadre est conçu pour la SF-UDA, où le modèle est entraîné sur un domaine source (par exemple, CT) et testé sur un domaine cible (par exemple, IRM) sans accès aux images source lors de l'adaptation cible.

• La méthode exploite la propriété d'affinité des caractéristiques : les embeddings d'images de SAM2 pour une même classe tendent à se regrouper étroitement, même à travers différents domaines.
• Le prior de diffusion entraîné sur la source est appliqué directement aux tranches du domaine cible sans fine-tuning supplémentaire, générant des pseudo-embeddings de mémoire pour le domaine cible.

3. Contributions clés

1. Cadre novateur : Introduction de DiffuSAM, la première adaptation basée sur la diffusion de SAM2 pour la segmentation d'images médicales sans invites.
2. Opération dans l'espace latent : Le modèle opère dans l'espace d'embedding latent de SAM2, lui permettant de capturer des distributions de caractéristiques de haut niveau sans modifier le backbone lourd.
3. Intégration volumétrique : Intégration avec le mécanisme d'attention de mémoire de SAM2 pour assurer la continuité anatomique à travers les tranches 3D via un conditionnement inter-tranches.
4. Efficacité : L'approche nécessite l'entraînement d'un UNet léger plutôt que le fine-tuning complet du modèle de base, entraînant une faible utilisation de la mémoire (<5 Go de GPU) et une inférence rapide (seulement 2 étapes de diffusion).

4. Résultats expérimentaux

La méthode a été évaluée sur deux jeux de données majeurs : BTCV (scanners CT) et CHAOS (scanners IRM), en se concentrant sur la segmentation de la rate, des reins et du foie.

Configuration en peu d'exemplaires (BTCV)

• Configuration : Entraîné sur 10 % des données (3 volumes), testé sur 90 % (27 volumes).
• Performance : DiffuSAM a obtenu un score Dice moyen de 87,24 %.
• Comparaison : Il a surpassé SAM/SAM2 natifs (qui reposent sur des invites) et les modèles UNet/UNETR supervisés entraînés sur le même petit jeu de données. Il était compétitif avec FATE-SAM (qui nécessite des tranches de vérité terrain au moment du test) mais a atteint cela sans invites manuelles ni supervision au moment du test.
• Ablation : L'ajout d'un conditionnement inter-tranches volumétrique a amélioré le Dice moyen de 82,72 % (DiffuSAM-no 3D) à 87,24 %.

Configuration SF-UDA (CT $\to$ IRM)

• Configuration : Entraîné sur BTCV (CT), testé sur CHAOS (IRM) sans accès aux images source.
• Performance : DiffuSAM a obtenu un Dice moyen de 85,2 %.
• Comparaison : Il a nettement surpassé d'autres baselines sans source (DPL : 56,0 %) et a performé au même niveau que DFG (85,2 %), une méthode robuste basée sur l'affinité des caractéristiques.
• Ablation : Le conditionnement volumétrique a amélioré les performances de 80,4 % à 85,2 %.

5. Importance et conclusion

• Applicabilité clinique : En éliminant le besoin d'invites utilisateur, DiffuSAM comble le fossé entre les modèles de base et les flux de travail cliniques automatisés.
• Efficacité des données : La méthode démontre qu'une segmentation de haute qualité peut être obtenue avec un minimum de données d'entraînement (peu d'exemplaires) et sans accès aux données source lors de l'adaptation cible (SF-UDA).
• Efficacité computationnelle : Contrairement au fine-tuning complet de grands modèles, DiffuSAM entraîne un prior de diffusion léger, le rendant accessible pour des environnements aux ressources limitées.
• Travaux futurs : Les auteurs reconnaissent que bien que MedSAM (entraîné de bout en bout avec des invites) performe mieux, DiffuSAM offre un avantage unique en évitant les invites et le fine-tuning complet. Les travaux futurs exploreront l'intégration de techniques supplémentaires de SF-UDA.

En résumé, DiffuSAM exploite avec succès la puissance générative des modèles de diffusion pour "halluciner" les embeddings de mémoire nécessaires à SAM2, permettant une segmentation d'images médicales robuste, sans invites et adaptable aux domaines.