Majority clustering for imbalanced image classification

Input dataset

The proposed model, MCIIC, is used to solve the problem of classification where the dataset is class imbalanced. The number of instances belonging to one class greatly outnumbers the instances in another class. The dataset has binary classes, and the occurrence of one class is much less frequent than others. Several machine learning algorithms exhibit a bias towards the majority class, resulting in suboptimal predictive performance for the minority class. To deal with this problem, data need to be pre-processed before classifying. In MCIIC, the dataset used are of large size, so first stage involves extraction of features.

Feature extraction

Feature extraction plays important role for effectively handling large image datasets by reducing dimensionality, helps to enhance the generalization of the image. To perform this task, the model uses Residual Network (ResNet-18). The key advancement in ResNet is the introduction of Residual blocks. ResNet is a combination of traditional CNNs with residual blocks, where the residual blocks allow the network to extract more advanced and hierarchal features. These residual connections allow the model to propagate gradients effectively through deep layers, allowing the learning of complex features. The ResNet-18 model significantly decreases the complexity of the image by reducing its large dimensions to just 512 features. The ResNet model is pre-trained on an ImageNet dataset and can classify images into one thousand object categories. Figure 2 shows the layered structure of a normal ResNet-18 model (Escobar Díaz Guerrero et al., 2024).

In the proposed model a ResNet-18 model (originally trained on ImageNet data set) is used as base. The ResNet-18 model comprises of eighteen layers, where the first 17 layers are used for feature extraction and final layer is responsible for predicting the image label. These initial layers are crucial for extracting significant features (512 features) from the input images. These extracted features are stored in an array. Subsequently, k-means clustering is used to divide majority class into multiple classes based on the extracted features.

Elbow method

In the proposed model, to address class imbalance in the data, the majority class is partitioned into multiple clusters using k-means clustering. The optimal number of clusters (k) is determined using the elbow method. This approach requires minimizing a sum of squared errors (SSE): (1)${\displaystyle SSE=\sum _{k=1}^{cluster\text{\_}n}\sum _{x_j\in X}∥x_i-c_k∥}$

where k = N_c is the optimal number of clusters where the elbow is formed in the graph, x_i is pixel value of an image in the dataset and c_k is the cluster centre which is initialized manually and cluster_n is the maximum number of clusters to be considered. The elbow method helps to finalize the value of cluster_n for k-means clustering where the elbow point in Fig. 3 points to the optimal number of clusters to considered for the process of clustering.

The algorithm for the elbow method is given below Algorithm 1:

Figure 3 shows the graph of elbow method for COVID/no-COVID test dataset providing the optimal value of k as the output. The below given graph forms a curve at k = 5. The obtained value of k is the number of classes in which the majority class will be divided using k-means clustering.

According to the elbow method, the determined number of clusters is applied to partition the majority class. The dataset M1 representing the majority class is segmented into Nc clusters using k-means clustering, with each cluster assigned a new class label. This procedure results in a reduction of the number of images per class as the number of majority classes increases to N_c. Consequently, this transforms the original binary-class problem into a multi-class problem with N_c+1 total classes, where the (N_c+1)^th class represents the minority class.

Training the model using transfer learning

Once the data is classified into N_c+1 classes, it can be used to train the required model. These models can be trained from scratch i.e., training a neural network without relying on pre-existing knowledge or pre-trained weights. Alternatively, there is an option of employing machine learning techniques such as transfer learning to train new models based on previously trained models, particularly useful when dealing with smaller datasets.

Transfer learning is a concept of machine learning which involves using a pretrained neural network model to improve the learning process on another task, i.e., model trained on one task is adapted for second task (Shorten & Khoshgoftaar, 2019). The process of transfer learning involves selecting a pretrained model (ResNet-18 model, trained on a large data set), removing the ending layers responsible for specific tasks, and finally introducing new layers after the original layers. These layers are responsible for training the model for new tasks. The last step involves training the model for the new task using new data sets.

Transfer learning is often considered better than training a model from scratch because it provides data efficiency, faster convergence, and improved generalization. In many real-world scenarios, obtaining a large labelled dataset for a specific task can be challenging and expensive. By leveraging a pre-trained model’s learned features from a related task, transfer learning enables the adaptation to a target task even with limited labelled data. As the pre-trained model has already learnt generic features from a different but related task, transfer learning provides quicker convergence and reduces resource requirements. Figure 4 shows the difference between training a model from scratch vs transfer learning as learned knowledge is transferred while training the new model resulting into improved efficiency (Zhong et al., 2020).

Thus, the goal of using transfer learning for training the model is to leverage a pretrained model’s features without extensively retraining the complete model. The dataset is split into 80-10-10 ratio for training, validation and testing. Training is now performed on the new multiclass dataset using transfer learning. ResNet-18 pretrained on ImageNet dataset, is used for training the dataset to achieve the desired outcomes.

Validating the training model and tuning of the hyper-parameters

During the training process, the model’s weights are continuously updated to minimize a chosen loss function (e.g., cross-entropy loss) on the training data. However, to prevent overfitting and ensure the model generalizes well to unseen data, the model’s performance on the validation set is monitored. The best model weights are those that result in the highest accuracy (or lowest loss) on the validation set and with the less computational cost. The standardized parameters are chosen for the process of initialization to produce more optimized results.

Analysing the performance

Once the model gives output for multiclass, this is then converted to binary class to obtain the final classification. To test the accuracy of the model, the multi-class data is transformed into binary class data using the process of mapping. The mapping is based on the labels, such that classes from the majority class are mapped into one class and label with the minority class are mapped into desired class. The performance of the model is evaluated on a binary validation dataset by mapping its predictions to binary values, and selecting the best model weights based on the highest accuracy achieved on the validation set during training.

The architecture of the proposed methodology for the classification of imbalanced image dataset is shown in Fig. 5. The input dataset is imbalanced dataset of tuberculosis images and normal images, and the proposed methodology aims to classify minority class images (tuberculosis images) correctly.

Analysing the time-complexity

To determine the time-complexity of the proposed method assuming the following parameters:

M1-Represents images in the majority cluster

N_c-Represents number of clusters in which majority class is divided

d- Dimension of the dataset

L-Number of layers

P-Number of parameters

G-Number of hyper-parameter combination

T-Number of training runs

K-Number of time k-means clustering algorithm run for different values of N_c

Performing the elbow method using k-means clustering takes O(M1N_cKd). The classification process using the model trained with transfer learning takes O(M1LP) and hyper-parameters tuning requires O(GT). The overall time-complexity of the proposed algorithm takes O(M1N_cKd+M1LP+GT).

Computing infrastructure:

OS: Ubuntu 22.04.4 LTS

CPU: Intel(R) Core(TM) i5-6600

GPU: NVIDIA GeForce RTX 3090

RAM: 32 GB

To explore the efficiency of the proposed model MCIIC, the proposed model is further integrated with weighted data augmentation and weighted data classification.

1. The first approach involves conducting the analysis without any additional changes to the dataset. This means that no adjustments are made to the dataset before the training and testing of the model. All the data points are used without any modifications. This provides a baseline for understanding how the model performs using the given data without any additional processing. This allows for the evaluation of model’s raw performance on the dataset without any adjustments.

2. The second approach incorporates the proposed algorithm MCIIC with the weighted classification into the evaluation process to assess its impact on the results. Since most machine learning algorithms are not good with biased classes, the concept of weights is applied by giving weights to both majority class and minority class (Schaudt et al., 2023). This difference in weights will influence the classification during training process. It reduces the misclassification made by minority class by setting a higher-class weight and simultaneously reducing the weights of the majority class. Here weighted loss function is used instead of normal weight function to train the model. The weights are used to assign the higher penalty while miss-classification of the minority class. This will increase the cost of miss-classification of the minority class.

3. In the third approach, the MCIIC is integrated with weighted augmentation. In weighted augmentation the weights are assigned to the classes based on the number of the samples in the specific class. The minority class is provided with higher weights and majority class with lower weights. The class with the higher weights is having high probability of the data needs to be augmented. To solve the problem of class imbalance, the minority class is biased with higher weights to pay more attention for the augmentation. For example, if class a, b and c have 10, 20 and 30 samples each, respectively, then augmentation probability of class a will be 1, class b will be 0.5 and class c will be 0.33. Different types of image augmentation techniques include image rotation, image shifting, image blurring, image flipping, and image noising. These techniques help in efficient training of model, even in the absence of a large dataset.

For training and testing the model using the described approaches, various datasets are required. These datasets are explained below.

Datasets

1. Bald/no-bald dataset: This dataset contains the images of faces. It is divided into two types of images: one with the faces of bald individuals and other with the faces of individuals with hair on their scalp. This dataset is acquired from CelebA dataset (Li, 2018). CelebA dataset has 40 different classes, where bald/no-bald attribute category stands out having the most notable difference in the number of samples of these two classes. The size of the training data includes number of bald images are 158,442 and number of no-bald images are 3,637. The imbalance ratio is 2.3%. Some of the images belonging to the “bald” and “no-bald” classes are shown in Figs. 6 and 7.

2. Normal/tuberculosis dataset: The second dataset consists of medical chest X-ray images. One category, named as “Normal” consists of healthy chest images, while the other category named “Tuberculosis”, consists of X-rays depicting chest images of people infected with tuberculosis. These images have been obtained from the imbalanced Tuberculosis and Pneumonia dataset (Roshan, 2022). For binary classification, only the “normal” and “tuberculosis” classes are considered. The size of the training data includes number of normal images without tuberculosis are 7,430 and number of images with tuberculosis are 1,510. The imbalance ratio is 20.3%. Some of the images belonging to the “Normal” and “Tuberculosis” classes are shown in Figs. 8 and 9.

3. COVID/no-COVID dataset: This dataset includes chest X-ray images used for covid detection. One category, named “COVID” consists of chest X-ray images showing signs of COVID infection, while the “no-COVID” category consists of images without any COVID infection. These images are obtained from the COVIDx CRX-3 dataset (Roshan, 2022). To introduces class imbalance, approximately 20% of the images from the COVID class are randomly selected. The size of the training data includes number of no-COVID images are 11,352 and number of COVID images are 681. The imbalance ratio is 6%. Some of the images belonging to the “COVID” and “no-COVID” classes are shown in Figs. 10 and 11. Table 1 shows the summary of the features of the datasets used for evaluation.

Evaluation metrics

To assess the effectiveness of the proposed algorithm compared to other classification techniques, confusion matrices are employed to analyze how well the predictions align with the actual ground truth values. When making predictions, instances from the minority class are rare occurrences where a positive outcome is indicated. In such cases, the majority class is typically considered as negative. The metrics used to validate the performance of the proposed algorithms are:

Accuracy: In a classification model, accuracy measures the correctly predicted labels, indicating the proportion of correctly predicted cases out of all instances. It is calculated as Eq. (2).

where TP (true positive) denotes correctly predicted positive cases, TN (true negative) denotes correctly predicted negative cases, FP (false positive) indicates incorrectly predicted positive cases, and FN (false negative) indicates incorrectly predicted negative cases.

Precision: It evaluates the precision of positive predictions, indicating how reliably the model identifies positive instances correctly. The better value of high precision indicates the low value of incorrectly predicted positive instances. Equation (3) represents the measure of the precision as (3)${\displaystyle Precision=\frac{TP}{TP+FP}.}$

Recall: The recall metric evaluates how well the model can correctly identify positive instances among all the actual positive instances. Equation (4) represents the measure of the recall as (4)${\displaystyle Recall=\frac{TP}{TP+FN}.}$

The higher vale of recall implies low value of the false negative instances predicted by the model.

F1_score: This measure balances the precision and recall when both are equally important. It is calculated as the harmonic mean of the precision and recall as Eq. (5) (5)${\displaystyle F{1}_{score}=\frac{2\ast Precision\ast Recall}{\left(Precision+Recall\right)}.}$

Furthermore, in our research, the confusion matrix plot proves highly valuable due to its insensitivity to sample balance, making it particularly adept at handling skewed data distributions. Moreover, these plots demonstrate greater responsiveness to variations in model performance, with higher values indicating a greater likelihood that the classifier will effectively rank positive samples, thereby enhancing classification accuracy (Yuan et al., 2021).

To choose the consistency of the experimentation, the training parameters and configuration chosen are as follows (Box 1):

Results and Discussions

To evaluate the performance of the proposed model, analysis is done with different integrated models considering the class imbalance problem. The classification is done using simple baseline model using ResNet-18, ResNet-18 with weighted augmentation, ResNet-18 with weighted classification, ResNet-18 with MCIIC, ResNet-18 with MCIIC and weighted augmentation, and ResNet-18 with weighted classification. After a direct comparison with the baseline method, and other modified models it is observed from Table 2. That classification results for bald/no-bald dataset where MCIIC model shows better accuracy as compared to other models. When MCIIC is applied with weighted augmentation is shows even better results.

Table 3 shows the results of classification performed on normal or tuberculosis dataset. It shows that classification performed using MCIIC with weighted augmentation outperforms in accuracy.

Recall and F1-score are better when the classification is performed using MCIIC with weighted classification as compared to base model of ResNet-18 and other hybrid models with ResNet-18.

Table 4 depict the results of comparison between ResNet model and the proposed MCIIC model for the COVID/no-COVID dataset, respectively. It can be observed form the results that proposed model outperforms in accuracy, precision and F1-score while hybrid model of MCIIC with weighted classification performs well giving better recall values.

Figure 12 shows the confusion matrix plots for bald/no-bald dataset. Figure 12A shows the classification performed using normal ResNet-18 model, Fig. 12B shows classification performed using proposed MCIIC, where the number of samples in the minority class (bald samples) has increased for all three classes with the increase in the value of true positives and the number of majority class samples (no-bald samples) has reduced as false negative samples are reduced. Figure 12C shows classification performed using MCIIC integrated with weighted augmentation and Fig. 12D shows classification performed using MCIIC integrated with weighted classification. The proposed method outperforms when integrated with weighted augmentation giving better results on imbalanced datasets. Figure 13 shows the confusion matrix plots of the COVID/no-COVID dataset. Here Fig. 13D shows the best results where classification accuracy of minority class increases with the proposed approach integrated with weighted classification. Figure 14 shows the confusion matrix plots for tuberculosis/normal dataset. Here the proposed MCIIC results shown in Fig. 14B achieves maximum accuracy.

Upon analysing the outcomes of the proposed model against those of the baseline method, it becomes evident that the proposed model outperforms the baseline across all approaches. The results prove the effectiveness of the model for the process of image classification even with skewed datasets. Therefore, proving its efficiency and scope in various applications.