Enhancing Imbalanced Data Classification Using Style-Based Generative Adversarial Network-Based Data Augmentation: A Case Study of Computed Tomography Images of Brain Stroke ^†

Abstract

Stroke is a leading cause of death and disability. However, brain computed tomography image classification using machine learning and deep learning algorithms frequently suffers from a class imbalance problem, making it difficult to effectively extract deep-detailed features from instances of minority stroke lesions. In this study, we systematically implement three style-based generative adversarial network (StyleGAN)-based data augmentation approaches: StyleGAN2, StyleGAN3, and conditional StyleGAN3 to address class imbalance in brain stroke classification. Furthermore, we deploy an ensemble learning-based deep neural network to enhance the effect of those data augmentation algorithms on downstream classification tasks. Experimental results show that StyleGAN3 effectively outperforms the other two StyleGAN data augmentation approaches in terms of precision, recall, and F1-score when addressing highly imbalanced brain stroke classification. Overall, this paper demonstrates the efficacy of three StyleGAN-based data augmentation approaches in addressing imbalanced brain stroke detection.

1. Introduction

Stroke is a leading cause of death and disability worldwide, particularly in elderly populations [1]. According to the World Health Organization, approximately 15 million people worldwide experience stroke annually. The report indicates one-third die and another one-third are left with permanent disabilities [2]. Recently, computed tomography (CT) scans have been widely used in the clinical domain for brain stroke diagnosis due to their speed and ability to provide real-time brain structure information. This facilitates differentiation between hemorrhagic and ischemic strokes [3]. However, in actual medical applications, stroke examples are significantly fewer than normal samples, leading to an imbalance in class distribution. This causes traditional classification models to lean towards normal examples during training, thereby negatively affecting stroke detection accuracy.

To address the issue of class imbalance, He and Garcia [4] addressed traditional classification models that tend to overlook minority classes when class distributions are uneven. To solve the class imbalance problem, many studies have proposed sampling-level approaches, including under-sampling and over-sampling. Under-sampling involves randomly removing examples from the majority class to balance data distribution. While this approach can reduce class size disparities, it may lead to important information loss [5]. In contrast, the over-sampling method is used to directly extend the training data size by generating synthetic minority class samples.

In recent years, GAN [6] has gained great achievements in deep learning. GANs offer high-quality synthetic images through adversarial training between generators and discriminators. According to Frid-Adar et al.’s study [7], by generating synthetic images, the GAN model significantly improved imbalanced liver lesion recognition rates with convolutional neural networks. Although GANs have been widely applied to various image generation tasks, Saad et al. [8] indicated that typical GANs often suffer from mode collapse during training. To address this issue, some studies have proposed several improved GAN-based architectures. Among these architectures, Mirza and Osindero [9] proposed a conditional GAN, which incorporates label conditions into inputs to enable generators to produce images for specific label(s). However, due to classical GAN architecture limitations, CGAN will generate examples with limited diversity. To capture complex and deep features, Karras et al. [10] proposed the StyleGAN architecture. This method introduces style control mechanisms into generators and uses adaptive instance normalization (AdaIN) to adjust image style layer by layer during training. This enables enhanced diversity in generated images. However, the original StyleGAN still faces problems like feature mixing and local fake artifacts.

To address these limitations, Karras et al. [11] proposed StyleGAN2, which improves path length regularization in a generator structure, thereby effectively enhancing generated image quality. Additionally, to address overfitting issues in small-sample cases, Karras et al. [12] proposed the StyleGAN2-based adaptive discriminator augmentation (ADA) strategy, which significantly improves generation quality for small datasets. Recently, StyleGAN3 has introduced alias-free convolution to effectively avoid aliasing issues for generated images [13], ensuring virtual image correctness after rotation or translation. Furthermore, conditional StyleGAN3 can produce virtual images according to specific label(s).

In this study, we applied a StyleGAN-based model to boost imbalanced stroke detection performance by balancing class distribution. Furthermore, we adopted an ensemble learning predictive model to integrate prediction results based on multiple deep learning models for improving overall classification performance. We used evaluation metrics such as precision, recall, and F1-score to assess imbalanced data classification accuracy. Our experimental results demonstrated that among the three StyleGAN-based approaches, StyleGAN3 achieved the best classification performance in dealing with imbalanced stroke data.

2. Related Works

2.1. StyleGAN2

The core of StyleGAN2 (SG2) transforms a random latent vector $z\sim N\left(0,1\right)$ through a nonlinear mapping network to convert it into a latent space vector $w=f_{mapping}\left(z\right)$, where $f_{mapping}$ is a multi-layer perceptron. This latent vector $w$ is injected layer by layer into the synthesis layers of the generation network at different resolutions using adaptive instance normalization (AdaIN). Each layer’s feature map $x_l$ is expressed as Equation (1):

$$x_i=g_i\left(x_i-1,w_i\right)$$

(1)

where $g_i$ is the synthesis function of the ith layer, and $w_i$ is the style vector of the ith layer. This design effectively controls detailed features at different resolutions, improving image generation quality and diversity. The adversarial loss function of StyleGAN2 is Wasserstein GAN-GP (WGAN-GP), which can be defined as follows:

Discriminator loss is calculated as follows:

$${\mathcal{L}}_{\mathcal{D}}=E_{\tilde{x}\sim P_g}\left[D\left(\tilde{x}\right)\right]-E_{x\sim P_r}\left[D\left(x\right)\right]+\lambda E_{\widehat{x}\sim P_{\widehat{x}}}\left[{\left(|{\nabla }_{\widehat{x}}D\left(\widehat{x}\right){|}_2-1\right)}^2\right]$$

(2)

Generator loss is calculated as follows:

$${\mathcal{L}}_G=-E_{z\sim p\left(z\right)}\left[D\left(G\left(z\right)\right)\right]$$

(3)

It includes the following path length regularization loss.

$$L_{\mathrm{path}}=E_{w,y}{\left(|J_w\left(G\left(w\right)\right)·y{|}_2-a\right)}^2$$

(4)

Total generator loss is calculated as follows:

$${\mathcal{L}}_G^{total}={\mathcal{L}}_{\mathcal{G}}+{\lambda }_{\mathrm{path}}·L_{\mathrm{path}}$$

(5)

2.2. StyleGAN3

StyleGAN3 (SG3) builds upon StyleGAN2 with the same loss function and path regularization mechanism. Differing from StyleGAN2, it abandons the latent injection at every layer and introduces alias-free convolution. Through continuously convolutional kernels, it avoids spectral aliasing during image generation, enabling more stable and realistic capture of structural features in images. Additionally, StyleGAN3 incorporates ADA, which prevents discriminator overfitting and model collapse. The data augmentation probability $p$, is calculated using Equation (6).

$$p_{t+1}=p_t+\mathsf{\lambda }·\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{n}\left(r_t-r_{\mathrm{target}}\right)$$

(6)

where $\lambda$ is the step size, $p_t$ is the current augmentation probability, $r_t$ is the accuracy on augmented real images, and $r_{\mathrm{target}}$ is the target accuracy maintenance rate. If $r_t$ exceeds $r_{\mathrm{target}}$, $p$ increases; conversely, if it falls below this target accuracy maintenance rate, p decreases.

2.3. Conditional StyleGAN3

The conditional StyleGAN3 (CSG3) is trained with original images and conditional label learning. CSG3’s input consists of a one-hot vector c corresponding to the image’s label and a latent vector z. They are fed into the mapping network as follows:

$$w=f_{\mathrm{mapping}}\left(\left[z,\hspace{0.17em}c\right]\right)$$

(7)

The conditional style vector $w$ is injected layer by layer into the generator. Each synthesis function can be expressed as follows:

$$x_i=g_i\left(x_{i-1},\hspace{0.17em}{w}_i\right)$$

(8)

where $g_i$ is the synthesis function of the ith layer, and $w_i$ is the style vector containing conditional information for the ith layer. This design enables the generation network to control generated image content features according to target label(s). The loss functions of CSG3 are expressed as follows:

• Discriminator loss:

$$\begin{array}{c}{\mathcal{L}}_{\mathcal{D}}=E_{x,c\sim p_{\mathrm{dt}}\left(x,c\right)}\left[D\left(x,c\right)\right]-E_{z\sim p\left(z\right),c\sim p_{\mathrm{dt}}\left(c\right)}\left[D\left(G\left(z,c\right),c\right)\right]\\ \hspace{1em}\hspace{1em}+{\displaystyle \frac{2}{\gamma }}{E}_{x,c\sim p_{\mathrm{dt}}\left(x,c\right)}\left[|{\nabla }_{x}D\left(x,c\right){|}_2^2\right]\end{array}$$

(9)

• Generator loss:

$${\mathcal{L}}_{\mathcal{G}}=-E_{z\sim p\left(z\right),c\sim p_{\mathrm{dt}}\left(c\right)}\left[D\left(G\left(z,c\right),c\right)\right]+{\lambda }_{\mathrm{path}}·L_{\mathrm{path}}$$

(10)

where ${\mathit{L}}_{\mathrm{path}}$ is defined as Equation (4).

3. Methodology

The research method for generating synthetic stroke images is depicted in Figure 1.

Figure 1. Stroke data augmentation using StyleGAN.

3.1. StyleGAN Models for Stroke Image Generation

The developed method comprised two modules: StyleGAN models, such as StyleGAN2 (SG2), StyleGAN3 (SG3), and conditional StyleGAN3 (CSG3), to extend training data size, and an ensemble learning module to boost detection accuracy of imbalanced stroke. The parameter settings of the three StyleGAN models are provided in Table 1.

Table 1. StyleGAN model parameter configuration.

Parameters	StyleGAN2	StyleGAN3	Conditional StyleGAN3	Description
gpus	1	1	1	Number of GPUs
batch	8	8	16	Batch size
gamma	10	1	6	R1 regularization coefficient
aug	noaug	ada	ada	Data augmentation type
mirror	1	1	1	Horizontal flip augmentation
kimg	1000	1000	1000	Total training images (K)
metrics	fid50k	fid50k	fid50k	Frechet inception distance (FID)
cfg	paper256 [14]	stylegan3-t [15]	stylegan3-t [15]	Configure settings
cbase	-	16,384	16,384	Base channel count
dlr	-	0.0025	0.0025	Discriminator learning rate

3.2. Ensemble Learning Classification Model for Stroke Detection

After performing data augmentation, this study adopted an ensemble learning classification algorithm to improve the stability and generalization capability of the brain CT image for binary classification. To avoid overfitting, the final classification outcome was obtained from three deep learning networks: EfficientNet-B0, Xception, and InceptionV3. In each sub-network, we added batch normalization, ReLU activation, and dropout layers to learn brain stroke data characteristics with a 256 × 256 CT scan image size. All sub-networks used hyperparameter settings during training: a batch size of 16, Adam optimizer with learning rate $1\times {10}^{-4}$, and 40 epochs. We set this to a loss function as “CrossEntropyLoss” indicator. In the ensemble learning stage, we used soft voting to fuse predictions from all sub-models. For each sub-model $M_i$, the predicted result of input x denoted as label $k$, as seen in Equation (11). We averaged the predictive outputs across all sub-networks to obtain average probability for each label, as seen in Equation (12). We selected the label with highest average probability as the final predictive result, as seen in Equation (13).

$$p_k^{\left(i\right)}=P\left(y=k∣x;M_i\right)$$

(11)

$$\overline{p_k}={\displaystyle \frac{1}{N}}\sum _{i=1}^N{p}_k^{\left(i\right)}$$

(12)

$$\widehat{y}=\arg \underset{k}{\max }\overline{p_k}$$

(13)

4. Experiment

All experiments were conducted on a computer equipped with an Intel(R) Core(TM) i7-14700KF CPU (27-processor) and Nvidia GeForce RTX 4090 graphics (24 gigabytes random access memory) running Ubuntu 24.04 LTS. The deep learning model was built using PyTorch. We used the PyTorch container version 25.03-py3 from NVIDIA NGC, which includes PyTorch version (2.7.0a0 + 7c8ec84dab) with compute unified device architecture 12.8.1 support.

4.1. Case Description

The brain stroke CT image dataset is publicly available at https://www.kaggle.com/datasets/iashiqul/brain-stroke-prediction-ct-scan-image-dataset%20 (accessed on 28 January 2026). We modified it to establish an imbalanced scenario. To avoid data leakage, we split the original dataset into the training sample set which included 1080 normal images and 108 stroke images (i.e., an imbalance ratio of ${\displaystyle \frac{N_{\mathrm{normal}}}{N_{\mathrm{stroke}}}}=10$), the validation sample set which included 307 normal images and 130 stroke images, and the testing sample set which included 157 normal images and 64 stroke images.

4.2. Experiment Setting

In this study, we compared three types of StyleGAN data augmentation approaches for brain CT image classification under an imbalance ratio of IR = 10. The three approaches included “Geometric augmentation (GEO) + SG2”, “GEO + SG3”, and “GEO + CSG3”. We comprehensively analyzed their effects for stroke detection. SG2 employed a generative adversarial network to augment minority class stroke images. SG3 utilized an alias-free convolution mechanism to enhance the generation stability of images. CSG3 further incorporated conditional label variables to control the label of artificial images. By defining the stroke class as a positive class and the normal class as a negative class, we calculated true positive (TP), false positive (FP), false negative (FN), and true negative (TN). Four evaluation metrics were used to assess stroke detection performance (Equations (14)–(17)).

$$\mathrm{A}\mathrm{c}\mathrm{c}\mathrm{u}\mathrm{r}\mathrm{a}\mathrm{c}\mathrm{y}={\displaystyle \frac{\mathrm{T}\mathrm{P}+\mathrm{T}\mathrm{N}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}+\mathrm{T}\mathrm{N}+\mathrm{F}\mathrm{N}}}$$

(14)

$$\mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}={\displaystyle \frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{N}}}$$

(15)

$$\mathrm{Precision}={\displaystyle \frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}}}$$

(16)

$$\mathrm{F}1\text{-}\mathrm{s}\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{e}=2\times {\displaystyle \frac{\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}\mathrm{n}\times \mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}}{\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}\mathrm{n}+\mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}}}$$

(17)

4.3. Results

We averaged predictive results over 30 independent experiments. The results of average and standard deviation (SD) are listed in Table 2.

Table 2. Results of average and standard deviation.

Method	StyleGAN2 (SG2)		StyleGAN3 (SG3)		Conditional StyleGAN3 (CSG3)
Performance	Average	SD	Average	SD	Average	SD
Accuracy	0.829	0.021	0.852	0.016	0.837	0.015
Normal_Recall	0.998	0.003	0.998	0.004	0.999	0.003
Normal_Precision	0.807	0.019	0.830	0.016	0.815	0.015
Normal_F1-score	0.892	0.012	0.906	0.009	0.897	0.009
Stroke_Recall	0.414	0.072	0.496	0.057	0.441	0.056
Stroke_Precision	0.988	0.021	0.991	0.017	0.993	0.016
Stroke_F1-score	0.580	0.073	0.659	0.051	0.609	0.052
Macro_F1-score	0.736	0.042	0.782	0.030	0.753	0.030
Weighted_F1-score	0.802	0.029	0.834	0.021	0.814	0.021

Among the three StyleGAN-based approaches, SG3 achieved the best classification accuracy score of 0.852 ± 0.01. CSG3 presented a better classification accuracy score of 0.837 ± 0.01, and SG2 achieved a score of 0.829 ± 0.02, as shown in Table 2 and Figure 2. For the stroke label, SG3 demonstrated a superior score of 0.496 ± 0.05 compared to the SG2 score of 0.414 ± 0.07 and the CSG3 score of 0.441 ± 0.05 in terms of recall. The F1-score had been improved from 0.580 ± 0.07 by SG2 to 0.659 ± 0.05 by SG3. This indicated that SG3’s alias-free architecture effectively enhanced performance for the difficult-to-classify minority stroke class. For the normal class, precision and recall remained high across all three StyleGAN-based approaches. This demonstrated that the data augmentation strategies did not compromise the classification performance of the majority class. From Table 2, SG3 outperformed the other two approaches in terms of Macro_f1-score at 0.782 ± 0.03 and Weighted_f1-score at 0.834 ± 0.02. This indicated SG3’s ability to deal with imbalance stroke class cases.

Figure 2. Classification accuracy using StyleGAN2, StyleGAN3, and conditional StyleGAN3.

4.4. Comparison for FID Results

We selected FID to measure image generation quality using Equation (18).

$$\mathrm{FID}={\parallel {\mathrm{\mu }}_r-{\mathrm{\mu }}_g\parallel }^2+\mathrm{Tr}\left({\mathsf{\Sigma }}_r+{\mathsf{\Sigma }}_g-2{\left({\mathsf{\Sigma }}_r{\mathsf{\Sigma }}_g\right)}^{1/2}\right)$$

(18)

where $\parallel ·\parallel$ represents the Euclidean norm, $\mathrm{Tr}$() denotes the matrix trace (sum of diagonal elements), and ${\left({\mathsf{\Sigma }}_{\mathit{r}}{\mathsf{\Sigma }}_{\mathit{g}}\right)}^{\mathbf{1}\mathbf{/}\mathbf{2}}$ represents the matrix square root of ${\mathsf{\Sigma }}_{\mathit{r}}{\mathsf{\Sigma }}_{\mathit{g}}$. ${\mathsf{\mu }}_{\mathit{r}}$ and ${\mathsf{\mu }}_{\mathit{g}}$ represent the average vectors of real and generated images, respectively. ${\mathsf{\Sigma }}_{\mathit{r}}$ and ${\mathsf{\Sigma }}_{\mathit{g}}$ are the corresponding covariance matrices. FID was used to measure the distance between generated and real image distributions. Lower FID scores indicated that generated images were closer to real images. As shown in Figure 3, CSG3 (in the red line) achieved the lowest FID.

Figure 3. FID results for StyleGAN2, StyleGAN3, and conditional StyleGAN3.

Here, SG2 achieved the lowest FID score of 65.91 with 400 k images. SG3 reached FID’s score at 101.04 with 1000 k images. CSG3 achieved the best FID at 63.62 with 1000 k images. According to these results, we suggest that incorporating the majority class features can effectively help improve the identification of minority class instances in brain structures.

5. Conclusions

To address the class imbalance problem in stroke classification, we compared StyleGAN-based over-sampling models: StyleGAN2, StyleGAN3, and conditional StyleGAN3. We selected the best generative model among them. Although the SG3 model achieved the best classification performance with an ensemble learning model, on generation quality and model generalization, StyleGAN3 showed overfitting. This deteriorated the SG3 model’s effectiveness under certain parameter settings. In this study, we evaluated classification performance for a highly imbalanced ratio value of 10 between stroke and normal classes. Experimental results validated that StyleGAN3 outperformed both StyleGAN2 and conditional StyleGAN3 methods in terms of precision, recall, and F1-score metrics. Two future research directions include expanding SG3’s validation across different medical datasets and using Keras tuner to optimize hyperparameters of the StyleGAN model.