SwinTextUNet: Integrating CLIP-Based Text Guidance into Swin Transformer U-Nets for Medical Image Segmentation

1. Planteamiento del Problema

La segmentación precisa de imágenes médicas es fundamental para el diagnóstico asistido por computadora y la planificación del tratamiento. Sin embargo, los modelos tradicionales, como las U-Net basadas en CNN, presentan limitaciones significativas:

• Dependencia exclusiva de características visuales: Ignoran la información textual (informes de radiología, notas diagnósticas) que los clínicos utilizan rutinariamente para interpretar las imágenes.
• Falta de robustez: Cuando las características de la imagen son ambiguas, de bajo contraste o ruidosas, los modelos unimodales tienden a fallar.
• Brecha en la integración multimodal: Aunque existen modelos de visión-idioma (VLM) como CLIP, su integración en arquitecturas de segmentación médica supervisada y end-to-end sigue siendo poco explorada.

2. Metodología

El artículo propone SwinTextUNet, un marco de segmentación multimodal que fusiona embeddings textuales derivados de CLIP con un backbone de U-Net basado en Transformadores Swin. La arquitectura consta de cuatro componentes principales:

• Codificación de Texto (CLIP): Se utiliza un codificador de texto CLIP preentrenado (congelado) para convertir prompts diagnósticos (ej. "infección pulmonar bilateral") en embeddings semánticos. Estos se proyectan a la dimensión de los tokens visuales mediante una capa lineal.
• Codificación Visual (Swin Transformer): La imagen de entrada se procesa mediante un codificador jerárquico de Swin Transformer, que captura dependencias de largo alcance y contexto global mediante ventanas desplazadas (shifted windows), manteniendo la eficiencia computacional.
• Mecanismos de Fusión Multimodal:
  • Atención Cruzada (Cross-Attention): En cada etapa del codificador, los tokens visuales se refinan atendiendo a los tokens de texto. Esto permite que el contexto semántico guíe las características visuales sin abrumar las estructuras locales.
  • Fusión Convolucional (ConvFuse): En el decodificador, los mapas de características de salto (skip connections) se fusionan con las características del decodificador mediante bloques convolucionales para mejorar la reconstrucción de bordes y la continuidad espacial.
• Entrenamiento: El modelo se optimiza con una función de pérdida híbrida que combina Entropía Cruzada y Coeficiente Dice, utilizando el optimizador AdamW.

3. Contribuciones Clave

1. Arquitectura Multimodal Innovadora: Presentación de una variante de U-Net con Transformador Swin que integra explícitamente embeddings textuales de CLIP para la segmentación médica.
2. Mecanismo de Condicionamiento de Texto: Desarrollo de un mecanismo que fusiona embeddings semánticos con tokens de imagen a través de las etapas de codificador y decodificador, logrando una alineación cruzada mejorada.
3. Validación de la Integración Visión-Lenguaje: Demostración de que la incorporación de guías textuales mejora la robustez y la precisión en la segmentación de lesiones ambiguas, superando a los enfoques unimodales.

4. Resultados Experimentales

El modelo se evaluó en el conjunto de datos QaTa-COV19 (9,258 radiografías de tórax con anotaciones de lesiones de COVID-19 y descripciones textuales).

• Rendimiento Óptimo: La variante de 4 etapas logró el mejor equilibrio entre rendimiento y complejidad, obteniendo:
  • Dice Score: 86.47%
  • IoU (Intersección sobre Unión): 78.2%
• Estudios de Ablación:
  • La eliminación de la guía de texto redujo el Dice en un 7.3% y el IoU en un 8.8%, confirmando la importancia crítica de los priores semánticos.
  • La eliminación de la fusión convolucional (ConvFuse) y la atención cruzada también resultó en caídas significativas de rendimiento.
• Comparación con Baselines: SwinTextUNet superó consistentemente a modelos basados en CNN (U-Net, Attention U-Net) y otros modelos basados en Transformadores (LViT, Swin-UNet), demostrando una mayor precisión en la delimitación de bordes y una reducción de falsos positivos.
• Análisis de Profundidad: Se compararon variantes de 3, 4 y 5 etapas. La de 4 etapas ofreció el mejor compromiso, mientras que la de 5 etapas introdujo una sobrecarga de parámetros con retornos decrecientes.

5. Significado e Impacto

Este trabajo subraya el potencial de la integración visión-lenguaje en la inteligencia artificial médica.

• Avance Clínico: Al incorporar el contexto textual (informes médicos) directamente en el proceso de segmentación, el modelo se alinea más estrechamente con el flujo de trabajo clínico real, donde los médicos interpretan imágenes junto con notas descriptivas.
• Robustez: La capacidad de utilizar guías semánticas ayuda a resolver ambigüedades en imágenes de baja calidad o con patrones poco claros, lo cual es crucial para el diagnóstico preciso.
• Futuro: Aunque el estudio se limitó a imágenes 2D y un solo conjunto de datos, sienta las bases para futuras extensiones hacia conjuntos de datos 3D, múltiples enfermedades y la validación en flujos de trabajo clínicos reales, potencialmente mejorando las herramientas de diagnóstico asistido.