LesionDiffusion: Towards Text-controlled General Lesion Synthesis

1. O Problema

As métodos de reconhecimento de lesões em imagens médicas baseados em aprendizado supervisionado completo enfrentam desafios significativos devido à dependência de grandes conjuntos de dados anotados. A coleta e o compartilhamento desses dados são caros e difíceis devido a preocupações de privacidade e à escassez de dados públicos que contenham tanto descrições de lesões quanto máscaras correspondentes.

Embora a geração sintética de lesões tenha surgido como uma solução promissora, os modelos existentes (como GANs e modelos de difusão anteriores) apresentam limitações críticas:

• Falta de Escalabilidade: Eles geralmente treinam cada tipo de lesão independentemente, limitando a generalização.
• Controle Insuficiente: Há pouca capacidade de controle fino sobre os atributos da lesão (forma, densidade, localização).
• Estruturas Simples: A maioria dos modelos foca em lesões sólidas e esféricas, negligenciando estruturas complexas (como lesões ocosas em câncer gástrico ou colorretal) e texturas variadas.

2. Metodologia

O LesionDiffusion é um framework de síntese de lesões controlado por texto para imagens 3D de Tomografia Computadorizada (CT). O sistema gera tanto a lesão quanto a máscara correspondente, utilizando um template de relatório estruturado como condição textual.

O framework é composto por duas redes principais e um processo de treinamento em duas etapas:

A. Dados e Template Estruturado

• Foi criado um novo conjunto de dados com 1.505 exames de CT anotados, cobrindo 14 tipos de lesões em 8 órgãos diferentes.
• Cada lesão é acompanhada por um relatório estruturado contendo 10 atributos divididos em:
  • Atributos de Máscara: Forma, localização, relação com órgãos vizinhos, tamanho.
  • Atributos de Imagem: Densidade, variações de densidade, características de superfície, status de realce.

B. Arquitetura do Modelo

O processo é dividido em dois estágios:

1. Estágio I: Geração de Máscara de Lesão (LMNet)

   • Utiliza um modelo de difusão para gerar a máscara binária da lesão.
   • É guiado por atributos da máscara (extraídos do relatório) e por uma caixa delimitadora (bounding box) e semântica do órgão.
   • Introduz uma função de perda ponderada pela caixa delimitadora para garantir que a geração ocorra dentro da região restrita especificada.
   • Os atributos textuais são codificados usando o BiomedCLIP e fundidos via mecanismos de atenção cruzada.

2. Estágio II: Preenchimento de Imagem de Lesão (LINet)

   • Realiza o inpainting (preenchimento) da lesão na imagem CT original, substituindo a região da lesão por uma nova textura sintética.
   • Utiliza Latent Diffusion (Difusão em Espaço Latente) para lidar com o alto custo de memória de imagens 3D. Um autoencoder (VAE-GAN) comprime as imagens CT em um espaço latente.
   • O processo de difusão ocorre apenas dentro da região da lesão definida pela máscara.
   • É guiado por atributos de imagem (densidade, heterogeneidade, etc.) para controlar a aparência da lesão gerada.

C. Processo de Inferência

1. Gera-se um relatório estruturado (manualmente ou via LLM).
2. Extrai-se os atributos para gerar uma caixa delimitadora (bbox) aleatória dentro do órgão correto.
3. O LMNet gera a máscara da lesão.
4. O LINet realiza o inpainting na imagem CT codificada, usando a máscara e os atributos de imagem como condições.
5. Decodifica-se a imagem latente para obter a CT final com a lesão sintética.

3. Principais Contribuições

• Framework Geral e Controlável: Primeiro modelo capaz de sintetizar uma ampla variedade de lesões (sólidas e ocosas) em múltiplos órgãos com controle granular via texto.
• Novo Conjunto de Dados: Criação de um dataset de 1.505 CTs com máscaras e relatórios estruturados, preenchendo uma lacuna na disponibilidade de dados multimodais anotados.
• Arquitetura Dual: Separação eficiente entre a geração da geometria (máscara) e a textura (imagem), permitindo maior precisão e controle.
• Generalização Zero-Shot: Capacidade de gerar lesões em órgãos e tipos de lesões não vistos durante o treinamento, baseando-se apenas na descrição textual dos atributos.

4. Resultados

Os experimentos foram conduzidos focando na melhoria do desempenho de modelos de segmentação (usando nnUNet-v2) treinados com dados sintéticos.

• Desempenho de Segmentação: O uso de dados sintéticos do LesionDiffusion para pré-treinamento resultou em um aumento médio de 6,9% na pontuação DSC (Dice Similarity Coefficient) em comparação com o modelo concorrente DiffTumor.
• Comparação com Dados Reais: Após ajuste fino (fine-tuning) com dados reais, os modelos pré-treinados com LesionDiffusion superaram os modelos treinados apenas com dados reais em 2,4% e superaram os pré-treinados com DiffTumor em 3,5%.
• Generalização (Teste em Cérebro): O modelo foi testado na geração de lesões de hemorragia cerebral (hematencephalon), um órgão e tipo de lesão não presentes no treinamento. O LesionDiffusion manteve um desempenho robusto, enquanto o DiffTumor sofreu uma queda drástica de 46% no DSC devido à lacuna de domínio.
• Ablação de Texto: A remoção das condições textuais resultou em uma queda média de 15,4% no DSC, demonstrando que o controle textual é essencial para a precisão, especialmente em órgãos com múltiplos tipos de lesões.
• Alinhamento Textual: Métricas de convexidade e esfericidade confirmaram que o modelo responde corretamente a prompts como "redondo" vs. "irregular".

5. Significado e Conclusão

O LesionDiffusion representa um avanço significativo na síntese de dados médicos. Ao resolver o problema da escassez de dados anotados e oferecer um controle fino sobre a geração de lesões complexas, ele permite o desenvolvimento de modelos de diagnóstico mais robustos e generalizáveis.

Sua capacidade de generalizar para órgãos e patologias não vistos durante o treinamento sugere que ele pode servir como a base para um modelo fundamental (foundation model) de geração de lesões. Isso tem implicações profundas para a pesquisa médica, permitindo a criação de conjuntos de dados sintéticos diversificados para treinar algoritmos de IA, reduzindo a dependência de dados sensíveis de pacientes e acelerando o desenvolvimento de ferramentas de diagnóstico assistido.