Diffusion Autoencoder for Unsupervised Artifact Restoration in Handheld Fundus Images

Resumo Técnico: Autoencoder de Difusão para Restauração Não Supervisionada de Artefatos em Imagens de Fundoscopia Portátil

1. O Problema

O avanço de dispositivos de fundoscopia portátil (handheld) e baseados em smartphones democratizou o acesso ao diagnóstico oftalmológico e à triagem de doenças. No entanto, as imagens capturadas por esses dispositivos frequentemente sofrem de artefatos severos, como reflexos de flash, variações de exposição (super ou subexposição) e desfoque induzido por movimento.

• Desafio: Esses artefatos degradam a qualidade da imagem e impedem análises downstream confiáveis (como segmentação de vasos ou diagnóstico de retinopatia diabética).
• Limitação das Soluções Atuais: A maioria dos modelos de restauração baseados em aprendizado profundo (GANs, modelos de difusão convencionais) depende de supervisão pareada (pares de imagem suja/limpa) ou de máscaras de artefato pré-definidas. Isso os torna pouco adaptáveis às degradações não estruturadas e variadas encontradas em imagens reais de dispositivos portáteis.

2. Metodologia Proposta

Os autores propõem um Autoencoder de Difusão Não Supervisionado (DiffAE) que integra um codificador de contexto ao processo de remoção de ruído (denoising). O sistema opera em duas fases principais:

• Fase de Treinamento (Não Supervisionada):

  • O modelo é treinado apenas com imagens de fundoscopia de alta qualidade adquiridas em sistemas de mesa (table-top), sem usar imagens portáteis defeituosas.
  • Utiliza uma arquitetura de Autoencoder de Difusão:
    1. Um Codificador de Contexto ($Enc_\phi$) mapeia a imagem limpa ($x_0$) para um vetor latente compacto ($z$). Este vetor captura estruturas retinianas globais e serve como um "prior semântico".
    2. Uma rede UNet ($f_\theta$) atua como o processo de difusão reversa, aprendendo a reconstruir a imagem original a partir de ruído, condicionada tanto pelo vetor latente $z$ quanto pelo tempo de difusão.
  • O objetivo é aprender uma representação semântica rica que permita a geração de estruturas retinianas realistas.

• Fase de Inferência (Restauração):

  • Para restaurar imagens portáteis com artefatos, o modelo utiliza o vetor latente $z$ (extraído da imagem degradada ou de uma representação limpa associada) e incorpora máscaras de artefato ($M$) durante o processo de difusão reversa.
  • Em cada passo de tempo $t$, a imagem é atualizada combinando a região mascarada (que precisa ser preenchida) com o ruído, guiada pelo contexto semântico $z$.
  • Isso permite a inpainting (preenchimento) de regiões degradadas sem necessidade de pares de treinamento, preservando a consistência anatômica (ex: alinhamento do disco óptico).

3. Contribuições Principais

1. Framework DiffAE Unsupervised: Introdução de um autoencoder de difusão que integra um codificador de contexto ao processo de denoising para aprender representações semanticamente significativas, focado na restauração de imagens de fundoscopia portátil.
2. Mecanismo de Inferência Não Supervisionado: Desenvolvimento de uma estratégia onde o modelo aprende a gerar apenas imagens de alta qualidade (table-top) e, na inferência, utiliza máscaras de artefato para restaurar imagens degradadas, eliminando a necessidade de dados de treinamento pareados.
3. Validação Abrangente: Validação quantitativa e qualitativa em conjuntos de dados não vistos durante o treinamento, incluindo tarefas downstream como segmentação de vasos e classificação de doenças.

4. Resultados e Avaliação

O modelo foi comparado com baselines de ponta (GANs como EdgeConnect, DeepFillV2, ShiftNet, MISF e modelos de difusão como ClarityNet).

• Métricas de Qualidade de Imagem (Conjunto Sintético):

  • O modelo proposto alcançou o melhor PSNR (38.64 dB) e SSIM (0.97), superando o melhor baseline (DeepFillV2) em 1.11 dB de PSNR.
  • Na Segmentação de Vasos (Dice Score), obteve 0.891, demonstrando superioridade na preservação da continuidade vascular em comparação aos outros métodos (que frequentemente introduziam artefatos ou borrões nas bordas dos vasos).

• Avaliação em Dados Reais (Conjunto de Teste mBRSET):

  • Avaliação de Qualidade: O modelo produziu a maior proporção de imagens classificadas como "Boas" (35.8%) e "Usáveis" (52.6%), com a menor taxa de rejeição (15.8%).
  • Classificação de Doença (Retinopatia Diabética): A restauração das imagens aumentou a precisão da classificação de DR para 81.17%, superando a linha de base sem restauração (77.11%) e todos os outros métodos de restauração. Isso indica que a restauração preserva informações contextuais críticas para o diagnóstico.

• Estudo de Ablação:

  • A comparação entre usar apenas a representação da imagem com artefato versus uma representação interpolada (mistura de imagem suja e limpa) mostrou que o uso de um contexto enriquecido melhora significativamente a qualidade da restauração (aumento de PSNR de 38.98 para 41.46 dB), validando a importância do codificador de contexto.

5. Significado e Conclusão

Este trabalho demonstra que a formação baseada em autoencoders de difusão é uma ferramenta poderosa para tarefas de restauração não supervisionada em cenários médicos complexos.

• Impacto Clínico: Ao permitir a restauração eficaz de imagens de dispositivos portáteis de baixo custo, o método viabiliza o uso de IA para triagem de doenças oculares em locais remotos ou com recursos limitados, onde a qualidade da imagem é historicamente um gargalo.
• Inovação Técnica: A abordagem supera a dependência de dados pareados e de estruturas de artefato pré-definidas, oferecendo uma solução robusta para degradações não estruturadas, mantendo a fidelidade anatômica necessária para diagnósticos precisos.