PrefPaint: Enhancing Medical Image Inpainting through Expert Human Feedback

1. Planteamiento del Problema

La inpainting (reconstrucción de imágenes) es una tarea crucial en el procesamiento de imágenes médicas, especialmente para generar datos sintéticos de pólipos que sirvan para entrenar modelos de IA clínica. Sin embargo, existen desafíos significativos:

• Falta de datos anotados: Generar imágenes sintéticas realistas y anatómicamente precisas es difícil sin validación de expertos.
• Riesgo de errores clínicos: Las imágenes generadas por modelos estándar pueden contener patrones anatómicos incorrectos, lo que lleva a diagnósticos erróneos en tareas posteriores.
• Costo computacional de RLHF: Las técnicas tradicionales de Aprendizaje por Refuerzo con Retroalimentación Humana (RLHF) requieren entrenar modelos de recompensa intermedios, lo cual es computacionalmente costoso y poco práctico para instituciones médicas con recursos limitados.

2. Metodología Propuesta: PrefPaint

El sistema PrefPaint integra retroalimentación humana directa en el proceso de entrenamiento de modelos de difusión (específicamente Stable Diffusion Inpainting) mediante un enfoque optimizado y una interfaz interactiva.

A. Algoritmo de Entrenamiento (D3PO)

En lugar de utilizar RLHF completo, el sistema emplea D3PO (Direct Preference Optimization basado en preferencias), una variante que evita la necesidad de un modelo de recompensa entrenado:

• Proceso: Se conceptualiza el proceso de denoising como un Proceso de Decisión de Markov (MDP) de múltiples pasos.
• Mecanismo: El modelo genera dos imágenes basadas en un prompt y una máscara. Un experto médico selecciona la imagen preferida (binario: "buena" o "mala").
• Actualización: Se utiliza una función de pérdida (DDPO loss) para actualizar directamente los parámetros del modelo basándose en estas preferencias binarias, eliminando el sobrecosto computacional de entrenar un modelo de recompensa intermedio.

B. Arquitectura del Sistema Interactivo

Se desarrolló una plataforma web robusta para facilitar la participación de expertos médicos (incluso sin conocimientos técnicos profundos):

• Backend: Construido en Golang con arquitectura Server-Side Rendering (SSR) y base de datos PostgreSQL para manejar operaciones CRUD en tiempo real.
• Frontend: Utiliza templ (motor de plantillas tipado) y Tailwind CSS para una interfaz reactiva y ligera.
• Gestión de Tareas: Implementa una cola FIFO y hilos asíncronos para ejecutar scripts de entrenamiento en Python sin bloquear la interfaz de usuario, permitiendo que los expertos continúen trabajando mientras el modelo se ajusta en segundo plano.

C. Interfaz de Usuario (UI) y Concepto HCI

• Model Tree (Árbol de Modelos): Una interfaz novedosa que visualiza la evolución del modelo como un árbol jerárquico. El nodo raíz es el modelo base, y las ramas representan iteraciones de ajuste fino. Esto permite a los expertos rastrear el progreso y la procedencia de las mejoras.
• Recopilación de Feedback: Los expertos califican las imágenes generadas con un sistema binario simple ("Me gusta" / "No me gusta"), minimizando la carga cognitiva.
• Inferencia Interactiva: Permite a los usuarios subir imágenes médicas, dibujar máscaras sobre regiones dañadas y proporcionar prompts de texto para generar reconstrucciones.

3. Contribuciones Clave

1. Solución Eficiente para Entornos Clínicos: Introduce un método práctico que integra feedback de expertos en la inpainting médica sin los costos computacionales prohibitivos del RLHF tradicional, utilizando D3PO.
2. Interfaz "Model Tree": Propone un concepto de HCI que visualiza la evolución y procedencia de los modelos ajustados, facilitando la gestión de versiones y la comprensión del progreso por parte de los clínicos.
3. Validación Rigurosa: Demuestra mediante métricas cuantitativas y estudios de usuarios que la integración de feedback experto mejora significativamente la precisión anatómica y la realismo visual en comparación con métodos state-of-the-art.

4. Resultados Experimentales

El sistema fue evaluado utilizando el conjunto de datos Kvasir-SEG (pólipos sésiles y pedunculados) y comparado contra DreamBooth, Stable Diffusion Inpainting (SDi) y SD2i.

• Métricas Cuantitativas: PrefPaint superó a todos los métodos basales en todas las métricas clave:
  • Menor error L1 y L2 (mejor precisión de píxeles).
  • Mayor SSIM (mejor similitud estructural).
  • Menor LPIPS (mayor realismo perceptual).
  • Ejemplo: En el subconjunto de pólipos sésiles, el error L1 se redujo de 0.0227 (SD2i) a 0.0169 (PrefPaint).
• Estudios de Usuario (Pilot Study):
  • Participaron 28 personas, incluyendo 3 médicos especialistas.
  • Validación Médica: Los médicos otorgaron puntuaciones significativamente más altas a PrefPaint en las categorías de alineación, realismo, suavidad y calidad general.
  • Feedback de Usabilidad: La interfaz web fue calificada positivamente por su adaptabilidad y facilidad de uso, incluso para usuarios no técnicos.

5. Significado e Impacto

El trabajo de PrefPaint es significativo porque cierra la brecha entre los flujos de trabajo complejos de aprendizaje automático y los usuarios clínicos finales.

• Fiabilidad Clínica: Al incorporar directamente el conocimiento de oncólogos y gastroenterólogos, el sistema genera datos sintéticos que respetan la anatomía real, reduciendo el riesgo de que los modelos de IA aprendan patrones erróneos.
• Accesibilidad: Al eliminar la necesidad de modelos de recompensa costosos y proporcionar una interfaz web intuitiva, hace que el ajuste fino de modelos de IA sea viable para hospitales e instituciones médicas con recursos limitados.
• Futuro de la IA Médica: Establece un nuevo paradigma para la generación de datos sintéticos en medicina, donde la validación humana continua es parte integral del ciclo de entrenamiento del modelo.