Multiclass Railway Wheel Defect Classification Using Machine Learning Algorithms: A Comparative Study

Abstract

Railway wheel condition monitoring plays a critical role in ensuring safe and reliable rail operations, as defects in wheel geometry can lead to excessive wear, vibration, and potential derailment. However, real-world inspection data commonly exhibit severe class imbalance, where several defect types are underrepresented, making accurate multiclass classification challenging. This study aims to develop a machine learning–based framework for identifying seven defective wheel conditions using three geometric profile parameters which is flange thickness, flange height, and flange inclination (Qr). Kernel Density Estimation (KDE) with a Gaussian kernel (bandwidth = 0.2) was employed to generate synthetic data samples for minority classes. This augmentation increased the dataset from a skewed distribution to a balanced set of 1827 samples, preserving original statistical properties while providing a more robust foundation for model training. Three supervised classifiers namely, Random Forest, Linear-Support Vector Machine, and K-Nearest Neighbors were implemented and evaluated on the augmented dataset. Model performance was validated using a 5-fold cross-validation technique to ensure accuracy and generalizability. While all models achieved high predictive performance with accuracies exceeding 90%, the Random Forest model yielded the highest validated performance with a mean accuracy of 93.1% and a macro F1-score of 93.5%. Confusion matrix analysis confirmed that most defect types were correctly classified with minimal misclassification. These findings demonstrate that KDE-based augmentation effectively mitigates class imbalance, significantly enhancing the reliability of multiclass wheel defect classification in predictive maintenance systems.