Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Résumé technique de Brain-SAM

Énoncé du problème
L'imagerie par résonance magnétique (IRM) est fondamentale pour la neuroimagerie, pourtant la segmentation précise des structures cérébrales et des lésions reste un défi. Bien que les modèles de fondation comme le Segment Anything Model (SAM) et son successeur SAM2 aient démontré une remarquable capacité de généralisation sur les images naturelles, leur application à l'IRM cérébrale se heurte à des obstacles importants. Les adaptations médicales existantes (par exemple, MedSAM, SAM-Med3D) sont généralement entraînées sur des ensembles de données d'organes mixtes contenant des structures telles que le cœur, le foie et les poumons. Ces modèles peinent à s'adapter aux caractéristiques anatomiques uniques du cerveau, telles que l'organisation complexe des tissus, les limites indistinctes entre la substance grise et la substance blanche, et la grande diversité des morphologies de lésions. De plus, la plupart des adaptations basées sur SAM convertissent les scans volumétriques 3D en coupes 2D, négligeant ainsi le contexte spatial entre les coupes. Par ailleurs, le recours à des invites interactives manuelles limite l'évolutivité, car les cliniciens peuvent éprouver des difficultés à fournir des invites précises pour des pathologies subtiles comme les lésions épileptiques. Enfin, le décodeur de masque original de SAM, conçu pour des images naturelles 2D aux contours nets, est souvent sous-optimal sur les données d'IRM cérébrale caractérisées par des bords flous et une hétérogénéité tissulaire.

Méthodologie
Les auteurs proposent Brain-SAM, un modèle conçu spécifiquement pour la segmentation de lésions en IRM cérébrale 3D. L'architecture étend le cadre de SAM2 par trois modifications principales :

1. Adaptation 3D de l'encodeur d'image : L'encodeur Hiera 2D original (un Transformer de vision hiérarchique) est adapté pour traiter directement des données volumétriques 3D. Cela implique l'extension de quatre modules clés :

   • Patch Embedding : Remplacement des convolutions 2D par des convolutions 3D pour partitionner les entrées volumétriques en patchs de voxels.
   • Position Embedding : Extension des représentations de 2D à 3D pour capturer les relations spatiales à travers la profondeur, la hauteur et la largeur.
   • Token Sequence : Aplatissement des volumes 3D selon les trois dimensions pour créer des séquences de tokens plus longues, adaptées à l'attention volumétrique.
   • Q-Pooling : Remplacement de MaxPool2d par MaxPool3d pour permettre le sous-échantillonnage des caractéristiques à travers les dimensions de profondeur et spatiales.

2. Décodeur inspiré de UNETR : Pour remédier aux limites du décodeur original de SAM dans la gestion des textures locales fines et des détails de bordure, les auteurs introduisent un décodeur inspiré de UNETR. Ce composant intègre la capacité de modélisation globale des Transformers avec la représentation locale fine des CNN via une fusion de caractéristiques multi-échelles. Contra�à des décodeurs UNETR standards qui nécessitent des transformations de résolution, cette conception exploite les caractéristiques hiérarchiques intrinsèques de l'encodeur Hiera pour une fusion directe et un suréchantillonnage progressif. L'architecture prend en charge une sortie à double tête : une pour la segmentation entièrement automatique et une autre pour générer des caractéristiques intermédiaires pour la segmentation basée sur des invites (prompts).

3. Stratégie d'entraînement et de fine-tuning :

   • Pré-entraînement : Le modèle est initialisé avec les poids de SAM2 et subit une étape de pré-entraînement sur le jeu de données BraTS2021 (1 000 cas, 4 modalités) exclusivement pour la segmentation automatique afin d'établir des représentations médicales robustes.
   • Fine-tuning basé sur des adaptateurs : Pour s'adapter à plusieurs ensembles de données de lésions cérébrales (tumeurs, AVC, épilepsie) tout en évitant l'oubli catastrophique, des modules d'adaptateurs légers (architecture bottleneck avec convolutions 3D) sont insérés avant chaque bloc Hiera. Pendant le fine-tuning, le backbone reste gelé et seuls les paramètres partagés de l'adaptateur sont mis à jour.
   • Mécanisme d'invite (Prompt) : Des invites de points sont utilisées comme principale méthode d'interaction, étendue à la 3D. La stratégie de génération de masque est modifiée : au lieu d'utiliser des clés mises à jour dans l'opération produit scalaire, la conception fusionne les requêtes traitées avec les représentations de caractéristiques intermédiaires du décodeur inspiré de UNETR.
   • Raffinement : Un mécanisme de raffinement post-traitement utilise des clics d'utilisateurs pour ajuster localement les probabilités de masque via des noyaux gaussiens 3D.

Contributions clés

1. Extension 3D de SAM2 : Ce travail étend avec succès le cadre de SAM2 des images naturelles 2D vers l'analyse d'images médicales 3D, permettant un traitement volumétrique direct sans approximation coupe par coupe.
2. Architecture de décodeur hybride : L'introduction d'un décodeur inspiré de UNETR permet au modèle de supporter à la fois la segmentation automatique et interactive au sein d'une architecture unifiée, capturant efficacement les détails locaux fins et les dépendances contextuelles globales.
3. Curation complète de données : Les auteurs ont constitué un ensemble de données à grande échelle couvrant les tumeurs cérébrales, l'AVC et l'épilepsie, permettant au modèle de soutenir un large éventail d'applications cliniques.

Résultats
Des expériences approfondies ont été menées sur sept ensembles de données publics (BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI) couvrant six modalités d'IRM. Brain-SAM a été comparé à des modèles spécifiques à la tâche (nnU-Net, UNETR, Swin UNETR) et à des modèles de fondation (SAM-Med3D, VISTA3D).

• Performance globale : Brain-SAM a obtenu des scores Dice supérieurs et des distances de Hausdorff à 95 % (HD95) plus faibles que tous les modèles de référence. Plus précisément, par rapport à nnU-Net, il a réalisé des améliorations de Dice de 22 % sur les lésions épileptiques, 9 % sur les métastases cérébrales et 6 % sur les méningiomes.
• Segmentation de petites lésions : Le modèle a démontré des avantages significatifs dans la segmentation de petites lésions (ex. épilepsie), atteignant des scores Dice 15 % à 18 % plus élevés que les modèles de référence robustes. Il a maintenu un meilleur équilibre entre rappel et précision pour les cibles de petite taille (< 5 000 voxels) par rapport aux autres méthodes.
• Études d'ablation :
  • Le remplacement du backbone 3D-ViT par un 3D-Hiera a amélioré les scores Dice jusqu'à 14 % pour les petites lésions et a considérablement réduit l'HD95.
  • L'intégration du décodeur inspiré de UNETR a augmenté le score Dice moyen de 8 % par rapport à la version sans ce dernier.
  • La stratégie de fine-tuning basée sur des adaptateurs s'est révélée très efficace en termes de paramètres (39,4 M de paramètres ajustables contre 72,02 M pour un fine-tuning complet) tout en atteignant des performances comparables au fine-tuning complet.

Signification et affirmations
L'article affirme que Brain-SAM offre un modèle pré-entraîné robuste pour les tâches d'analyse d'IRM cérébrale en aval. En abordant les défis anatomiques et structurels spécifiques des tissus cérébraux — tels que les contours indistincts et la taille réduite des lésions — il fournit une segmentation plus fiable et précise que les modèles de fondation généralistes existants. Les auteurs postulent que cette capacité peut favoriser une meilleure prise de décision clinique, particulièrement dans le diagnostic précoce et la planification des traitements pour des pathologies comme l'épilepsie et les tumeurs cérébrales. La capacité du modèle à fonctionner à la fois en mode automatique et guidé par des invites vise à réduire la charge de travail d'annotation manuelle tout en maintenant une haute précision. Les auteurs notent que l'étude actuelle est limitée aux données mono-modalités et suggèrent que les travaux futurs pourraient explorer l'intégration multi-modale pour améliorer davantage la caractérisation des lésions.