A Novel Approach for Sperm Morphology Classification Using Generative Adversarial Network and Attention Network

Sperm morphology refers to the size and shape of sperm cells. The structure of sperm is very important for its function and fertilization [1]. It consists of three major parts: the Head, the Midpiece, and the Tail. The head of the sperm contains genetic material. The midpiece contains mitochondria for energy [2]. Likewise, the tail portion is used for swimming. Based on the survey, 40% of men have fertility issues. The major cause of fertility problems is increased scrotal temperature and oxidative stress. It will damage developing sperm and lead to morphological problems. In addition, chronic prostatitis, epididymitis, and sexually transmitted infections can cause inflammation. This inflammation damages sperm production and maturation processes. The proper sperm morphology analysis is used to identify these issues and supports fertility treatments and interventions.

The traditional method of sperm morphology analysis is labour-intensive and time-consuming. Also, it is dependent on the expertise of a trained analyst. However, recently, AI-based approaches have provided automated solutions to analyze sperm morphology [3]. In AI, the models can be trained on labeled datasets of sperm images to classify sperm into categories, such as "normal" and "abnormal." These classifications are based on features like head shape, midpiece, and tail structure. There are two categories of AI models that can be used for classification. The machine learning (ML) models like support vector machine (SVM), decision tree, and CatBoost, etc. [4]. Likewise, the deep learning (DL) models of Convolutional Neural Networks (CNNs), ResNet, and U-Net architectures are used for image classification [5].

Despite significant advancements in AI models, there are several critical challenges that remain, which affect the performance of the models. Existing sperm image datasets are limited in size and may not adequately represent the full spectrum of sperm morphology. This dataset suffers from class imbalances. Most of the models lack fine-grained feature extraction. They struggled to capture complex and multiscale features inherent in sperm images.

In this work, a three-fold model is proposed to address the above challenges. The Generative Adversarial Network (GAN) model is proposed for data augmentation. Then, new models are proposed for feature extraction and classification. The proposed model is validated for multiple datasets with Explainable AI (XAI) techniques for model interpretations.

In this work, a multi-stage model is proposed for sperm morphology classification. The overall workflow is shown in Figure 1. Initially, for data augmentation, the Modified Generative Adversarial Network (M-GAN) is proposed for realistic sperm image generation. Then, a Biomorphological Attention Feature Network (BMAF-Net) is proposed for effective feature extraction. For better feature optimisation, the Modified Waterwheel Plant Algorithm (M-WWPA) is proposed. Finally, the Gradient Boosted Random Forest (GBRF) is applied for classification. This multi-component model is proposed to increase the accuracy and reliability of automated sperm morphology analysis in medical settings.

Figure 1. Workflow of proposed classification

3.1 Proposed Generative Adversarial Network for sperm image generation

In this work, a modified GAN architecture is proposed to generate images. This GAN model incorporates a morphology-based conditioning mechanism. In this GAN, the generator produces images specific to each of the four sperm morphology classes. The generator considers a latent vector as well as a morphology class index to generate realistic sperm images conditioned on specific sperm types. It involves conditional generation, embedding layers, and attention mechanisms to improve high-fidelity generation.

The proposed model consists of two primary components:

1. Generator: This component generates sperm images conditioned on the input morphology label.
2. Discriminator: This component classifies images as real or fake and also determines the morphology class of the input image.

3.1.1 Generator architecture

The generator uses a latent vector $Z$ and a morphology embedding $E$ as input. The morphology embedding encodes the four distinct sperm morphology types: normal, abnormal 1, abnormal 2, and defective. By conditioning the generator on this morphology label, it can generate specific sperm images corresponding to each class. The generator starts by embedding the morphology index $m \in\{0,1,2,3\}$ into a continuous space using an embedding layer. It can be expressed as a morphology vector $E_m$ as follows:

$E_m=$ Embedding $(m)$ where $m \in\{0,1,2,3\}$             (1)

The latent vector $z$ and the morphology embedding $E_m$ are concatenated to form an input vector for the first fully connected layer of the generator.

$x_1=\operatorname{ReLU}\left(W_f \cdot\left[z, E_m\right]+b_f\right)$           (2)

The feature vector $x_1$ is passed through a series of deconvolutional layers. It is used to progressively upscale the feature map to generate a high-resolution sperm image. In addition, attention mechanisms like non-local attention and spatial attention are applied for the model to focus on crucial regions like the sperm head and tail. The non-local attention mechanism captures long-range dependencies. The generated image captures the global structure of the sperm image.

$A=\operatorname{Softmax}\left(Q K^T\right) V$           (3)

where, Q, K, and V represent the query, key, and value matrices, respectively. The spatial attention mechanism focuses on specific regions in the image, such as the sperm head and tail.

$A_{\text {spatial }}=\operatorname{Sigmoid}(\operatorname{Conv}(x))$           (4)

where, Conv(x) is a convolutional layer that generates the attention map.

3.1.2 Discriminator architecture

Here, the discriminator is responsible for two tasks: determining whether an image is real or fake and predicting the morphology class of the input image. The multiple convolutional layers are used to extract features from the input image, with a fully connected layer. The probabilities of the architecture are the Real/Fake probability and the morphology classification. For Real/Fake Classification:

$p_{\text {real }}=\operatorname{Sigmoid}\left(\mathrm{FC}_1(x)\right)$           (5)

For Morphology Classification:

$p_{\text {morphology }}=\operatorname{Softmax}\left(\mathrm{FC}_2(x)\right)$           (6)

where, $\mathrm{FC}_1$ and $\mathrm{FC}_2$ are fully connected layers in the discriminator.

3.1.3 Loss functions

The loss functions are applied to the generator to produce high-quality images. The adversarial loss and morphological loss are used to guide the training process. The generator tries to minimize the adversarial loss, which is calculated as the binary cross-entropy between the discriminator's prediction for real/fake:

$\begin{aligned} \mathcal{L}_G=-\mathbb{E}\left[\log D_{\text {real}}\left(x_{\text {fake }}\right)\right]+\lambda & \cdot \mathcal{L}_{\text {morphology }}\left(x_{\text {fake}}, x_{\text {real}}\right)\end{aligned}$            (7)

where, $\mathcal{L}_{\text {morphology }}$ is the morphological loss. The discriminator's loss is computed as the sum of the real/fake classification loss and the morphology classification loss:

$\begin{aligned} & \mathcal{L}_D=-\left[\mathbb{E}\left[\log D_{\text {real}}\left(x_{\text {real}}\right)\right]+\mathbb{E}[\log (1\right.\left.\left.-D_{\text {real}}\left(x_{\text {fake }}\right)\right)\right]\left.+\mathbb{E}\left[\log D_{\text {morphology}}\left(x_{\text {real}}\right)\right]\right]\end{aligned}$             (8)

To improve the structural quality of the generated images, the edge loss and SSIM (Structural Similarity Index) loss is added. These loss functions are expressed as follows:

$\mathcal{L}_{\text {edge}}=\left\|\operatorname{Sobel}\left(x_{\text {fake}}\right)-\operatorname{Sobel}\left(x_{\text {real}}\right)\right\|_2^2$            (9)

$\mathcal{L}_{\mathrm{SSIM}}=1-\operatorname{SSIM}\left(x_{\text {fake}}, x_{\text {real}}\right)$               (10)

Compared to traditional conditional GANs, the M-GAN incorporates a morphology-based conditioning mechanism to generate images that are not just generic but conditioned on the specific sperm morphology types. This unique conditioning is used to generate sperm images that reflect the morphological diversity seen in real-world sperm images. Moreover, the latent vector in M-GAN is combined with a morphology class index. It is used for more controlled generation and enhancing the ability to model variations in size, shape, and scale.

In addition, attention mechanisms such as non-local attention and spatial attention are integrated into the generator. It is used to focus on critical regions of sperm images and address the scale variations and morphological diversity. Thus, M-GAN's novelty lies in its morphology-specific conditioning and the use of advanced attention mechanisms to address specific challenges such as scale variations and morphological diversity.

3.2 Biomorphological Attention Feature Network

In this work, BMAF-Net is proposed to extract morphological and texture features from sperm images. This architecture uses attention mechanisms and learnable Gabor filters to increase the classification accuracy of the model. The architecture of BMAF-Net involves several key layers to extract useful features and produce a compact feature vector for the classification task. The architecture of BMAF-Net is shown in Figure 2.

Figure 2. Architecture of Biomorphological Attention Feature Network (BMAF-Net)

3.2.1 Learnable Gabor layer

The image has dimensions $H \times W \times 3$, where $H$ and $W$ are the height and width of the image. Initially, learnable Gabor filters are applied to the input image to extract morphological features. The filter learns spatial frequencies and orientations that highlight structural elements in the sperm image. The output of this layer can be expressed as:

$F_{\mathrm{morph}}=\operatorname{Conv}\left(I, W_{\mathrm{gabor}}\right)+b$             (11)

where, $I$ is the input image, $W_{\text {gabor }}$ is the learnable Gabor filter, $b$ is the bias term, $F_{\text {morph}}$ is the output feature map that contains the morphological features of the sperm.

The inclusion of learnable Gabor filters in BMAF-Net is motivated by the need to capture both morphological and textural features of sperm images. The learnable Gabor filters are used for the network to adaptively learn the most relevant filter parameters for sperm morphology. It includes the specific structural patterns of sperm heads, midpieces, and tails. This flexibility enhances the network's ability to focus on both shape and texture.

3.2.2 Batch Normalization + ReLU

Next, Batch Normalisation is applied to stabilise the training process. It normalises the activations across the mini-batch. Then, the ReLU activation is used to introduce non-linearity. It can be written as:

$F_{\text {norm}}=\operatorname{ReLU}\left(\right.$ BatchNorm $\left.\left(F_{\text {morp }}\right)\right)$           (12)

where, $F_{\text {morph }}$ is the morphological feature map from the previous layer, and $F_{\text {norm}}$ is the normalized and activated feature map.

3.2.3 Morphological Attention Mechanism

The Morphological Attention Mechanism applies channel-wise attention to focus on important morphological features like the sperm head and tail. The attention map $A_{\text {morph}}$ is learned through a convolutional layer with a sigmoid function to produce attention weights for each channel. The attention map is given by:

$A_{\mathrm{morph}}=\sigma\left(\operatorname{Conv}\left(F_{\mathrm{norm}}\right)\right)$            (13)

where, $\sigma$ is the sigmoid function, $\operatorname{Conv}\left(F_{\text {norm}}\right)$ is the convolution applied to the normalized feature map, $A_{\text {morph }}$ is the attention map. Then, the feature map is recalibrated by multiplying it with the attention map:

$F_{\text {att}}=F_{\text {norm}} \times A_{\text {morph}}$           (14)

where, $F_{\mathrm{att}}$ is the feature map after applying attention.

3.2.4 Max pooling

Max pooling is applied to downsample the feature map and reduce its spatial dimensions. A 2 × 2 pooling operation is used. It is used to reduce the computational complexity and retain important spatial information. The output of the pooling operation is:

$F_{\mathrm{pool}}=\operatorname{MaxPool}\left(F_{\mathrm{att}}\right)$              (15)

where, $F_{\mathrm{pool}}$ is the downsampled feature map after max pooling.

3.2.5 Learnable Gabor layer

In BMAF-Net, the learnable Gabor layer is applied to capture texture features from the pooled feature map. This layer learns filters that capture texture patterns. The output feature map is:

$F_{\text {texture}}=\operatorname{Conv}\left(F_{\text {pool}}, W_{\text {gabor_texture}}\right)+b$         (16)

where, $F_{\text {pool}}$ is the pooled feature map from the previous layer, $W_{\text {gabor_texture}}$ is the learnable texture filter, $b$ is the bias term, $F_{\text {texture}}$ is the texture feature map.

3.2.6 Batch Normalization + ReLU

After extracting texture features, Batch Normalization and ReLU activation are applied to the texture feature map to normalize and introduce non-linearity. The output after these operations is:

$F_{\text {texture_norm}}=\operatorname{ReLU}\left(\right.$ BatchNorm $\left.\left(F_{\text {texture}}\right)\right)$            (17)

where, $F_{\text {texture}}$ is the texture feature map from the previous layer, and $F_{\text {texture_norm}}$ is the normalized and non-linearly activated texture feature map.

3.2.7 Morphological Attention

The second Morphological Attention Mechanism is applied to the texture feature map to focus on important texture regions in the sperm image. This attention mechanism recalibrates the texture feature map by emphasizing relevant texture patterns. The attention map for the texture features is:

$A_{\text {texture}}=\sigma\left(\operatorname{Conv}\left(F_{\text {texture_norm}}\right)\right)$              (18)

where, $A_{\text {texture }}$ is the texture-specific attention map.

The texture feature map is recalibrated by multiplying it by the attention map:

$F_{\text {texture_att}}=F_{\text {texture_norm}} \times A_{\text {texture}}$             (19)

3.2.8 Max pooling

A second max pooling (2 × 2) operation is applied to the texture attention features to reduce the spatial dimensions further while retaining the most discriminative texture features. The output after max pooling is:

$F_{\text {texture_pool}}=\operatorname{MaxPool}\left(F_{\text {texture_att}}\right)$             (20)

3.2.9 Fusion of morphological and texture features

The features from both the morphological and texture paths are fused together. This fusion is used for the network to learn a comprehensive representation of the sperm morphology. The fused feature map is:

$F_{\text {fusion}}=F_{\text {morph}} \oplus F_{\text {texture_pool}}$                    (21)

where, ⊕ represents concatenation. It combines the morphological and texture features.

3.2.10 Adaptive Average Pooling

Adaptive Average Pooling is applied to the fused feature map to aggregate the features into a fixed-size representation. The output after adaptive pooling is:

$F_{\text {avg }}=$ AdaptiveAvgPool $\left(F_{\text {fusion}}\right)$                           (22)

The aggregated features are passed through fully connected layers for feature recalibration. These layers are used to reduce the dimensionality of the feature vector and enable the model to focus on the most important features for classification or analysis. The output after the fully connected layers is:

$F_{\text {recal }}=\mathrm{FC}\left(F_{\text {avg}}\right)$                                (23)

The final output of the network is a feature vector that represents the sperm morphology. The final output is:

$F_{\text {output }}=\mathrm{FC}_2\left(F_{\text {recal}}\right)$                          (24)

where, $F_{\text {output}}$ is the final output feature vector that represents the sperm morphology.

3.3 Modified Waterwheel Plant Algorithm

The Waterwheel Plant Algorithm (WWPA) is inspired by the Aldrovanda vesiculosa carnivorous plant [25]. This plant hunts its prey using specialized traps. In optimization, WWPA finds an optimal solution based on exploration and stages. However, WWPA has two major drawbacks: slow Convergence and Local Optima Trapping.

3.3.1 Modified Waterwheel Plant Algorithm

To overcome these issues, in this work, a Modified Waterwheel Plant Algorithm is proposed. The strategy of differential mutation and a self-adaptive exploration-exploitation balancing is added to improve the performance of WWPA. In addition, memetic refinement is included to improve the convergence.

3.3.2 Position recognition and hunting of insects (exploration)

In this phase, the primary objective is to explore a wide area of the solution space. This phase simulates the behavior of the waterwheel plant as it searches for prey. The differential mutation introduces variability by perturbing the search agents’ positions using a mutation term. The waterwheels move to new positions based on their current position and velocity to avoid local minima. The position update in the exploration phase is modelled as follows:

$\mathrm{W}=r_1 \cdot(\mathrm{P}(t)+2 K)$               (25)

where, W is the velocity of the search agent, $r_1$ is a random number in the interval $[0,1], \mathrm{P}(t)$ is the current position of the search agent, $K$ is an exponential factor controlling the search area. Then, the new position of the waterwheel is updated as:

$\mathrm{P}(t+1)=\mathrm{P}(t)+\mathrm{W} \cdot\left(2 K+r_2\right)+\Delta \mathrm{P}$                      (26)

where, $r_2$ is another random variable, $\Delta \mathrm{P}=\beta \cdot\left(\mathrm{P}_r-\mathrm{P}_s\right)$ is the differential mutation term, where $\beta$ is a scaling factor, and $\mathrm{P}_r, \mathrm{P}_s$ are randomly selected positions from the population. If the optimization does not improve for three consecutive iterations, the position is reinitialized with a Gaussian distribution to ensure continued exploration:

$\mathrm{P}(t+1)=\mathcal{N}\left(\mu_P, \sigma\right)+r_1 \cdot(\mathrm{P}(t)+2 K \cdot \mathrm{~W})$                       (27)

where, $\mathcal{N}\left(\mu_P, \sigma\right)$ is a Gaussian distribution with mean $\mu_P$ and standard deviation $\sigma$.

3.3.3 Carrying the insect in a suitable tube (exploitation)

In this phase, WWPA focuses on refining the candidate solutions. It simulates the waterwheel plant’s behavior of carrying prey to a feeding tube. To find the best solution, the algorithm shifts its focus broadly. The waterwheel moves closer to the best solution found so far. If there is no improvement, a mutation term is applied to avoid stagnation in local optima. The position of each waterwheel is updated to move towards the best solution found so far:

$\mathrm{W}=r_3 \cdot\left(K \cdot \mathrm{P}_{\text {best }}(t)+r_3 \cdot \mathrm{P}(t)\right)$                       (28)

where, $\mathrm{P}_{\text {best}}(t)$ is the best solution found at iteration $t, r_3$ is a random factor in the interval $[0,2]$. The new position of the waterwheel is updated as follows:

$\mathrm{P}(t+1)=\mathrm{P}(t)+K \cdot \mathrm{~W}$                            (29)

If the solution does not improve after three iterations, a mutation term is applied to explore nearby areas:

$P(t+1)=\left(r_1+K\right) \cdot \sin (F C \theta)$                      (30)

where, $F$ and $C$ are random variables in the range $[-5,5], \theta$ is the angle guiding the search direction. The exploration rate decays to focus more on exploitation. This is controlled by the equation:

$K=\left(1+2 \cdot \frac{t^2}{T_{\max }}+F\right)$                            (31)

where, $t$ is the present iteration, $T_{\max }$ is the number of iterations.

3.3.4 Memetic refinement in exploitation

To accelerate convergence, a memetic refinement technique is incorporated. It performs simulated annealing or gradient descent to improve the best solutions found so far. It can be modelled as follows:

$\mathrm{P}(t+1)=\mathrm{P}(t)-\alpha \cdot \nabla F(\mathrm{P}(t))$                         (32)

where, $\alpha$ is a learning rate, $\nabla F(\mathrm{P}(t))$ is the gradient of the objective function at the current solution. Alternatively, Simulated Annealing can be applied as follows:

$\mathrm{P}(t+1)=\mathrm{P}(t)+\alpha \cdot \delta \mathrm{P}$                     (33)

where, $\delta \mathrm{P}$ is a random perturbation, $\alpha$ decreases over time to help fine-tune the solution. After feature extraction, MWWPA is applied for feature selection to remove irrelevant features. The selected features are forwarded to GBRF for classification.

Here, the four major components, like M-GAN, BMAF-Net, M-WWPA, and GBRF, are essential to address specific challenges in sperm morphology classification. M-GAN handles data augmentation and class imbalance by generating realistic sperm images. The BMAF-Net uses attention mechanisms and learnable filters for detailed feature extraction. Also, M-WWPA is applied to handle the high-dimensional nature of sperm image features. M-WWPA overcomes slow convergence and local optima issues in the feature selection process.

3.4 Gradient Boosted Random Forest

For classification, the GBRF-based hybrid ML is proposed. This model combines the strengths of RF and Gradient Boosting to improve predictive performance. The main innovation in GBRF is based on the integration of RF as base learners within a Gradient Boosting model. It uses a boosting mechanism to reduce the error of the entire ensemble of trees.

In conventional gradient boosting, each tree is trained sequentially to correct the errors made by the previous trees, which uses gradient descent on the residual errors. Likewise, RF constructs multiple decision trees independently and combines their results. The proposed GBRF combines these two techniques by incorporating RF as the base learner within the boosting framework. Here, the trees are trained iteratively with the goal of minimizing the ensemble's error.  In GBRF, Random Forests is used as a base learner for gradient boosting. At each iteration, instead of training a single decision tree, an RF is trained on the negative gradient of the current ensemble. It can be mathematically expressed as follows:

$\hat{y}^{(m)}=\hat{y}^{(m-1)}+\eta \cdot \frac{1}{T_m} \sum_{t=1}^{T_m} f_{t, m}(x)$                     (34)

where, $\hat{y}^{(m-1)}$ is the prediction from the previous iteration, $f_{t, m}(x)$ represents the prediction of the $t$-th tree from the $m$ th Random Forest, $T_m$ is the number of trees in the $m$-th Random Forest.

Here, at each boosting iteration, instead of adding a single decision tree, an RF is added, which is trained on the residual errors from the previous iteration. The residual error at each iteration is computed as:

$r_m=y-\hat{y}^{(m-1)}$                     (35)

where, $y$ is the true label, $\hat{y}^{(m-1)}$ is the predicted value from the previous iteration. Then, the RF is trained on the negative gradient of the residuals at each step:

$f_{t, m}(x)=$ TrainRandomForest $\left(r_m\right)$                         (36)