Comprehensive Data Augmentation Approach Using WGAN-GP and UMAP for Enhancing Alzheimer’s Disease Diagnosis

Abstract

In this study, the Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP) was used to improve the diagnosis of Alzheimer’s disease using medical imaging and the Alzheimer’s disease image dataset across four diagnostic classes. The WGAN-GP was employed for data augmentation. The original dataset, the augmented dataset and the combined data were mapped using Uniform Manifold Approximation and Projection (UMAP) in both a 2D and 3D space. The same combined interaction network analysis was then performed on the test data. The results showed that, for the test accuracy, the score was 30.46% for the original dataset (unbalanced), whereas for the WGAN-GP augmented dataset (balanced), it improved to 56.84%, indicating that the WGAN-GP augmentation can effectively address the unbalanced problem.

1. Introduction

Alzheimer’s disease (AD) is a progressive, neurodegenerative disorder that is particularly common in the elderly, and its diagnosis remains a major challenge [1,2,3,4,5,6]. While early and accurate diagnosis is essential for treatment and care planning, the medical imaging data used to diagnose AD often face the problem of diagnostic class imbalance. This imbalance is caused by a markedly different ratio of normal to abnormal samples in the dataset and is known to adversely affect the performance of machine learning models [7]. In particular, disease samples belonging to minority diagnostic classes may not be adequately trained, leading to reduced diagnostic accuracy. To solve this imbalance problem and improve the accuracy of AD diagnosis, we propose a data augmentation method using the Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP). WGAN-GP is a data augmentation method that has more stable learning than conventional GANs and can provide high-quality generated data, which can be applied in the field of medical imaging.

MRI morphometry has long been used to analyze the brain structure in patients with AD. Studies like that of Matsuda [8] have highlighted the importance of brain atrophy measurements, especially in the hippocampus, which is one of the earliest affected regions in AD. Such morphometric analysis is essential in staging the disease and tracking its progression. However, there are limitations in conventional imaging analysis, particularly when faced with variability in imaging data quality and segmentation techniques. For instance, Tubi et al. [9] have demonstrated the influence of segmentation algorithms on the accuracy of white matter hyperintensity (WMH) analysis, a crucial biomarker for cognitive decline associated with AD. Variability in data preprocessing can significantly impact the diagnostic outcomes, particularly when deep learning models are applied. To address these challenges, researchers have explored various advanced machine learning methods. So et al. [10] proposed a deep learning approach using texture features from MRI scans, focusing on hippocampal changes. This research has shown promise in automating the diagnosis process by leveraging imaging features that are often overlooked in traditional visual assessments. Similarly, Kamal et al. [11] employed machine learning frameworks to analyze brain MRI scans, with the aim of enhancing the detection of AD. However, the major issue remains that these models can suffer from data imbalance, as disease samples (especially at early stages) are underrepresented in datasets. This imbalance leads to less accurate models, as the underrepresented minority diagnostic class is not adequately trained. Several studies have proposed data augmentation methods to overcome these limitations. For example, Bhateja et al. [12] introduced a multispectral image fusion technique to enhance the quality of medical imaging data. The combination of multiple imaging modalities, such as MRI and PET, has proven to be a valuable approach for improving diagnostic accuracy [13]. Nevertheless, data augmentation using traditional techniques is often insufficient to fully solve the problem of imbalance. Generative adversarial networks (GANs) have emerged as a powerful tool for generating synthetic data to augment existing datasets.

In recent years, the development of the Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP) has shown potential in generating high-quality synthetic medical imaging data. The WGAN-GP has been specifically designed to stabilize the learning process of GANs and can produce realistic, high-resolution images, making it particularly suitable for medical imaging applications [14]. By integrating the WGAN-GP with neuroimaging data, the imbalance problem in AD diagnosis can be mitigated, allowing for more robust training of machine learning models. For instance, Sajjad et al. [15] utilized a deep convolutional generative adversarial network (DCGAN) for synthetic data augmentation in PET imaging for AD classification, demonstrating significant improvements in diagnostic performance. Similarly, other studies have applied advanced generative techniques to augment MRI data and enhance the training of machine learning models in AD research [16,17,18,19,20]. It is also being applied to algorithms for early detection [20,21] and to research into disease screening methods in combination with other signals such as Electroencephalograms (EEGs) [22,23].

In this paper, a combined dataset of generated and original data was mapped into a 2D and 3D space using Uniform Manifold Approximation and Projection (UMAP) for visual analysis.

2. Methods

2.1. Dataset

The Alzheimer’s MRI dataset utilized in this study was sourced from Kaggle (https://www.kaggle.com/datasets/lukechugh/best-alzheimer-mri-dataset-99-accuracy/data accessed on 13 September 2024). This dataset comprises a collection of MRI images categorized into four distinct diagnostic classes relevant to Alzheimer’s disease diagnosis: Non-Demented, Very Mildly Demented, Mildly Demented and Moderately Demented. The original dataset presented an imbalance among the diagnostic classes, which could negatively impact the accuracy of classification models. Each category included 100, 70, 28 and 2 patients, respectively, and each patient’s brain was scanned on a 32-slice horizontal-axis MRI. MRI images were acquired using a 1.5 Tesla MRI scanner in a T1-weighted sequence. The images have a resolution of 128 × 128 pixels and are in “.jpg” format. All images were preprocessed to remove the skull.

2.2. Data Augmentation Using WGAN-GP

To address the issue of diagnostic class imbalance, we employed the Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP). The WGAN-GP is a variant of GANs designed to generate high-quality synthetic data while maintaining stability during training. The model was trained on the original MRI images to generate synthetic samples for the underrepresented diagnostic classes. These synthetic data were then combined with the original dataset to create a balanced augmented dataset. The WGAN-GP was implemented using Python and the TensorFlow library, with specific hyperparameters tuned to optimize the quality of generated images.

To perform this process, we first defined the paths and labels of the dataset to be used for Alzheimer’s disease diagnosis. The dataset contains both the original actual images and the images generated by WGAN-GP. Labels were set to distinguish between the actual data and the generated data. Next, feature extraction was performed. As image loading and preprocessing, images in the dataset were loaded and converted to grayscale. After conversion, the images were changed to the ubyte format for easy handling by the algorithm. The GLCM (Gray Level Co-occurrence Matrix) was then computed. The GLCM was computed for the grayscale images, and contrast and correlation were extracted as texture features. For each image in the original dataset, the aforementioned feature extraction was performed. The features were labeled as “non-GAN” (real data). Features were extracted for the combined dataset of the original dataset and the dataset generated by the GAN in the same way; the data generated by the GAN were labeled as “generated by GAN (generated data)”. The extracted features for all images were compiled into a list, which was converted into a NumPy array and scaled using StandardScaler to normalize the extracted features. This homogenizes the value range of the features and improves the performance of subsequent dimensionality reduction methods. After dimensionality reduction was performed on the scaled data using PCA and t-SNE, scatter plots of the original dataset, the GAN-generated dataset and the combined dataset were plotted using 2D and 3D dimensionality reduction with UMAP. This compares the differences in the visual distribution of each dataset. Finally, the scaled feature data were output and saved in a format that can be used for the next processing steps and algorithms (Figure 1).

Learning in WGAN-GP is characterized by the fact that it does not use the format commonly used in identification, in which the correct answer label is paired with the expected result [24,25], but rather defines its own loss function, rather than using an already defined function such as binary_cross_entropy. The procedure for defining a unique loss function and passing it to the optimizer for training is as follows: (1) create a model, (2) define the loss function, (3) instantiate the optimizer, specify the weights to be trained in the updates method, and (4) create the input, output and instantiated. The code for WGAN-GP is available as an open source resource (keras-contrib/examples/improved_wgan.py accessed on 13 September 2024) and has recently been applied in the field of medical image processing [26,27].

2.3. Dimensionality Reduction and Visualization with UMAP

After data augmentation, we applied Uniform Manifold Approximation and Projection (UMAP) for dimensionality reduction and visualization. UMAP is a non-linear dimensionality reduction technique that preserves the global structure of high-dimensional data when projected into lower-dimensional spaces. We used UMAP to project both the original and augmented datasets into 2D and 3D spaces. This mapping allowed us to visually assess the distribution and separability of the different diagnostic classes in both the original and augmented datasets.

2.4. Classification and Evaluation

To evaluate the effectiveness of WGAN-GP augmentation, we employed a classification model based on a Random Forest classifier, which was chosen for its ability to handle imbalanced datasets. The model was trained using both the original (unbalanced) dataset and the WGAN-GP augmented (balanced) dataset. We utilized supervised learning with cross-validation to ensure the robustness of the model, splitting the data into training and validation sets. Data preprocessing was a critical step before training the model. Two different preprocessing pipelines were compared: (1) standard scaling, where features were normalized to have zero mean and unit variance, and (2) min–max scaling, which rescaled features to a [0, 1] range. The classification model was trained on the datasets prepared using both methods, and their impact on model performance was evaluated. The model’s performance was tested on a separate test set, where accuracy, precision, recall, and F1-scores were computed for both the original and augmented datasets. Additionally, we used 2D UMAP projections to visualize the diagnostic class separability and assess classification accuracy across different diagnostic class distributions. The improvement in classification metrics from the original to the augmented dataset was used as a key indicator of the success of the augmentation process.

The performance before and after data enhancement was objectively evaluated using the Learned Perceptual Image Patch Similarity (LPIPS) metric with a trained network. Deep learning is a machine learning technique that has attracted attention for its extremely high accuracy in image classification tasks. LPIPS, a typical measure of distance between images, is calculated based on the deep features output by the convolutional layer of a trained image classification network, such as AlexNet or VGG, for an input image [28]. LPIPS, which uses feature extraction based on deep learning, has been shown to be more accurate than conventional methods that focus only on pixel brightness and contrast. In recent years, LPIPS has shown high versatility in quantifying the color patterns of living organisms [29,30,31] and has become popular for various comparative analyses because it can be equally accurate with a single simple and objective criterion based on deep learning.

2.5. Software and Tools

The entire process, from data augmentation to dimensionality reduction and classification, was implemented using Python. Key libraries included TensorFlow for the WGAN-GP implementation, UMAP-learn for dimensionality reduction and Scikit-learn for classification and evaluation.

3. Results

The results of the 2D UMAP and 3D UMAP analyses are shown in Figure 2 and Figure 3. The classification heatmaps are shown in Figure 4. 2D and 3D UMAP analysis results show that the UMAP projection improves diagnostic class separation and clustering in the augmented and combined datasets compared to the original datasets, reflecting the benefits of data augmentation by the WGAN-GP. The results reflect that, in particular, the improved separation in 3D space further supports the utility of the WGAN-GP to address diagnostic class imbalances and improve diagnostic accuracy. The classification accuracy of the original unbalanced dataset was 30.46%, whereas the classification accuracy of the diseases after applying the WGAN-GP for data augmentation was 56.84%, a significant improvement. The performance before and after data augmentation was evaluated using LPIPS, with an average LPIPS score of (0, 0.205).

4. Discussion

There are some studies on disease visualization by machine learning for patients with chronic kidney disease and on screening for obesity [32,33]. However, there are few examples of studies on Alzheimer’s classification, and only a few studies on brain circuitry in health and disease from linear and non-linear approaches using PCA, multidimensional scaling (MDS), iso-mapping, local linear embedding (LLE), Laplacian eigenmaps (LEM), t-SNE neural manifolds of dynamics [34] and proposing a framework for removing confounding effects from distance-based dimensionality reduction methods [35]. In this work, we have already shown that similarity- and graph structure-preserving dimensionality reduction tools such as t-SNE and UMAP can capture complex biological patterns in high-dimensional data. We use the WGAN-GP for the purpose of enhancing pathological diagnosis, which is novel in this respect.

Overall, the results confirm that the WGAN-GP augmentation leads to significant improvements in the classification performance of Alzheimer’s disease diagnosis models. The increased accuracy from 30.46% to 56.84% in the 2D UMAP analysis demonstrates the potential of this approach to mitigate the issues associated with unbalanced datasets and to enhance the reliability of medical imaging diagnoses. WGAN-GP successfully addressed the problem of diagnostic class imbalance within the Alzheimer’s disease image dataset and, when combined with the UMAP technique, significantly improved diagnostic accuracy. The most significant improvement was observed in 2D UMAP mapping, and this significant improvement highlights the importance of data enhancement in medical imaging, especially when dealing with diagnostic class imbalances that are common in real-world datasets. The effectiveness of WGAN-GP in generating high-quality synthetic data may have contributed to enhanced discrimination of Alzheimer’s disease among the four diagnostic classes, as observed in the UMAP visualization. The expanded dataset provided a more balanced representation of the diagnostic classes and allowed the model to better capture underlying patterns in the data. In addition, the dimensionality reduction provided by UMAP played an important role in visualizing the structure of the dataset in both the2D and 3D space; the combination of the WGAN-GP and UMAP not only improved diagnostic accuracy, but also allowed for a clearer understanding of the complex relationships among the different diagnostic classes in the dataset. The results of this study were as follows. This approach may be particularly useful when exploring other medical imaging datasets that suffer from similar problems of diagnostic class imbalance and high dimensionality. The method has the potential for broader application in medical imaging and other domains where data imbalance and high dimensionality are important issues.

However, while the results are promising, further research is needed to improve the augmentation process and to assess its generality across different datasets and diagnostic tasks. In addition, integrating the WGAN-GP with other data augmentation techniques to further increase the robustness of the model is a future challenge. Some limitations of this work are that only UMAP is used for visualization and not t-SNE, a method that preserves both global and local structures well. t-SNE has execution time issues, and there have been proposals, e.g., LargeVis [36,37,38], to solve them, but the execution time increases exponentially as the embedding dimension increases. UMAP is adopted in this study because its execution time is almost constant regardless of the number of embedded dimensions, but studies comparing it with t-SNE have been increasing in recent years [39,40,41], and this is a future issue. Furthermore, the role of genetic factors in the development of Alzheimer’s disease as well as other neurodegenerative diseases such as Parkinson’s disease will need to be examined in the future. Half of Parkinson’s patients develop dementia, and many studies suggest that Alzheimer’s disease is hereditary (e.g., a small proportion of all Alzheimer’s disease is familial, and in familial Alzheimer’s disease, if one of the parents has a genetic mutation, the odds of the child developing the disease are high) [42,43,44,45,46,47,48,49,50,51]. Applications for early screening for Parkinson’s disease would be medically useful. Incorporating state-of-the-art image enhancement methods and applications related to restoration to eliminate the negative effects of noise, such as single image norm normalization [52] and multiobjective haze removal methods [53], is a concurrent challenge.

5. Conclusions

In this study, we presented a comprehensive data augmentation approach using the Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP) combined with Uniform Manifold Approximation and Projection (UMAP) to enhance the diagnostic accuracy of Alzheimer’s disease (AD) through medical imaging. The main challenge addressed in this research was the inherent diagnostic class imbalance in the Alzheimer’s disease dataset, which significantly affects the performance of machine learning models in medical diagnosis.

Our results demonstrated that the WGAN-GP is an effective method for generating high-quality synthetic data, which, when combined with the original data, helps to mitigate the diagnostic class imbalance problem. Specifically, the use of WGAN-GP to augment the dataset led to a significant improvement in test accuracy, increasing from 30.46% with the original unbalanced dataset to 56.84% with the augmented dataset. This indicates that the synthetic data generated by WGAN-GP not only increased the quantity of data but also preserved critical features necessary for accurate diagnosis.