Diffusion Autoencoder for Unsupervised Artifact Restoration in Handheld Fundus Images

Abstract: The advent of handheld fundus imaging devices has made ophthalmologic diagnosis and disease screening more accessible, efficient, and cost-effective. However, images captured from these setups often suffer from artifacts such as flash reflections, exposure variations, and motion-induced blur, which degrade image quality and hinder downstream analysis. While generative models have been effective in image restoration, most depend on paired supervision or predefined artifact structures, making them less adaptable to unstructured degradations commonly observed in handheld fundus images. To address this, we propose an unsupervised diffusion autoencoder that integrates a context encoder with the denoising process to learn semantically meaningful representations for artifact restoration. The model is trained only on high-quality table-top fundus images and infers to restore artifact-affected handheld acquisitions. We validate the restorations through quantitative and qualitative evaluations, and have shown that diagnostic accuracy increases to 81.17% on an unseen dataset and multiple artifact conditions

• The paper presents an unsupervised diffusion autoencoder that leverages context-aware denoising for artifact restoration in handheld fundus images.
• It integrates a forward diffusion and reverse denoising process with a UNet to extract semantically rich latent representations using only clean table-top images.
• Experimental evaluations demonstrate improved image quality, segmentation accuracy, and diabetic retinopathy classification without reliance on paired artifact-clean supervision.

Diffusion Autoencoder for Unsupervised Artifact Restoration in Handheld Fundus Images

Introduction

This work addresses the significant challenge of restoring ophthalmic fundus images acquired using handheld and mobile devices, which are prone to a range of image artifacts, including flash reflections, exposure inconsistencies, and motion blurring. These degradations impede the efficacy of downstream automated analysis tasks crucial for disease screening and diagnosis, particularly diabetic retinopathy (DR). Unlike table-top fundus cameras, which capture high-quality images amenable to classical computer vision models, handheld acquisitions exhibit unstructured and spatially diverse artifacts that classical GAN-based or CNN-based inpainting and restoration methods struggle to overcome. Existing supervised methods are limited by their reliance on paired training data or strict artifact masks, rendering them less effective for unconstrained, real-world degradations.

Methodology

The proposed solution is an unsupervised diffusion autoencoder (DiffAE) integrating a context encoder with the denoising UNet, facilitating the extraction of semantically meaningful latent representations for artifact restoration. Training is performed exclusively on clean, high-quality table-top fundus images, with no access to artifact-affected pairs or handcrafted mask supervision.

Diffusion Autoencoder Framework

The encoder maps each clean image $x_0$ to a global latent space $z = \mathrm{Enco}(x_0)$ capturing retinal structural properties. The forward diffusion process adds noise to $x_0$, generating noisy samples over $T$ steps, while the reverse process reconstructs the clean image conditioned on both the intermediate noisy image and the latent vector $z$. The UNet denoiser is trained to minimize the mean squared error between predicted and true clean images, ensuring accurate and contextually rich latent embeddings. This setup allows the model to generalize to unseen artifact types by inferring with the learned latent priors.

Mask-Conditioned Restoration

Artifact restoration during inference employs artifact masks $M$, generated through adaptive image processing or synthesized based on real artifact regions. Restoration is formulated as denoising-time inpainting, where each denoising step reconstructs artifact-masked regions with context from both image and latent priors. This integrates the restoration process with the underlying anatomical and vascular context, crucial for regulatory downstream applications.

Experimental Evaluation

Experiments utilize the EyePACS dataset for training (clean images only), the mBRSET mobile fundus dataset for real artifact evaluation, and the DRIVE dataset for controlled synthetic artifact introduction. Baseline models include state-of-the-art GANs (EdgeConnect, DeepFillV2), CNN-based (MISF, ShiftNet), and a previous diffusion-based approach (ClarityNet), all retrained in a comparable unsupervised setting.

Quantitative Results

• Image Quality (Synthetic Set): The proposed method surpasses DeepFillV2 by +1.11 dB in PSNR, with comparable SSIM, demonstrating superior restoration fidelity.
• Vessel Segmentation: Dice score for vessel segmentation is highest for DiffAE (0.891), outperforming all baselines, indicating optimal restoration of fine vascular structures necessary for clinical validity.
• Quality Assessment (Handheld Test Set): The Good/Usable/Reject rate is superior for the proposed model, with 35.8% "Good" and 52.6% "Usable", reflecting improved subjective and automated grading.
• DR Classification: The restored images yield a DR classification accuracy of 81.17%, improving upon the uncorrected baseline (77.11%) and all other methods.

Notably, this performance is achieved entirely without paired artifact-clean supervision, suggesting robust generalization enabled by context-modulated autoencoding.

Qualitative Analysis

The visual superiority of the method is demonstrated in handling both synthetic and genuine artifacts. Competing GAN-based methods often introduce unrealistic structures or fail to reconstruct complex anatomical regions, while the proposed method produces anatomically consistent, high-fidelity reconstructions for a diverse array of degradation patterns.

Ablation

Ablation studies on latent representations reveal that improved context encoding, via interpolation between representations for artifact-laden and clean images, substantially boosts PSNR to 41.46 dB, reinforcing the pivotal role of semantically informed latent priors in the restoration pipeline.

Practical and Theoretical Implications

This work demonstrates that diffusion-based autoencoding, augmented with context-aware representation learning, enables robust unsupervised restoration for medical image domains plagued by unpredictable and large-scale artifacts. By decoupling reconstruction from explicit artifact supervision, the proposed framework dramatically reduces annotation bottlenecks, facilitating scalable deployment of handheld fundus imaging in low-resource and point-of-care settings. Practically, improvements in artifact restoration translate directly to enhanced disease screening accuracy, supporting regulatory acceptance for clinical use.

Theoretically, these results suggest new directions for integrating representation learning with diffusion processes in unsupervised restoration, particularly for biomedical domains where paired data is scarce or infeasible. The use of contextual priors at each denoising step could extend beyond artifact inpainting to other ambiguous inverse problems, such as semi-supervised segmentation or domain adaptation in biomedical imaging.

Conclusion

This study introduces an unsupervised diffusion autoencoder with a context encoder for artifact restoration in handheld fundus images, demonstrating state-of-the-art quantitative and qualitative results across image quality, segmentation fidelity, and diagnostic utility. The integration of semantic context into the diffusion denoising process enables robust unsupervised restoration without explicit artifact supervision. Future work could enhance artifact mask generation using AI-driven localization and extend this approach to other medical imaging modalities suffering from out-of-distribution artifact profiles.