Domain and Task-Focused Example Selection for Data-Efficient Contrastive Medical Image Segmentation

1. Problema y Contexto

La segmentación de imágenes médicas es fundamental para el diagnóstico y la monitorización de enfermedades, pero los enfoques de aprendizaje supervisado completo requieren grandes cantidades de datos anotados a nivel de píxel. Este proceso es:

• Costoso y lento: Requiere la intervención de radiólogos expertos.
• Propenso a errores: La anotación manual puede ser inconsistente.
• Limitante: La escasez de datos anotados frena el desarrollo de modelos robustos, especialmente en escenarios de datos limitados o cuando se cambia de dominio (diferentes protocolos de adquisición de imágenes).

Aunque el aprendizaje auto-supervisado (SSL) y el aprendizaje contrastivo han surgido como alternativas, las estrategias existentes a menudo dependen de aumentos de datos artificiales que no capturan las relaciones inherentes entre imágenes médicas diferentes, o requieren que los datos estén perfectamente alineados, lo cual es difícil en tomografías computarizadas (CT) de diferentes pacientes.

2. Metodología Propuesta: PolyCL

Los autores proponen PolyCL, un nuevo marco de aprendizaje contrastivo auto-supervisado diseñado específicamente para imágenes médicas. La metodología se divide en dos etapas principales y una integración posterior con modelos de segmentación foundation.

A. Estrategias de Selección de Ejemplos (Pre-entrenamiento)

A diferencia de los métodos tradicionales que crean pares positivos mediante aumentos de una misma imagen, PolyCL selecciona pares de imágenes basándose en la relación intrínseca entre diferentes cortes (slices) de escaneos CT. Se proponen tres estrategias:

1. PolyCL-S (Basado en el Escaneo):
   • Lógica: Un ejemplo positivo se selecciona aleatoriamente de la misma exploración (escaneo) que la imagen ancla, y un negativo de un escaneo diferente.
   • Objetivo: Aprender relaciones intra-escaneo sin necesidad de conocer la anatomía específica.
2. PolyCL-O (Basado en el Órgano):
   • Lógica: Requiere saber si un corte contiene el órgano de interés (ej. hígado). Si la imagen ancla tiene el órgano, el positivo también debe tenerlo; si no, el positivo tampoco debe tenerlo. Los negativos se eligen de la categoría opuesta.
   • Objetivo: Enseñar al encoder a representar la estructura objetivo específica antes de ver etiquetas de segmentación.
3. PolyCL-M (Mixto):
   • Lógica: Combina las dos anteriores. Un ejemplo positivo debe provenir del mismo escaneo Y tener la misma presencia/ausencia del órgano que la imagen ancla.
   • Objetivo: Aprender tanto la invarianza inter-escaneo como la coherencia intra-escaneo, mejorando la discriminación de contextos similares pero diferentes.

B. Entrenamiento y Ajuste Fino (Fine-Tuning)

• Pre-entrenamiento: Se utiliza una función de pérdida contrastiva (InfoNCE) para alinear las representaciones de los pares positivos y separar las negativas en el espacio latente.
• Ajuste Fino: Se añade un decodificador (tipo U-Net) al encoder pre-entrenado y se entrena con una pequeña cantidad de datos etiquetados utilizando una función de pérdida basada en el coeficiente Dice.

C. Integración con Segment Anything Model (SAM)

Para mejorar aún más los resultados con datos limitados, se integran dos técnicas basadas en SAM:

1. Refinamiento de Máscara (Post-procesamiento): Las máscaras gruesas generadas por PolyCL se convierten en cuadros delimitadores (bounding boxes) que se usan como "prompts" para SAM. SAM refina los bordes de la máscara, mejorando la precisión anatómica.
2. Propagación de Máscara (SAM 2): Utilizando la capacidad de video de SAM 2, se toma una única slice 2D anotada (o generada por PolyCL) y se propaga a través de todo el volumen 3D del CT. Esto permite generar segmentaciones volumétricas completas con una anotación mínima.

3. Contribuciones Clave

1. Marco PolyCL: Un nuevo enfoque de aprendizaje contrastivo que utiliza estrategias de selección de ejemplos basadas en el dominio (órgano) y la tarea, evitando la dependencia de aumentos de datos agresivos.
2. Estrategias de Selección Innovadoras: Introducción de las variantes PolyCL-S, PolyCL-O y PolyCL-M, demostrando que la selección de pares basada en la anatomía y el escaneo es superior a los métodos estándar.
3. Integración Híbrida con SAM: Uso novedoso de SAM no solo como modelo de ceros (zero-shot), sino como módulo de refinamiento post-entrenamiento y como mecanismo de propagación 3D a partir de una sola slice.
4. Validación Exhaustiva: Evaluación en múltiples conjuntos de datos públicos (LiTS, TotalSegmentator, MSD, KiTS, BTCV) y escenarios de dominio cruzado.

4. Resultados Experimentales

Las evaluaciones se realizaron en tareas de segmentación de hígado y riñones en imágenes CT:

• Rendimiento In-Domain (LiTS/KiTS):
  • PolyCL (especialmente las variantes M y O) iguala o supera a los modelos totalmente supervisados (como U-Net, nnU-Net, TransUNet) en términos del coeficiente Dice, incluso con datos de entrenamiento limitados.
  • Logra mejoras significativas en la distancia de Hausdorff (precisión de los bordes), reduciendo errores de localización en comparación con métodos de aprendizaje auto-supervisado anteriores (como SimCLR o SupCon).
• Generalización Cross-Domain:
  • En pruebas de dominio cruzado (entrenar en LiTS, probar en TotalSegmentator o MSD), PolyCL demostró una superioridad clara sobre los modelos totalmente supervisados y otros métodos SSL.
  • Por ejemplo, en el dataset TotalSegmentator, la variante PolyCL-S con backbone ResU-Net mejoró el Dice en un 10.7% sobre el mejor modelo totalmente supervisado.
• Eficiencia con Datos Limitados (Reduced-Label):
  • Con solo el 5% de datos etiquetados, la combinación PolyCL-M + SAM alcanzó un Dice de 0.682, superando a la versión sin refinar en casi un 19% y a otros métodos de refinamiento (SAMRefiner).
  • La propagación con SAM 2 permitió generar segmentaciones 3D de alta calidad utilizando solo una slice anotada, manteniendo un rendimiento robusto incluso en configuraciones de ultra-baja anotación.
• Análisis Estadístico: Las mejoras en la distancia de Hausdorff fueron estadísticamente significativas (p < 0.01) para las estrategias PolyCL-O y PolyCL-M frente a las líneas base.

5. Significado e Impacto

El trabajo de PolyCL es significativo por varias razones:

• Reducción de la Carga de Anotación: Demuestra que es posible lograr un rendimiento de segmentación de clase mundial con una fracción mínima de datos etiquetados, lo cual es crucial para la adopción clínica donde la anotación es un cuello de botella.
• Robustez en el Cambio de Dominio: La capacidad de generalizar a nuevos datasets sin reentrenamiento específico sugiere que el modelo aprende características transferibles y relevantes para la tarea, no solo patrones específicos del dataset de entrenamiento.
• Sinergia con Modelos Foundation: La integración exitosa de SAM y SAM 2 en un flujo de trabajo de aprendizaje auto-supervisado abre nuevas vías para combinar la eficiencia del entrenamiento con la precisión de los modelos foundation, ofreciendo una solución práctica para la segmentación volumétrica 3D con mínima intervención humana.
• Viabilidad Clínica: Aunque introduce un pequeño costo computacional adicional (pre-procesamiento para selección de órganos y refinamiento con SAM), este costo es insignificante comparado con el tiempo que ahorraría a los radiólogos en la anotación manual de miles de cortes.

En conclusión, PolyCL representa un avance importante hacia la segmentación médica eficiente en datos, combinando estrategias de aprendizaje contrastivo inteligentes con la potencia de los modelos foundation modernos.