This code repository is the official code repository of paper "Data Augmentation for Bearing Fault Diagnosis Using a Novel Flow-Based Generative Model", and the account author is the original author of the paper and source code. The project's file usage instructions will be supplemented later. The code files of this repository are still being uploaded recently, and the weight files will be considered to be uploaded to Google Cloud Drive due to their large size. If you reference this project, please inform us in the issue, thank you!(2025-3-23)

We have uploaded the weight files, which can be downloaded and used at the link. At the same time, each folder has a brief introduction, which can be seen in the Readme in each folder.(2025-3-25)

To ensure the credibility of our paper, we decided to provide the pre-trained weight files of our model, and also to facilitate readers to directly complete image data synthesis by loading our weight files into the generator file. Our model pre-trained weight files can be obtained from this link: https://drive.google.com/drive/folders/1-nT8oyWBRETIntGGdyILk6AqfBFtFGXD?usp=drive_link Replace the downloaded checkpoint folder with the downloaded checkpoint folder, and then load the corresponding .pth file in the folder in the generator to run and synthesize data.

dl.yaml: This file contains all the environment configurations required for this project. You can use this file in conda to quickly deploy a virtual environment that can run our project.

config.yaml: This file contains the training and generation parameter configurations of the train.ipynb and generator.ipynb files, which can be adjusted according to your needs.

train.ipynb: This file can be used to train your model, but the data needs to be adjusted to a level acceptable to the model.py and dataset folders. The specific resolution should be adjusted according to the individual device, and 32x32 pixel images are recommended. generator.ipynb: This file is used to load and synthesize pre-trained files.

Bearings are critical components in industrial machinery, making fault diagnosis essential for ensuring operational reliability. Due to the rarity of fault conditions, the resulting signal data are extremely scarce compared to normal signal data, resulting in severe class imbalance and biasing traditional diagnostic methods toward the normal states. Researchers have turned to data augmentation to address the aforementioned challenge. However, existing approaches often struggle to capture the correlations between spatiotemporal and conditional information and suffer from inefficiency. To address these limitations, we propose a novel Spatiotemporal-Conditional information fusion Classifier-Free Guidance Rectified Flow generative model (SCIF-CFGRF). Specifically, the Continuous Wavelet Transform-based Feature Extraction Network block and the Residual Seaformer block, both carefully designed and integrated into the UNet architecture, enable high-fidelity modeling of high-dimensional Rectified Flow (RF) representations in 2D time–frequency domain signal feature heatmaps, while simultaneously embedding spatiotemporal-conditional information fusion. Furthermore, a dedicated loss function tailored to the characteristics of RF ensures fast and stable convergence during training. Finally, a novel spatiotemporal-conditional information fusion classifier-free guidance inference sampler and Quality Enhancer block complete the efficient and high-quality synthesis of samples. Experiments on two real-world bearing fault datasets and extensive ablation studies demonstrate that SCIF-CFGRF significantly outperforms mainstream methods in terms of synthesis quality and inference speed. Under a 1:5 class balance ratio, SCIF-CFGRF achieves a Cosine Similarity exceeding 0.94 and a diagnostic accuracy of 99.50% on synthetic samples while requiring only 6.97% of the inference time compared to the DDPM-based model.

The crucial contributions of this study are as follows:

(1) We propose SCIF-CFGRF, a model that fuses spatiotemporal-conditional information via a CFEN-block and RS-block within a UNet backbone built upon RF. The CFEN-block combines Continuous Wavelet Transform, RF, and a feature extractor to model RF in the 2D time–frequency domain, generating feature heatmaps and high-dimensional representations. The RS-block further fuses these features and spatiotemporal-conditional information to enable high-fidelity signal modeling.

(2) We design an optimization loss tailored to the path and velocity characteristics of RF, facilitating faster and more stable training of the SCIF-CFGRF backbone. This enhances both synthesis quality and training efficiency, supporting rapid industrial deployment.

(3) We propose a novel Classifier-Free Guidance inference sampler and a QE-block based on spatiotemporal-conditional RF information fusion, which eliminates the need for explicit classifiers while enhancing both generation quality and efficiency. The integration of the RF-based Euler synthesis method further improves performance and accelerates deployment.

(4) Experiments on 14 balanced datasets from two real world bearing datasets show that SCIF-CFGRF surpasses mainstream methods in synthesis quality. Under 1:400 imbalance, it achieves a cosine similarity of 0.94 and a downstream accuracy of 80%, with inference time reduced to 5% of DDPM.

In the future, we plan to test our method on additional industrial equipment, such as aircraft engines, chillers, and gearboxes. Moreover, we aim to explore its application in the financial sector for related research.

CPU recommendation: i7-14650HX(Laptop)

GPU recommendation: RTX-4060(8GB)(Laptop)

Memory：32GB

Solid State Drive：1TB

It is recommended to use a higher configuration than this. The author's device uses shared memory to prevent out of GPU memory. It is recommended that readers use devices with more than 12GB of GPU memory to prevent program crashes or reduced model performance.

Fig.1 Proposed SCIF-CFGRF model structure diagram.

Fig.2 Proposed CFEN-block structure diagram.

Fig.3 Proposed RS-block structure diagram.

Fig.4 Proposed methodological framework.

Fig.5 Result on CWRU dataset.

Fig.6 Result on SEU dataset.

Fig.7 GIF of the inference synthesis process on the CWRU dataset at BR1:400.

Fig.8 GIF of the inference synthesis process on the CWRU dataset at BR1:5.

Fig.9 GIF of the inference synthesis process on the SEU dataset at BR1:400.

Fig.10 GIF of the inference synthesis process on the SEU dataset at BR1:5.

Fig.11 Inference process efficiency comparison.

If your paper, research or project uses our research, please use this latex citation format:

@ARTICLE{11261892, author={Dai, Hong-Liang and Lin, Dong-Jie}, journal={IEEE Transactions on Instrumentation and Measurement}, title={Spatiotemporal-Conditional Information Fusion Classifier-Free Guidance Rectified Flow Model for Bearing Fault Diagnosis}, year={2025}, volume={}, number={}, pages={1-1}, keywords={Spatiotemporal phenomena;Radio frequency;Feature extraction;Fault diagnosis;Training;Continuous wavelet transforms;Diffusion models;Heating systems;Accuracy;Data models;Data imbalance;Fault diagnosis;Information Fusion;Deep Learning;Rectified Flow;Generative Models}, doi={10.1109/TIM.2025.3635328}}

If you plagiarize our research, we will pursue legal action.

We will continue to improve the code comments of the project later.

You can also follow our other fault data synthesis project DLRSD-SMOTE (https://github.com/amstlldj/DLRSD-SMOTE). The project will also be updated and maintained in the future.

If you have any questions, please contact the author's work email or leave a message in issues. The author will try his best to answer them at his convenience.If you are interested in seeking cooperation, you can also consult this email. Author contact email: 2112464120@e.gzhu.edu.cn