OOD-SEG: Exploiting out-of-distribution detection techniques for learning image segmentation from sparse multi-class positive-only annotations

Resumen Técnico: OOD-SEG

1. Planteamiento del Problema

El artículo aborda dos desafíos críticos en la segmentación de imágenes médicas y quirúrgicas basadas en redes neuronales profundas:

1. Escasez de anotaciones densas: Obtener etiquetas a nivel de píxel completas es costoso, consume mucho tiempo y requiere expertos de dominio. En consecuencia, muchos conjuntos de datos utilizan anotaciones esparcidas (solo se etiquetan regiones representativas) y exclusivamente positivas (solo se marcan las clases de interés, sin etiquetar el fondo o "background").
2. Falta de detección de datos fuera de distribución (OOD): Los sistemas de segmentación convencionales asumen que todos los datos de prueba provienen de la misma distribución que los de entrenamiento (in-distribution, ID). Esto es peligroso en entornos clínicos, ya que el modelo puede generar predicciones erróneas y confiables para píxeles que pertenecen a clases no vistas, nuevos instrumentos quirúrgicos o artefactos, sin alertar al usuario.

El problema central es cómo realizar una segmentación multiclase robusta utilizando solo anotaciones positivas esparcidas, sin asumir que los píxeles no etiquetados son necesariamente fondo, y al mismo tiempo detectar píxeles que sean OOD.

2. Metodología Propuesta

Los autores proponen OOD-SEG, un marco de trabajo que reformula la segmentación con anotaciones positivas esparcidas como un problema de aprendizaje Positivo-No Etiquetado (PU Learning) a nivel de píxel, integrando técnicas de detección OOD.

• Formulación del Problema:

  • Se asume que los píxeles etiquetados pertenecen a clases positivas (ID).
  • Los píxeles no etiquetados pueden pertenecer a clases positivas no anotadas o al fondo (negativo/OOD).
  • En lugar de tratar los píxeles no etiquetados como fondo (lo que introduce ruido), el modelo aprende solo de los píxeles positivos y utiliza la detección OOD para identificar qué píxeles no encajan en las distribuciones de las clases conocidas.

• Arquitectura y Entrenamiento:

  • Se utiliza una red base tipo U-Net con un codificador EfficientNet-B4 preentrenado.
  • Entrenamiento: La red se entrena únicamente con las anotaciones positivas esparcidas utilizando una función de pérdida de entropía cruzada estándar. La salida de la red tiene $C$ canales (donde $C$ es el número de clases positivas), sin un canal explícito para el fondo.
  • Inferencia y Detección OOD: Se introduce una función de puntuación ($S$) a nivel de píxel. Si la puntuación máxima de confianza para las clases ID supera un umbral $\tau$, el píxel se clasifica en la clase ID correspondiente. Si la puntuación es baja, el píxel se marca como OOD (fondo/anomalía).

• Métodos de Puntuación OOD Evaluados:
  El marco integra diferentes técnicas clásicas de detección OOD adaptadas a la segmentación:

  1. Baseline: Softmax directo de la red.
  2. ODIN: Aplicación de escalado de temperatura y perturbación adversaria (aunque en este trabajo se omitió la perturbación para reducir costos computacionales y mantener comparabilidad justa).
  3. Mahalanobis: Cálculo de la distancia de Mahalanobis entre las características de la imagen y las distribuciones gaussianas de las clases ID.
  4. GODIN (Generalized ODIN): Una adaptación novedosa propuesta por los autores que extiende el método GODIN (originalmente para clasificación) a la segmentación mediante el uso de capas convolucionales (en lugar de capas totalmente conectadas o lineales) para aprender parámetros de temperatura y similitud a nivel de píxel, capturando contexto espacial.

• Estrategia de Validación Cruzada de Dos Niveles:
  Para evaluar la detección OOD sin necesidad de conjuntos de datos externos (que son difíciles de obtener en medicina), los autores proponen una estrategia única:

  • Particiones por Sujeto (SP): Aseguran que no haya superposición de pacientes entre entrenamiento y prueba.
  • Particiones por Clase (CP): Durante el entrenamiento de un pliegue, se "retienen" (hold-out) algunas clases positivas anotadas. Estas clases retenidas se tratan como clases OOD para ese pliegue específico.
  • Esto permite evaluar la capacidad del modelo para detectar clases que no vio durante el entrenamiento, utilizando un solo conjunto de datos.

• Métricas de Evaluación:
  Se proponen métricas independientes del umbral: AUROC (Área bajo la curva ROC) y AUPR (Área bajo la curva Precisión-Recall), calculadas tratando las clases retenidas como negativas (OOD) y las clases entrenadas como positivas (ID).

3. Resultados Clave

Los experimentos se realizaron en tres conjuntos de datos:

1. Heiporspectral: Imágenes hiperespectrales de tejidos porcinos (20 órganos).
2. ODSI-DB: Imágenes hiperespectrales de tejidos dentales humanos (35 tejidos).
3. DSAD: Imágenes RGB laparoscópicas de cirugía rectal (11 órganos).

• Rendimiento de Detección OOD:

  • El método ODIN demostró consistentemente el mejor rendimiento en las métricas AUROC y AUPR en todos los conjuntos de datos. Esto sugiere que una puntuación de confianza bien calibrada es crucial para distinguir entre píxeles ID y OOD.
  • El método GODIN (con la adaptación convolucional) mostró un rendimiento sólido, superando al Mahalanobis en varios casos, lo que indica que la adaptación espacial es beneficiosa.
  • El método basado en Mahalanobis tuvo un rendimiento inferior, probablemente debido a la alta dimensionalidad de las características y la dificultad de estimar matrices de covarianza estables en espacios de alta dimensión con datos esparcidos.

• Precisión en Clases ID:

  • La introducción del umbral OOD ($\tau_m$) no degradó la precisión de clasificación para las clases in-distribution (ID) en comparación con el modelo sin detección de outliers ($\tau=0$).
  • El marco logra rechazar píxeles de fondo o anomalías sin sacrificar la detección de los tejidos de interés.

• Análisis Cualitativo:

  • Las visualizaciones muestran que el modelo con detección OOD genera máscaras más limpias, eliminando falsos positivos en áreas de fondo o estructuras no etiquetadas que el modelo base (sin OOD) clasificaba erróneamente como alguna de las clases conocidas.

4. Contribuciones Principales

1. Nuevo Marco de Aprendizaje: Presentan el primer enfoque que formula la segmentación multiclase con anotaciones positivas esparcidas como un problema de aprendizaje PU a nivel de píxel, utilizando herramientas de detección OOD.
2. Protocolo de Evaluación Innovador: Desarrollan una estrategia de validación cruzada de dos niveles (sujeto y clase) que permite evaluar la detección OOD dentro de un único conjunto de datos, eliminando la necesidad de conjuntos de datos externos OOD costosos de anotar.
3. Adaptación de GODIN: Introducen una adaptación convolucional del método GODIN, extendiendo su aplicabilidad a tareas de segmentación en lugar de solo clasificación.
4. Validación Multidisciplinaria: Demuestran la robustez del método en diferentes modalidades de imagen médica (hiperespectral y RGB) y diferentes contextos quirúrgicos.

5. Significado e Impacto

• Seguridad Clínica: Al permitir que los sistemas de segmentación identifiquen cuándo están operando fuera de su distribución de entrenamiento (ej. nuevos instrumentos, tejidos patológicos no vistos), se reduce el riesgo de errores médicos catastróficos en entornos automatizados.
• Viabilidad de Datos: Facilita el uso de datos clínicos existentes que solo tienen anotaciones parciales y positivas, democratizando el desarrollo de modelos de IA para áreas donde la anotación densa es imposible.
• Estándar de Evaluación: La propuesta de métricas y validación cruzada establece un nuevo estándar para evaluar la robustez de los modelos de segmentación médica frente a anomalías, un área históricamente descuidada.

En conclusión, OOD-SEG ofrece una solución práctica y robusta para la segmentación médica en condiciones de anotación limitada, integrando la seguridad de la detección de anomalías directamente en el proceso de inferencia sin comprometer la precisión de las clases conocidas.