PrefPaint: Enhancing Medical Image Inpainting through Expert Human Feedback

1. Problema

O inpainting (preenchimento de partes faltantes ou corrompidas de uma imagem) é uma tarefa crucial na geração de imagens médicas, especialmente para a criação de dados sintéticos de treinamento para Inteligência Artificial (IA) clínica. O foco específico deste trabalho é a geração de imagens sintéticas de pólipos (precursoras do câncer colorretal) que sejam anatomicamente precisas.

Os desafios principais identificados são:

• Inacurácia Anatómica: Modelos de geração padrão (como Stable Diffusion) frequentemente produzem pólipos sintéticos com padrões anatômicos incorretos ou texturas irreais, o que pode levar a falsos padrões e erros no diagnóstico clínico.
• Custo Computacional do RLHF: A técnica tradicional de Reinforcement Learning from Human Feedback (RLHF) requer o treinamento de um modelo de recompensa intermediário, o que demanda grandes conjuntos de dados, arquiteturas otimizadas e um custo computacional massivo, tornando-a inviável para instituições médicas com recursos limitados.
• Falta de Validação de Especialistas: A geração de dados sintéticos sem a validação direta de especialistas (como gastroenterologistas) resulta em modelos que aprendem padrões errôneos.

2. Metodologia

O sistema proposto, PrefPaint, integra feedback humano direto no processo de fine-tuning (ajuste fino) de um modelo de inpainting baseado em Stable Diffusion, utilizando uma abordagem eficiente em recursos.

• Algoritmo de Otimização (D3PO): Em vez do RLHF tradicional, o sistema utiliza D3PO (uma variação baseada em Direct Preference Optimization e Denoising Diffusion Policy Optimization).
  • O processo de denoising é tratado como um Processo de Decisão de Markov (MDP).
  • O modelo é ajustado diretamente com base em preferências binárias humanas (como/desgostei), eliminando a necessidade de treinar um modelo de recompensa intermediário. Isso reduz drasticamente o custo computacional.
• Arquitetura do Sistema:
  • Interface Web Interativa: Desenvolvida com Golang (backend), templ e Tailwind CSS (frontend), permitindo que médicos sem background técnico utilizem a ferramenta.
  • Gestão de Tarefas Assíncrona: Utiliza uma fila FIFO e threads dedicadas para gerenciar tarefas pesadas de GPU (inferência e fine-tuning) sem bloquear a interface do usuário.
  • Model Tree (Árvore de Modelos): Um conceito inovador de HCI (Interação Humano-Computador) que visualiza a evolução do modelo. Cada nó na árvore representa uma versão do modelo ajustada, permitindo aos especialistas rastrear o histórico de melhorias e selecionar a melhor versão para feedback adicional.
• Fluxo de Trabalho:
  1. O modelo gera imagens de pólipos com base em máscaras e prompts.
  2. Especialistas médicos avaliam as imagens usando um sistema binário simples (Like/Dislike).
  3. O feedback é usado para atualizar diretamente os parâmetros do modelo via D3PO.
  4. O processo é iterativo, gerando uma nova versão do modelo na "Árvore de Modelos".

3. Principais Contribuições

1. Solução Eficiente para IA Clínica: Apresenta um método viável para integrar feedback de especialistas em ambientes clínicos com recursos computacionais limitados, contornando a barreira do custo do RLHF tradicional.
2. Interface "Model Tree": Introduz uma nova interface de visualização hierárquica que permite aos médicos entenderem a proveniência e a evolução do modelo de IA, facilitando a confiança e a adoção da tecnologia.
3. Geração de Dados Sintéticos de Alta Fidelidade: Demonstra que o uso de feedback humano direto resulta em imagens de pólipos sintéticos que são anatomicamente corretas e visualmente realistas, superando métodos state-of-the-art.
4. Plataforma Integrada: Desenvolvimento de um sistema completo (backend, frontend e gestão de tarefas) que une especialistas médicos e engenheiros de IA em um único fluxo de trabalho iterativo.

4. Resultados

O sistema foi avaliado quantitativa e qualitativamente utilizando o conjunto de dados Kvasir-SEG (pólipos sésseis e pedunculados).

• Métricas Quantitativas: O método proposto superou consistentemente as bases de comparação (Stable Diffusion Inpainting, DreamBooth e SD2i) em todas as métricas:
  • Menor erro L1 e L2.
  • Maior SSIM (Índice de Similaridade Estrutural) e PSNR.
  • Menor LPIPS (Similaridade Perceptual), indicando maior realismo.
• Estudo de Usuário (Pilot Study):
  • Participaram 28 pessoas, incluindo 3 médicos e estudantes de ciência da computação.
  • Validação Médica: Especialistas médicos atribuíram pontuações significativamente mais altas ao PrefPaint em comparação com os métodos existentes (quase o dobro da pontuação em precisão clínica).
  • Métricas Subjetivas: O método obteve as melhores notas em Alinhamento, Realismo, Suavidade e Qualidade visual.
  • Usabilidade: A interface web recebeu avaliações positivas quanto à adaptabilidade e facilidade de uso, mesmo para usuários não técnicos.

5. Significado

O PrefPaint representa um avanço significativo na intersecção entre visão computacional e medicina clínica. Ao resolver o problema da geração de dados sintéticos imprecisos e ao criar um mecanismo acessível para a incorporação de conhecimento de domínio (médico) no treinamento de modelos de IA, o trabalho:

• Aumenta a Segurança Clínica: Reduz o risco de modelos de IA aprenderem padrões anatômicos falsos, o que é crítico para o diagnóstico precoce de câncer colorretal.
• Democratiza o Uso de IA Avançada: Ao eliminar a necessidade de modelos de recompensa complexos e caros, torna possível que hospitais e instituições de pesquisa menores utilizem técnicas de fine-tuning guiadas por humanos.
• Estabelece um Novo Padrão de Interação: A "Árvore de Modelos" oferece um novo paradigma para a colaboração humano-IA, onde a evolução do modelo é transparente e gerenciável por especialistas do domínio.

Em resumo, o artigo propõe uma solução prática e robusta para melhorar a qualidade de imagens médicas sintéticas, garantindo que a IA clínica seja treinada com dados que reflitam fielmente a realidade biológica humana.