Diffusion Autoencoder for Unsupervised Artifact Restoration in Handheld Fundus Images

1. Problem Statement

Handheld fundus imaging devices have revolutionized ophthalmologic screening by making it portable and cost-effective. However, images captured by these devices frequently suffer from severe artifacts, including:

• Flash reflections.
• Exposure variations (over- and under-exposure).
• Motion-induced blur.

These degradations render the images "out-of-distribution" for standard AI diagnostic models trained on high-quality, table-top fundus images. Existing restoration methods (GANs, CNNs) often rely on paired supervision (clean/degraded pairs) or predefined artifact masks, making them less adaptable to the unstructured and variable nature of handheld image defects. There is a critical need for an unsupervised approach that can restore these images using only high-quality data to enable reliable downstream analysis (e.g., disease grading).

2. Methodology

The authors propose a Diffusion Autoencoder (DiffAE) framework designed to learn semantic representations from high-quality images and apply them to restore artifact-ridden handheld images during inference.

Core Architecture

• Training Phase (Unsupervised): The model is trained only on high-quality, artifact-free table-top fundus images (from the EyePACS dataset). It does not require paired degraded images.

  • Context Encoder ($Enc_\phi$): A learnable encoder maps a clean image $x_0$ into a compact latent vector $z$ (512 dimensions). This vector captures global retinal structures and serves as a "semantic prior."
  • Denoising UNet ($f_\theta$): A standard diffusion model that learns to reverse the noise addition process. It is conditioned on the latent vector $z$ and the noisy image $x_t$.
  • Objective: The model minimizes Mean Squared Error (MSE) to reconstruct the clean image, ensuring the latent space $z$ encodes semantically meaningful retinal features.

• Inference Phase (Restoration):

  • The trained model is applied to handheld images containing artifacts.
  • Mask Conditioning: Artifact masks ($M$) are generated (via identification/binarization or simulation) and applied at every denoising timestep.
  • Process: The reverse diffusion process is guided by the artifact mask and the semantic prior $z$ (learned from high-quality data). The equation for inpainting at timestep $t$ is:
    $$x^{inpaint}_t = M \odot x_t + (1 - M) \odot x^{noisy}_0$$
  • This allows the model to "hallucinate" or reconstruct the missing/occluded retinal structures based on the learned global context rather than just local pixel interpolation.

3. Key Contributions

1. Unsupervised Diffusion Framework: Introduction of a DiffAE that integrates a context encoder with the denoising process, enabling artifact restoration without requiring paired degraded-clean training data.
2. Semantic Prior Learning: The model learns to generate high-quality retinal structures from table-top images and uses this knowledge to inpaint handheld images, preserving vascular continuity and anatomical consistency.
3. Comprehensive Validation: The approach is validated not just on pixel-level metrics, but also on downstream clinical tasks (vessel segmentation and Diabetic Retinopathy classification), proving its utility for medical diagnosis.

4. Experimental Results

The model was evaluated against state-of-the-art baselines (EdgeConnect, MISF, ShiftNet, DeepFillV2, ClarityNet) on two datasets: a Synthetic Set (DRIVE dataset with simulated artifacts) and a Test Set (real mobile fundus images from mBRSET).

Quantitative Performance

• Image Quality (Synthetic Set):
  • PSNR: The proposed method achieved 38.64 dB, outperforming the best baseline (DeepFillV2 at 37.53 dB) by 1.11 dB.
  • SSIM: Achieved 0.97, matching the best baseline.
  • Vessel Segmentation (Dice Score): Achieved 0.891, significantly higher than DeepFillV2 (0.874) and others, indicating superior preservation of fine vascular structures.
• Clinical Utility (Test Set):
  • Quality Assessment: The model produced the highest percentage of "Good" quality images (35.8%) and the lowest "Reject" rate (15.8%).
  • DR Classification Accuracy: When restored images were fed into a DR classifier, the proposed method achieved 81.17% accuracy. This is a significant improvement over the baseline accuracy of 77.11% (on un-restored images) and outperformed all other restoration methods.

Qualitative Analysis

• Visual comparisons showed that GAN-based methods often introduced unnatural textures (e.g., yellowish patches) or blurriness.
• The DiffAE approach successfully preserved the continuity of blood vessels and handled large artifacts (like flash reflections) more naturally, maintaining the anatomical structure of the optic disc and macula.

Ablation Study

• An ablation study confirmed the importance of the context encoder. Using an interpolated representation ($z_{interp}$) between the artifact image and a clean counterpart yielded higher PSNR (41.46 dB) than using only the artifact image's representation (38.98 dB), highlighting that the encoder effectively guides the restoration with clean semantic information.

5. Significance and Conclusion

This paper demonstrates that Diffusion Autoencoders are a powerful tool for unsupervised medical image restoration. By decoupling the learning of "what a healthy retina looks like" (via training on clean data) from the "inpainting process" (via inference on degraded data), the model overcomes the limitations of traditional supervised methods.

The primary significance lies in the clinical impact: the restored images are not just visually improved but are functionally superior for automated diagnosis, directly increasing the accuracy of Diabetic Retinopathy screening. This paves the way for more reliable, portable, and cost-effective ophthalmologic screening in resource-limited settings where handheld devices are prevalent. Future work suggests integrating AI-driven artifact localization to automate the mask generation process.