OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation with Diverse Shapes and Locations

Resumen Técnico de OTLesMix: Baricentro de Wasserstein y Mapa de Transporte Óptimo para la Generación de Lesiones Sintéticas

Planteamiento del Problema
El aprendizaje profundo ha revolucionado la segmentación de imágenes médicas, sin embargo, el entrenamiento supervisado sigue estando limitado por el alto costo y el tiempo requerido para las anotaciones de expertos, particularmente para tareas de segmentación 3D densas. Si bien el aumento de datos (data augmentation) es una técnica estándar para mejorar el entrenamiento de modelos, los métodos tradicionales (transformaciones espaciales/de intensidad) y los modelos generativos avanzados (GANs, Modelos de Difusión) presentan limitaciones específicas. Los modelos generativos requieren conjuntos de datos grandes y diversos para su propio entrenamiento, los cuales pueden no estar disponibles en contextos médicos. Por el contrario, las estrategias de aumento basadas en "mezcla" existentes (por ejemplo, MixUp, CutMix, CarveMix, LesionMix) a menudo no logran generar una variabilidad suficiente en las formas y ubicaciones de las lesiones. Típicamente, dependen del muestreo aleatorio simple o de transformaciones geométricas rígidas de las máscaras existentes, lo que resulta en muestras sintéticas que carecen de la diversidad necesaria para entrenar robustamente modelos de segmentación.

Metodología: OTLesMix
Los autores proponen OTLesMix, un novedoso método de síntesis de imágenes que aprovecha la teoría del Transporte Óptimo (OT) para generar lesiones sintéticas realistas y diversas. El método opera en dos etapas principales para crear una nueva muestra $(\tilde{X}, \tilde{Y})$ a partir de un conjunto de datos base $D = \{(X_i, Y_i)\}$:

1. Generación de la Máscara de la Lesión mediante el Baricentro de Wasserstein:

   • Se seleccionan aleatoriamente dos muestras base $(X_i, Y_i)$ y $(X_j, Y_j)$.
   • Las coordenadas de los píxeles no nulos en las máscaras de verdad de campo (ground truth) $Y_i$ y $Y_j$ se tratan como distribuciones probabilísticas.
   • Se genera una nueva máscara de lesión $\tilde{Y}_{les}$ calculando el baricentro de Wasserstein de estas dos distribuciones. Esto implica optimizar las coordenadas de los puntos de la nueva máscara para minimizar la suma ponderada de los costos de transporte hacia las máscas de origen:
     $$B(\lambda, \alpha_i, \alpha_j, \beta) = \lambda OT(\alpha_i, \beta) + (1-\lambda)OT(\alpha_j, \beta)$$
     donde $\lambda \sim U(0,1)$ es un parámetro de ponderación aleatorio.
   • Este proceso interpola entre las formas y ubicaciones de las dos máscaras originales, permitiendo que la nueva lesión exista en cualquier lugar entre las ubicaciones originales con una forma que es una interpolación suave de las originales.
   • Para asegurar la eficiencia computacional, los autores utilizan la divergencia de Sinkhorn sin sesgo (de-biased Sinkhorn divergence) y el algoritmo de Sinkhorn, lo que permite procesar hasta $10^6$ puntos en segundos en GPUs modernas.

2. Generación de Imagen y Etiquetas mediante el Mapa de Transporte Óptimo:

   • Una vez definida la nueva máscara $\tilde{Y}_{les}$, se deben determinar las intensidades de los píxeles y las etiquetas para la lesión sintética.
   • Se calcula un plan de transporte óptimo entre la nueva máscara $\tilde{Y}_{les}$ y una de las máscaras de origen (por ejemplo, $Y_j$).
   • Para cada píxel en la nueva máscara, el mapa de transporte $T$ determina su ubicación correspondiente en la imagen de origen $X_j$.
   • Síntesis de Intensidad: La intensidad del píxel sintético se obtiene mediante la interpolación lineal de las intensidades de la imagen de origen en la ubicación mapeada.
   • Síntesis de Etiquetas: Para tareas multiclas, la etiqueta de verdad de campo se asigna mediante la interpolación de vecino más cercano del mapa de etiquetas de origen.
   • La lesión sintetizada se inpainta (se integra) posteriormente en una imagen base (ya sea $X_i$, $X_j$, o una tercera imagen del conjunto de datos). Los pasos de post-procesamiento eliminan intersecciones con lesiones existentes o ubicaciones anatómicas irreales (por ejemplo, cráneo, LCR).

Contribuciones Clave

• Marco de Síntesis Novedoso: Introducción de OTLesMix, que combina de manera única los baricentros de Wasserstein para la interpolación de forma/ubicación y los mapas de transporte óptimo para la transferencia de intensidad/etiqueta.
• Diversidad Mejorada: El método aborda explícitamente la limitación de los métodos de mezcla existentes al generar lesiones con formas y ubicaciones diversas que no son meras copias o rotaciones simples de los datos existentes.
• Desempeño de Vanguardia (State-of-the-Art): El método demuestra un desempeño superior al mejorar las tareas de segmentación descendentes en comparación con las técnicas de aumento de datos basadas en mezcla estándar.

Resultados Experimentales
Los autores evaluaron OTLesMix en tres conjuntos de datos de segmentación de lesiones cerebrales utilizando el backbone nnUNet:

1. BraTS 2020: Segmentación de glioblastoma y glioma de bajo grado (cortes 2D).
2. ATLAS v2.0: Neuroimagen de accidente cerebrovascular (volúmenes 3D).
3. ISLES 2022: Segmentación de lesiones de accidente cerebrovascular isquémico (volúmenes 3D).

Hallazgos Cuantitativos:

• OTLesMix superó a los métodos de mezcla de vanguardia (MixUp, CutMix, CarveMix) y al baseline (nnUNet solo con aumento de datos tradicional) en todos los conjuntos de datos.
• Mejora sobre el Baseline: El Coeficiente de Similitud de Dice (DSC) mejoró entre 2.9 y s 6.6 puntos en comparación con los modelos entrenados sin datos sintéticos.
• Mejora sobre Competidores: OTLesMix logró un incremento de 0.8 a 1.8 puntos en el DSC sobre el segundo mejor método.
• Significancia Estadística: Las mejoras fueron estadísticamente significativas ($p \le 0.05$) en comparación con otros métodos generativos, siendo el valor p más alto 0.012 (vs. CarveMix en ATLAS).

Hallazgos Cualitativos:
La inspección visual confirmó que OTLesMix produce lesiones con mayor diversidad en forma y ubicación en comparación con otros métodos. El método manejó con éxito la segmentación multiclas (BraTS) y lesiones con múltiples componentes conectados (ATLAS, ISLES), manteniendo la consistencia en las formas interpoladas.

Significancia y Limitaciones
El artículo afirma que OTLesMix proporciona una solución robusta para el aumento de datos en imágenes médicas donde los datos son escasos, ofreciendo un punto medio entre la simplicidad de los métodos de mezcla y la complejidad de entrenar modelos generativos completos. Al aprovechar el transporte óptimo, asegura que las muestras sintéticas sean consistentes geométricamente y en intensidad con la distribución de los datos subyacentes.

Sin embargo, los autores reconocen limitaciones específicas:

• Costo Computacional: Generar una sola muestra toma aproximadamente 1 segundo para imágenes 2D y de 1 a 3 minutos para volúmenes 3D, dependiendo de la resolución.
• Restricciones de Realismo: El método depende de la interpolación lineal de las intensidades de los píxeles. Esto puede conducir a lesiones inpaintadas poco realistas si existe una brecha de dominio significativa (por ejemplo, diferentes escáneres o protocolos de adquisición) entre las dos imágenes de origen utilizadas para la mezcla.
• Dirección Futura: Los autores sugieren que el trabajo futuro podría combinar el procedimiento de generación de máscaras de OTLesMix con modelos de generación condicionados por máscara para mejorar aún más el realismo de las intensidades sintetizadas.