Gearbox Fault Detection Using Continuous Wavelet Transform and Vision Transformer (ViT)

Abstract

Managing the substantial volume of data generated by gearbox operations and the inherent complexities underscores the strategic advantage of employing deep learning methods. This study addresses fault detection in the gearbox using a Vision Transformer (ViT), whose images were prepared by continuous wavelet transformation. The dataset from the Open Energy Data Initiative (OEDI) serves as the foundation for our analysis. Data recorded by four vibration sensors, situated in the time domain and various operating loads, transform into two-dimensional images which, after some manipulations, are fed into a ViT, allowing us to tackle the classification challenge with 'Healthy' and 'Damaged' classes. Remarkably, we achieved an outstanding accuracy of 99.1554%, coupled with a reduction in the loss function to 0.040100. These results demonstrate the ViT's remarkable efficacy in detecting faults using images derived from signals recorded in the time domain.

Repository Contents

: Contains the source code for implementing the Continuous Wavelet Transform and Vision Transformer for gearbox fault detection.
: Contains the source code for Training the neural networks.
/model: Contains pre-trained neural networks and results.
: Contains the source code of the neural networks.
: some functions and plotter.

Dependencies

Python => 3.9
Pytorch => 2.0
NumPy => 1.4
Matplotlib
Wavelet Transform Library (PyWavelets)

Installation

Clone the repository:

git clone https://github.com/yriyazi/GearBox_Fault_detection-ICEE2024.git

Download the dataset from here and place it in the appropriate directory.
Run the .\dataloader\splitter.py to generate the CWT images.
```
python splitter.py
```
Run the scripts to preprocess the data and train the model:

python Train.py

Usage

depending on your specific need you may need to train, fine-tune, or just use some of the pre-trained weights.

Please make sure that the config.yaml configurations machs the model_weight.pt

Using the pre-trained weights.

the pre-trained weights are in the model folder. In Test.py change the model address (directory) and feed your data to the neural network.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use this code in your research, please cite:

[Ali Asadian, Yassin Riyazi, and Seyyed Mossa Ayati*. 2024. " Gearbox Fault Detection Using Continuous Wavelet Transform and Vision Transformer (ViT)." ICEE, 2024]

Acknowledgments

The Open Energy Data Initiative (OEDI) for providing the dataset.
The contributors to the libraries and frameworks used in this project.

Feel free to reach out for any questions or collaborations!

This README is generated with the assistance of ChatGPT.

yriyazi / GearBox_Fault_detection-ICEE2024

readme