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

Aquí se presenta un resumen técnico detallado del artículo "DIFFUSAM: DIFFUSION-BASED PROMPT-FREE SAM2 FOR FEW-SHOT AND SOURCE-FREE MEDICAL IMAGE SEGMENTATION."

1. Planteamiento del Problema

La segmentación de imágenes médicas es crítica para el diagnóstico clínico, pero enfrenta obstáculos significativos con los modelos base más avanzados actuales, como el Segment Anything Model (SAM) y su sucesor SAM2.

• Brecha de Dominio: SAM y SAM2 se entrenan en imágenes naturales y tienen dificultades para generalizar a datos médicos debido a las diferencias en textura, contraste y estructuras anatómicas.
• Dependencia de Prompts: Estos modelos requieren prompts proporcionados por el usuario (puntos, cajas o texto) para funcionar, lo cual es poco práctico para flujos de trabajo clínicos automatizados.
• Escasez de Datos: El ajuste fino supervisado requiere grandes conjuntos de datos anotados por expertos, cuya creación es costosa y consume mucho tiempo.
• Limitaciones de las Adaptaciones Existentes: Los intentos previos de adaptar SAM/SAM2 a datos médicos a menudo requieren un ajuste fino computacionalmente pesado o aún dependen de prompts. Además, las variantes sin prompts existentes no aprovechan las capacidades de memoria volumétrica de SAM2.

2. Metodología: DiffuSAM

Los autores proponen DiffuSAM, un marco que adapta SAM2 para la segmentación médica sin prompts utilizando un modelo de difusión latente condicional. La idea central es sintetizar "incrustaciones de memoria" que SAM2 pueda utilizar para generar máscaras, reemplazando efectivamente la necesidad de prompts del usuario.

Arquitectura Central

1. Esqueleto Congelado: El codificador de imágenes de SAM2 y el codificador de memoria permanecen congelados durante el entrenamiento. El modelo solo aprende un prior de difusión ligero.
2. Entrenamiento del Prior de Difusión:
   • Entrada: El modelo toma la incrustación de imagen ($z_{img}$) del codificador de imágenes congelado de SAM2 y la etiqueta de clase ($l$).
   • Objetivo: Aprende a generar la incrustación de memoria correspondiente ($z_{mem}$) que SAM2 produciría normalmente si se le diera un prompt.
   • Proceso: Un modelo de difusión basado en UNet se entrena bajo un programa de Modelo Probabilístico de Difusión con Eliminación de Ruido (DDPM). Aprende a eliminar el ruido de una incrustación de memoria ruidosa condicionada a la incrustación de imagen y a la etiqueta.
   • Función de Pérdida: El objetivo de entrenamiento combina una pérdida de prior (reconstruir la incrustación de memoria) y una pérdida de segmentación (asegurar que la incrustación generada produzca máscaras precisas al pasarla por el decodificador de máscaras de SAM2).
3. Inferencia (Generación sin Prompts):
   • En la inferencia, el modelo de difusión comienza con ruido puro y la incrustación de imagen.
   • Elimina el ruido iterativamente para sintetizar una incrustación de memoria ($\hat{z}_{mem}$).
   • Esta incrustación sintetizada se inyecta en la atención de memoria y el decodificador de máscaras de SAM2 para producir la máscara de segmentación final sin ninguna entrada del usuario.

Consistencia Volumétrica (Condicionamiento 3D)

Para abordar la naturaleza 3D de los volúmenes médicos (TC/RM), DiffuSAM impone consistencia espacial a través de las rebanadas:

• Al generar la incrustación de memoria para una rebanada específica $i$, el modelo se condiciona en las predicciones de segmentación de las rebanadas adyacentes ($adj$).
• Las predicciones de las rebanadas adyacentes se codifican mediante el codificador de memoria y se fusionan con los datos de la rebanada actual.
• Esto asegura la continuidad anatómica y previene inconsistencias entre rebanadas en el volumen 3D.

Adaptación de Dominio No Supervisada Libre de Origen (SF-UDA)

El marco está diseñado para SF-UDA, donde el modelo se entrena en un dominio fuente (por ejemplo, TC) y se prueba en un dominio objetivo (por ejemplo, RM) sin acceso a imágenes fuente durante la adaptación al objetivo.

• El método aprovecha la propiedad de afinidad de características: las incrustaciones de imagen de SAM2 para la misma clase tienden a agruparse estrechamente incluso a través de diferentes dominios.
• El prior de difusión entrenado en la fuente se aplica directamente a las rebanadas del dominio objetivo sin más ajuste fino, generando incrustaciones de memoria pseudo para el dominio objetivo.

3. Contribuciones Clave

1. Marco Novel: Introducción de DiffuSAM, la primera adaptación basada en difusión de SAM2 para la segmentación de imágenes médicas sin prompts.
2. Operación en Espacio Latente: El modelo opera en el espacio de incrustaciones latentes de SAM2, lo que le permite capturar distribuciones de características de alto nivel sin modificar el esqueleto pesado.
3. Integración Volumétrica: Integración con el mecanismo de atención de memoria de SAM2 para asegurar la continuidad anatómica a través de rebanadas 3D mediante condicionamiento entre rebanadas.
4. Eficiencia: El enfoque requiere entrenar un UNet ligero en lugar de ajustar fino todo el modelo base, resultando en un bajo uso de memoria (<5 GB de GPU) y una inferencia rápida (solo 2 pasos de difusión).

4. Resultados Experimentales

El método se evaluó en dos conjuntos de datos principales: BTCV (escaneos TC) y CHAOS (escaneos RM), centrándose en la segmentación del bazo, los riñones y el hígado.

Configuración de Few-Shot (BTCV)

• Configuración: Entrenado con el 10% de los datos (3 volúmenes), probado con el 90% (27 volúmenes).
• Rendimiento: DiffuSAM logró un puntuación Dice promedio del 87.24%.
• Comparación: Superó a SAM/SAM2 nativos (que dependen de prompts) y a modelos UNet/UNETR supervisados entrenados en el mismo pequeño conjunto de datos. Fue competitivo con FATE-SAM (que requiere rebanadas de verdad fundamental en el momento de la prueba), pero logró esto sin prompts manuales ni supervisión en el momento de la prueba.
• Ablación: Añadir condicionamiento volumétrico entre rebanadas mejoró el Dice promedio del 82.72% (DiffuSAM-no 3D) al 87.24%.

Configuración de UDA Libre de Origen (TC $\to$ RM)

• Configuración: Entrenado en BTCV (TC), probado en CHAOS (RM) sin acceso a imágenes fuente.
• Rendimiento: DiffuSAM logró un Dice promedio del 85.2%.
• Comparación: Superó significativamente a otras líneas base libres de origen (DPL: 56.0%) y rindió a la par con DFG (85.2%), un método fuerte basado en afinidad de características.
• Ablación: El condicionamiento volumétrico mejoró el rendimiento del 80.4% al 85.2%.

5. Significado y Conclusión

• Aplicabilidad Clínica: Al eliminar la necesidad de prompts del usuario, DiffuSAM cierra la brecha entre los modelos base y los flujos de trabajo clínicos automatizados.
• Eficiencia de Datos: El método demuestra que se puede lograr una segmentación de alta calidad con datos de entrenamiento mínimos (few-shot) y sin acceso a datos fuente durante la adaptación al objetivo (SF-UDA).
• Eficiencia Computacional: A diferencia del ajuste fino completo de modelos grandes, DiffuSAM entrena un prior de difusión ligero, haciéndolo accesible para entornos con recursos limitados.
• Trabajo Futuro: Los autores reconocen que, aunque MedSAM (entrenado de extremo a extremo con prompts) tiene un mejor rendimiento, DiffuSAM ofrece una ventaja única al evitar prompts y el ajuste fino completo. El trabajo futuro explorará la integración de técnicas adicionales de SF-UDA.

En resumen, DiffuSAM aprovecha con éxito el poder generativo de los modelos de difusión para "alucinar" las incrustaciones de memoria necesarias para SAM2, permitiendo una segmentación de imágenes médicas robusta, sin prompts y adaptable al dominio.