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

Resumen Técnico de Brain-SAM

Planteamiento del Problema
La resonancia magnética (RM) es fundamental para la neuroimagen, pero la segmentación precisa de las estructuras cerebrales y las lesiones sigue siendo un desafío. Aunque los modelos fundacionales como el Segment Anything Model (SAM) y su sucesor SAM2 han demostrado una notable capacidad de generalización en imágenes naturales, su aplicación a la RM cerebral enfrenta obstáculos significativos. Las adaptaciones médicas existentes (por ejemplo, MedSAM, SAM-Med3D) suelen entrenarse en conjuntos de datos de órganos mixtos que contienen estructuras como el corazón, el hígado y los pulmones. Estos modelos tienen dificultades para adaptarse a las características anatómicas únicas del cerebro, tales como la compleja organización tisular, los límites difusos entre la materia gris y blanca, y las morfologías de lesiones altamente diversas. Además, la mayoría de las adaptaciones basadas en SAM convierten los escaneos volumétricos 3D en cortes 2D, omitiendo el contexto espacial entre los cortes. Asimismo, la dependencia de prompts interactivos manuales limita la escalabilidad, ya que los clínicos pueden tener dificultades para proporcionar prompts precisos para patologías sutiles como las lesiones epilépticas. Finalmente, el decodificador de máscara original de SAM, diseñado para imágenes naturales 2D con bordes claros, suele tener un rendimiento subóptimo en los datos de RM cerebral, caracterizados por bordes difusos y heterogeneidad tisular.

Metodología
Los autores proponen Brain-SAM, un modelo diseñado específicamente para la segmentación de lesiones en RM cerebral 3D. La arquitectura extiende el marco de SAM2 mediante tres modificaciones principales:

1. Adaptación 3D del Codificador de Imágenes: El codificador Hiera 2D original (un Vision Transformer jerárquico) se adapta para procesar directamente datos volumétricos 3D. Esto implica la extensión de cuatro módulos clave:

   • Patch Embedding: Reemplazo de las convoluciones 2D por convoluciones 3D para particionar las entradas volumétricas en parches de vóxeles.
   • Position Embedding: Extensión de las representaciones de 2D a 3D para capturar las relaciones espaciales a través de la profundidad, altura y anchura.
   • Token Sequence: Aplanamiento de los volúmenes 3D a lo largo de las tres dimensiones para crear secuencias de tokens más largas adecuadas para la atención volumétrica.
   • Q-Pooling: Reemplazo de MaxPool2d por MaxPool3d para permitir el submuestreo de características a través de las dimensiones de profundidad y espacial.

2. Decodificador inspirado en UNETR: Para abordar las limitaciones del decodificador original de SAM al manejar texturas locales finas y detalles de los bordes, los autores introducen un decodificador inspirado en UNETR. Este componente integra la capacidad de modelado global de los Transformers con la representación local de grano fino de las CNN mediante la fusión de características multiescala. A diferencia de los decodificadores UNETR estándar que requieren transformaciones de resolución, este diseño aprovecha las características jerárquicas inherentes del codificador Hiera para una fusión directa y un upsampling progresivo. La arquitectura soporta una salida de doble cabeza: una para la segmentación totalmente automática y otra para generar características intermedias para la segmentación basada en prompts.

3. Estrategia de Preentrenamiento y Ajuste Fino (Fine-Tuning):

   • Preentrenamiento: El modelo se inicializa con los pesos de SAM2 y se somete a una etapa de preentrenamiento en el conjunto de datos BraTS2021 (1,000 casos, 4 modalidades) exclusivamente para segmentación automática, con el fin de establecer representaciones médicas robustas.
   • Ajuste Fino basado en Adaptadores: Para adaptarse a múltiples conjuntos de datos de lesiones cerebrales (tumores, accidentes cerebrovasculares, epilepsia) evitando el olvido catastrófico, se insertan módulos de adaptadores ligeros (arquitectura de cuello de botella con convoluciones 3D) antes de cada bloque Hiera. Durante el ajuste fino, el backbone permanece congelado y solo se actualizan los parámetros del adaptador compartido.
   • Mecanismo de Prompt: Se utilizan prompts de puntos como el método de interacción principal, extendido a 3D. La estrategia de generación de la máscara se modifica: en lugar de utilizar claves actualizadas en la operación de producto punto, el diseño fusiona las consultas procesadas con las representaciones de características intermedias del decodificador inspirado en UNETR.
   • Refinamiento: Un mecanismo de refinamiento de postprocesamiento utiliza clics del usuario para ajustar localmente las probabilidades de la máscara mediante núcleos gaussianos 3D.

Contribuciones Clave

1. Extensión 3D de SAM2: El trabajo extiende con éxito el marco de SAM2 de imágenes naturales 2D al análisis de imágenes médicas 3D, permitiendo el procesamiento volumétrico directo sin la aproximación de corte por corte.
2. Arquitectura de Decodificador Híbrido: La introducción de un decodificador inspirado en UNETR permite que el modelo soporte tanto la segmentación automática como la interactiva dentro de una arquitectura unificada, capturando eficazmente tanto los detalles locales de grano fino como las dependencias contextuales globales.
3. Curación Exhaustiva de Datos: Los autores curaron un conjunto de datos a gran escala que cubre tumores cerebrales, accidentes cerebrovasculares y epilepsia, permitiendo que el modelo soporte una amplia gama de aplicaciones clínicas.

Resultados
Se realizaron experimentos extensos en siete conjuntos de datos públicos (BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI) que cubren seis modalidades de RM. Brain-SAM se comparó contra modelos específicos de la tarea (nnU-Net, UNETR, Swin UNETR) y modelos fundacionales (SAM-Med3D, VISTA3D).

• Rendimiento General: Brain-SAM logró puntuaciones Dice superiores y distancias de Hausdorff al 95% (HD95) más bajas en comparación con todos los modelos de referencia. Específicamente, comparado con nnU-Net, logró mejoras en el Dice del 22% en lesiones epilépticas, 9% en metástasis cerebrales y 6% en meningiomas.
• Segmentación de Lesiones Pequeñas: El modelo demostró ventajas significativas en la segmentación de lesiones pequeñas (por ejemplo, epilepsia), logiendo puntuaciones Dice entre un 15% y 18% más altas que los modelos de referencia robustos. Mantuvo un mejor equilibrio entre la sensibilidad (recall) y la precisión para objetivos pequeños (<5,000 vóxeles) en comparación con otros métodos.
• Estudios de Ablación:
  • Reemplazar el backbone 3D-ViT por un 3D-Hiera mejoró las puntuaciones Dice hasta en un 14% para lesiones pequeñas y redujo significativamente el HD95.
  • La integración del decodificador inspirado en UNETR aumentó la puntuación Dice promedio en un 8% en comparación con la versión sin él.
  • La estrategia de ajuste fino basada en adaptadores demostró ser altamente eficiente en parámetros (39.4M de parámetros ajustables frente a 72.02M para el ajuste fino completo) logrando un rendimiento comparable al ajuste fino completo.

Significancia y Reivindicaciones
El artículo afirma que Brain-SAM ofrece un modelo preentrenado robusto para tareas de análisis de RM cerebral descendentes. Al abordar los desafíos anatómicos y estructurales específicos del tejido cerebral —como los bordes indistintos y el tamaño de las lesiones pequeñas—, proporciona una segmentación más confiable y precisa que los modelos fundacionales generalistas existentes. Los autores postulan que esta capacidad puede apoyar una toma de decisiones clínicas mejor informada, particularmente en el diagnóstico temprano y la planificación del tratamiento para condiciones como la epilepsia y los tumores cerebrales. La capacidad del modelo para funcionar tanto en modo automático como guiado por prompts tiene como objetivo reducir la carga de anotación manual manteniendo una alta precisión. Los autores señalan que el estudio actual se limita a datos de una sola modalidad y sugieren que trabajos futuros podrían explorar la integración de múltiples modalidades para mejorar aún más la caracterización de las lesiones.