DiffuPT: Class Imbalance Mitigation for Glaucoma Detection via Diffusion Based Generation and Model Pretraining

TL;DR

DiffuPT addresses class-imbalance in glaucoma detection by synthesizing data via diffusion models and pretraining classifiers on a balanced synthetic dataset before fine-tuning on real data. It introduces GlaucomaEgy, the largest national dataset for glaucoma detection, and demonstrates that diffusion-based augmentation outperforms GAN-based augmentation and other imbalance mitigation strategies. The results show harmonic mean improvements from $89.09\%$ to $92.59\%$ on GlaucomaEgy and consistent gains on AIROGS, indicating improved robustness and generalization. DiffuPT reduces variance in latent embeddings and provides a practical, annotation-free approach to mitigate class imbalance in medical imaging.

Abstract

Glaucoma is a progressive optic neuropathy characterized by structural damage to the optic nerve head and functional changes in the visual field. Detecting glaucoma early is crucial to preventing loss of eyesight. However, medical datasets often suffer from class imbalances, making detection more difficult for deep-learning algorithms. We use a generative-based framework to enhance glaucoma diagnosis, specifically addressing class imbalance through synthetic data generation. In addition, we collected the largest national dataset for glaucoma detection to support our study. The imbalance between normal and glaucomatous cases leads to performance degradation of classifier models. By combining our proposed framework leveraging diffusion models with a pretraining approach, we created a more robust classifier training process. This training process results in a better-performing classifier. The proposed approach shows promising results in improving the harmonic mean (sensitivity and specificity) and AUC for the roc for the glaucoma classifier. We report an improvement in the harmonic mean metric from 89.09% to 92.59% on the test set of our national dataset. We examine our method against other methods to overcome imbalance through extensive experiments. We report similar improvements on the AIROGS dataset. This study highlights that diffusion-based generation can be of great importance in tackling class imbalances in medical datasets to improve diagnostic performance.

Figures (10)

• Figure 1: Analysis of the embeddings on different training techniques. Normal Training shows high variance and low bias (less overlap between glaucoma and non-glaucoma classes. Pre-training alone shows much more stability with lower variance but with higher overlap. The fine-tuned model shows the best of both worlds, a robust classifier with lower variance and lower bias.
• Figure 2: Samples of the collected dataset. The two images on the left are from grade 0 (non-glaucomatous) samples. Both middle images are from grade 1 (glaucoma suspect). The two right images are from grade 2 (glaucoma).
• Figure 3: Framework for generating data using LDM. We show how to sample a generated dataset using LDM. Our method involves a filtering stage based on the classification output of the baseline classifier.
• Figure 4: Training procedure used for obtaining the final classifier.
• Figure 5: Samples from different generation methods. The left images show nonglaucoma samples from the dataset. Then we show the results of conditional generation from the DDPM and LDM. The bounded images are from GAN-based image-to-image translation from real nonglaucoma images to fake glaucoma using various methods.