v1:第一版修整代码

BFDS_train.py

@@ -2,8 +2,8 @@ import os
 import logging
 import warnings
 import json
-from datetime import datetime
 import requests
+from datetime import datetime
 
 if __name__ == "__main__":
     try:
@@ -17,7 +17,7 @@ if __name__ == "__main__":
         os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
         print(f"无法连接到 Hugging Face:换源到{os.environ['HF_ENDPOINT']}")
     if not os.path.exists("./cache"):
-        os.makedirs("./cache")  # 创建缓存目录
+        os.makedirs("./cache")
     os.environ["HF_DATASETS_CACHE"] = "./cache"
 
 from utils.logger import setlogger
@@ -35,14 +35,14 @@ class Argument:
         self.data_set = "BFDS-Project/Bearing-Fault-Diagnosis-System"  # 数据集huggingface地址
         self.conditions = fetch_all_conditions_from_huggingface(self.data_set)  # 数据集的配置和分割信息如果想要知道明确的信息来确定迁移方向请自行运行fetch_conditions.py
         self.labels = {"Normal Baseline Data": 0, "Ball": 1, "Inner Race": 2, "Outer Race Centered": 3, "Outer Race Opposite": 4, "Outer Race Orthogonal": 5}  # 标签
-        self.transfer_task = [["CWRU224", "12kDriveEnd"], ["CWRU224", "12kFanEnd"]]  # 迁移方向
-        self.target_domain_labeled = False  # 表示目标域在训练中是否带有标签
+        self.transfer_task = [["CWRURPM", "12kDriveEndrpm1730"], ["CWRURPM", "12kDriveEndrpm1750"]]  # 迁移方向
+        self.target_domain_labeled = True  # 表示目标域在训练中是否带有标签
 
         # 预处理
         self.normalize_type = None  # 归一化方式, mean-std/min-max/None
         self.stratified_sampling = True  # 是否分层采样, True/False
         # 模型
-        self.model_name = "ResNet"  # 模型名
+        self.model_name = "CNN"  # 模型名
         self.bottleneck = True  # 是否使用bottleneck层
         self.bottleneck_num = 256  # bottleneck层的输出维数
 
@@ -118,7 +118,8 @@ class Argument:
             "data_set": "BFDS-Project/Bearing-Fault-Diagnosis-System",
             "conditions": fetch_all_conditions_from_huggingface("BFDS-Project/Bearing-Fault-Diagnosis-System"),
             "labels": {"Normal Baseline Data": 0, "Ball": 1, "Inner Race": 2, "Outer Race Centered": 3, "Outer Race Opposite": 4, "Outer Race Orthogonal": 5},
-            "transfer_task": [["CWRU224", "12kDriveEnd"], ["CWRU224", "12kFanEnd"]],
+            "transfer_task": [["CWRURPM", "12kDriveEndrpm1730"], ["CWRURPM", "12kDriveEndrpm1750"]],
+            "target_domain_labeled": False,
             "normalize_type": None,
             "stratified_sampling": True,
             "model_name": "CNN",
@@ -126,7 +127,7 @@ class Argument:
             "bottleneck_num": 256,
             "batch_size": 64,
             "cuda_device": "0",
-            "max_epoch": 2,
+            "max_epoch": 100,
             "num_workers": 0,
             "checkpoint_dir": "./checkpoint",
             "print_step": 50,
@@ -136,8 +137,8 @@ class Argument:
             "lr": 1e-3,
             "lr_scheduler": "step",
             "gamma": 0.1,
-            "steps": [150, 250],
-            "middle_epoch": 0,
+            "steps": [25, 75],
+            "middle_epoch": 50,
             "distance_option": True,
             "distance_loss": "JMMD",
             "distance_tradeoff": "Step",

BFDS_web.py

@@ -1,6 +1,11 @@
+import logging
 import os
+import pandas as pd
 import requests
+import torch
+import warnings
 import zipfile
+from datetime import datetime
 
 if __name__ == "__main__":
     try:
@@ -20,18 +25,11 @@ if __name__ == "__main__":
 
 import gradio as gr
 from BFDS_train import Argument
-import torch
-from utils.predict import predict
-import pandas as pd
-
-import logging
-import warnings
-from datetime import datetime
-
-
 from utils.logger import setlogger
+from utils.predict import predict
 from utils.train import train_utils
 
+
 # 初始化 Argument 实例
 args = Argument()
 args.set_recommended_params()

Dockerfile.cpu

@@ -0,0 +1,31 @@
+FROM python:3.11-alpine
+
+# 设置环境变量以减少缓存和输出
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1
+
+# 设置工作目录
+WORKDIR /app
+
+# 安装必要的系统依赖
+RUN apk add --no-cache --virtual .build-deps \
+    gcc \
+    musl-dev \
+    libffi-dev \
+    openssl-dev \
+    make \
+    && apk add --no-cache \
+    bash \
+    git \
+    && pip install --upgrade pip
+
+# 安装 Python 依赖并清理构建依赖
+COPY requirements-cpu.txt .
+RUN pip install --no-cache-dir -r requirements-cpu.txt \
+    && apk del .build-deps
+
+# 复制项目文件
+COPY . .
+
+# 设置默认运行命令
+CMD ["python", "BFDS_web.py"]

README.md

@@ -1,7 +1 @@
-# Bearing-Fault-Diagnosis-System
-
-## Requirements
-- Python 3.10
-- matplotlib 3.10.0
-- numpy 2.2.2
-- PyWavelets 1.8.0
+# Bearing-Fault-Diagnosis-System

dataset/get_data.py

@@ -1,8 +1,8 @@
-import pandas as pd
-import numpy as np
-from datasets import load_dataset
 import librosa
 import mimetypes
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
 
 # ===============================================================
 # 加载有标签的数据集(n , m + 1)最后一列是标签

dataset/get_dataset.py

@@ -1,7 +1,9 @@
+from typing import Literal, Optional
+
 import pandas as pd
 import torch
-from torch.utils.data import Dataset, DataLoader, Subset, random_split
-from typing import Optional, Literal
+from torch.utils.data import DataLoader, Dataset, Subset, random_split
+
 from dataset.get_data import get_huggingface_dataset, get_local_dataset, get_user_dataset
 
 

models/CNN.py

@@ -33,7 +33,7 @@ class CNN(nn.Module):
             nn.AdaptiveMaxPool1d(4),
         )  # 128, 4,4
 
-        self.layer5 = nn.Sequential(nn.Linear(128 * 4,self.__in_features), nn.ReLU(inplace=True), nn.Dropout())
+        self.layer5 = nn.Sequential(nn.Linear(128 * 4, self.__in_features), nn.ReLU(inplace=True), nn.Dropout())
 
     def forward(self, x):
         x = self.layer1(x)

models/ResNet18_1d.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 
+
 class BasicBlock1D(nn.Module):
     expansion = 1  # 扩展倍数，用于调整输出通道数
 
@@ -32,6 +33,7 @@ class BasicBlock1D(nn.Module):
 
         return out
 
+
 class ResNet1D(nn.Module):
     def __init__(self, block=BasicBlock1D, layers=[2, 2, 2, 2]):
         super(ResNet1D, self).__init__()
@@ -58,8 +60,7 @@ class ResNet1D(nn.Module):
         # 如果需要调整通道数或步幅不为1，则定义下采样层
         if stride != 1 or self.in_channels != out_channels * block.expansion:
             downsample = nn.Sequential(
-                nn.Conv1d(self.in_channels, out_channels * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
+                nn.Conv1d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
                 nn.BatchNorm1d(out_channels * block.expansion),
             )
 
@@ -96,19 +97,16 @@ class ResNet1D(nn.Module):
     def output_num(self):
         # 返回输出特征维度
         return self.__in_features
-    
+
+
 def resnet1d18():
     # 构建 ResNet1D-18 模型
     return ResNet1D(layers=[2, 2, 2, 2])
 
+
 if __name__ == "__main__":
-    # 调试和测试模型
-    model = resnet1d18()  # 输出固定为 256 特征
+    model = resnet1d18()
     print(model)
-
-    # 创建一个随机输入张量，批量大小为 8，信号长度为 1024
     input_tensor = torch.randn(8, 1, 1024)
     output = model(input_tensor)
-
-    # 打印输出形状
     print("Output shape:", output.shape)

requirements-cpu.txt

@@ -0,0 +1,96 @@
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.11.16
+aiosignal==1.3.2
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+audioread==3.0.1
+certifi==2025.1.31
+cffi==1.17.1
+charset-normalizer==3.4.1
+click==8.1.8
+colorama==0.4.6
+contourpy==1.3.1
+cycler==0.12.1
+datasets==3.5.0
+decorator==5.2.1
+dill==0.3.8
+fastapi==0.115.12
+ffmpy==0.5.0
+filelock==3.18.0
+fonttools==4.57.0
+frozenlist==1.5.0
+fsspec==2024.12.0
+gradio==5.24.0
+gradio_client==1.8.0
+groovy==0.1.2
+h11==0.14.0
+httpcore==1.0.8
+httpx==0.28.1
+huggingface-hub==0.30.2
+idna==3.10
+Jinja2==3.1.6
+joblib==1.4.2
+kiwisolver==1.4.8
+lazy_loader==0.4
+librosa==0.11.0
+llvmlite==0.44.0
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+matplotlib==3.10.1
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.4.3
+multiprocess==0.70.16
+networkx==3.4.2
+numba==0.61.2
+numpy==2.2.4
+orjson==3.10.16
+packaging==24.2
+pandas==2.2.3
+pillow==11.1.0
+platformdirs==4.3.7
+pooch==1.8.2
+propcache==0.3.1
+pyarrow==19.0.1
+pycparser==2.22
+pydantic==2.11.3
+pydantic_core==2.33.1
+pydub==0.25.1
+Pygments==2.19.1
+pyparsing==3.2.3
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==14.0.0
+ruff==0.11.5
+safehttpx==0.1.6
+scikit-learn==1.6.1
+scipy==1.15.2
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+soundfile==0.13.1
+soxr==0.5.0.post1
+starlette==0.46.1
+sympy==1.13.1
+threadpoolctl==3.6.0
+tomlkit==0.13.2
+torch==2.6.0
+torchaudio==2.6.0
+torchvision==0.21.0
+tqdm==4.67.1
+typer==0.15.2
+typing-inspection==0.4.0
+typing_extensions==4.13.2
+tzdata==2025.2
+urllib3==2.4.0
+uvicorn==0.34.0
+websockets==15.0.1
+xxhash==3.5.0
+yarl==1.19.0

utils/fetch_conditions.py

@@ -18,7 +18,6 @@ if __name__ == "__main__":
     os.environ["HF_DATASETS_CACHE"] = "./cache"
 
 from datasets import get_dataset_config_names, get_dataset_split_names
-import json
 
 
 def fetch_all_conditions_from_huggingface(dataset_name):

utils/future_use.py

@@ -1,208 +0,0 @@
-# 暂时不使用的代码
-import numpy as np
-import pywt
-import matplotlib.pyplot as plt
-import os
-
-
-class ContinuousWaveletTransform:
-    def __init__(self, fs, signals, save_path=None, wavelet="cmor1.5-1.0", freqNum=224):
-        """
-        连续小波变换 (CWT) 计算类，支持单个信号或批量信号输入，并保存为 .npy 文件。
-
-        Args:
-            fs (_int_): 采样频率
-            signals (_np.array_): 输入信号，形状可以是 (signal_length,) 或 (batch_size, signal_length)
-            save_path (_str_): 如果提供，将保存 CWT 变换后的数据到 .npy 文件，默认为None，为None时不保存
-            wavelet (_str_): 连续小波类型（默认 'cmor1.5-1.0'）
-            freqNum (_int_): 频率点个数（默认 224）
-        """
-        self.fs = fs
-        self.save_path = save_path
-
-        # 确保路径存在
-        if save_path:
-            os.makedirs(save_path, exist_ok=True)
-
-        # 确保输入是 NumPy 数组
-        signals = np.asarray(signals, dtype=np.float32)  # 使用 float32 节省内存
-
-        # 处理 batch 维度
-        if signals.ndim == 1:
-            signals = signals[np.newaxis, :]  # (signal_length,) -> (1, signal_length)
-
-        self.batch_size, self.signal_length = signals.shape
-        self.time = np.arange(0, self.signal_length) / fs  # 时间轴
-        self.widths = np.geomspace(1, 512, num=freqNum).astype(np.float32)  # 频率尺度
-
-        # 预分配 CWT 结果矩阵
-        self.cwt_results = np.empty((self.batch_size, freqNum, self.signal_length), dtype=np.float32)
-
-        for i in range(self.batch_size):
-            signal = signals[i] - np.mean(signals[i])  # 去均值（去直流分量）
-            cwtmatr, freqs = pywt.cwt(signal, self.widths, wavelet, sampling_period=1 / fs)
-            cwt_result = np.abs(cwtmatr).astype(np.float32)  # 取模值，转换为 float32
-            self.cwt_results[i] = cwt_result
-
-            # 保存 CWT 结果到 .npy 文件
-            if save_path:
-                np.save(os.path.join(save_path, f"cwt_{i:04d}.npy"), cwt_result)
-                print(f"CWT 结果已保存到 {os.path.join(save_path, f'cwt_{i:04d}.npy')}")
-
-        self.freqs = freqs.astype(np.float32)  # 存储频率信息，节省内存
-
-    def plot(self, index=0, logspace=True, save_path=None):
-        """
-        绘制并可选保存 CWT 结果。
-
-        Args:
-            index (_int_): 选择绘制的信号索引
-            logspace (_bool_): 是否使用对数坐标绘制频率轴
-            save_path (_str_ 或 None): 如果提供路径，则保存 `.npy` 文件，否则不保存
-        """
-        if index >= self.batch_size:
-            raise ValueError(f"Index 超出范围！batch_size = {self.batch_size}, 但 index = {index}")
-
-        # 获取 CWT 结果
-        cwt_matrix = self.cwt_results[index]
-
-        # 选择是否保存 .npy
-        if save_path:
-            np.save(save_path, cwt_matrix)
-            print(f"CWT 结果已保存到 {save_path}")
-
-        # 绘图
-        fig, ax = plt.subplots(figsize=(10, 5))
-        pcm = ax.pcolormesh(self.time, self.freqs, cwt_matrix, shading="auto")
-
-        ax.set_yscale("log" if logspace else "linear")
-        ax.set_xlabel("Time (s)")
-        ax.set_ylabel("Frequency (Hz)")
-        ax.set_title(f"CWT Scaleogram (Signal {index})")
-        fig.colorbar(pcm, ax=ax)
-
-        plt.show()  # 显示图像
-
-
-if __name__ == "__main__":
-    fs = 1e3
-    N = 1e4
-    noise_power = 1e-3 * fs
-    time = np.arange(N) / float(fs)
-    mod = 2 * np.pi * 20 * np.cos(time)
-    carrier = np.sin(2 * np.pi * 100 * time + mod)  # 频率调制
-
-    rng = np.random.default_rng()
-    noise = rng.normal(scale=np.sqrt(noise_power), size=time.shape)
-    noise *= np.exp(-time / 5)
-    x = carrier + noise
-
-    CWT = ContinuousWaveletTransform(fs, x)
-    CWT.plot(0, logspace=False)
-
-    # plt.plot(time, x)
-    # plt.show()
-
-# ================================================================
-# %%
-import torch
-import torch.nn as nn
-import pandas as pd
-from dataset.get_dataset import SignalDatasetCreator
-from pathlib import Path
-
-data_set = "BFDS-Project/Bearing-Fault-Diagnosis-System"  # 数据集huggingface地址
-labels = {"Normal Baseline Data": 0, "Ball": 1, "Inner Race": 2, "Outer Race Centered": 3, "Outer Race Opposite": 4, "Outer Race Orthogonal": 5}  # 标签
-transfer_task = [["CWRU224", "12kDriveEnd"], ["CWRU224", "12kFanEnd"]]  # 迁移方向
-
-
-signal_dataset_creator = SignalDatasetCreator(data_set, labels, transfer_task, stratified_sampling=True)
-dataloaders = {}
-dataloaders["source_train"], dataloaders["source_val"], dataloaders["target_train"], dataloaders["target_val"] = signal_dataset_creator.data_split(
-    64, 0, device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
-)
-
-# %%
-import models
-
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = getattr(models, "ResNet")().to(device)
-bottleneck_layer = nn.Sequential(
-    nn.Linear(model.output_num(), 256),
-    nn.ReLU(inplace=True),
-    nn.Dropout(),
-).to(device)
-classifier_layer = nn.Linear(256, len(labels)).to(device)
-model_all = nn.Sequential(model, bottleneck_layer, classifier_layer).to(device)
-model_all.load_state_dict(torch.load("checkpoint/150_0/149-0.3942-best_model.bin"))  # 加载模型参数
-model_without_head = nn.Sequential(*list(model_all.children())[:-1])
-
-
-# %%
-from sklearn.manifold import TSNE
-import matplotlib.pyplot as plt
-import numpy as np
-from matplotlib.colors import ListedColormap
-
-# 定义一个固定的颜色映射
-num_classes = len(set(label for dataloader in dataloaders.values() for _, labels in dataloader for label in labels.numpy()))
-colors = plt.cm.get_cmap("tab10", num_classes)  # 使用 "tab10" 颜色映射
-cmap = ListedColormap(colors.colors)
-
-
-def plot_tsne(dataloader, title, ax):
-    model_all.eval()
-    with torch.no_grad():
-        for i, (inputs, labels) in enumerate(dataloader):
-            inputs = inputs.to(device)
-            labels = labels.to(device)
-            outputs = model_without_head(inputs)
-            # Collect all points across all batches
-            if i == 0:
-                all_points = outputs.cpu().numpy()
-                all_labels = labels.cpu().numpy()
-            else:
-                all_points = np.concatenate((all_points, outputs.cpu().numpy()), axis=0)
-                all_labels = np.concatenate((all_labels, labels.cpu().numpy()), axis=0)
-
-        # Apply t-SNE to reduce dimensions to 2D
-        tsne = TSNE(n_components=2, random_state=42)
-        reduced_points = tsne.fit_transform(all_points)
-
-        # Plot the reduced points
-        scatter = ax.scatter(reduced_points[:, 0], reduced_points[:, 1], c=all_labels, cmap=cmap, s=10)
-        ax.set_title(title)
-        ax.set_xlabel("Dimension 1")
-        ax.set_ylabel("Dimension 2")
-        return scatter, reduced_points, all_labels
-
-
-# Create a 2x2 subplot
-fig, axes = plt.subplots(2, 2, figsize=(12, 10))
-
-# Plot each dataloader
-# sc1,_ = plot_tsne(dataloaders["source_train"], "Source Train", axes[0, 0])
-sc2, reduced_points2, all_labels2 = plot_tsne(dataloaders["source_val"], "Source Val", axes[0, 1])
-# sc3,_ = plot_tsne(dataloaders["target_train"], "Target Train", axes[1, 0])
-# sc4,_ = plot_tsne(dataloaders["target_val"], "Target Val", axes[1, 1])
-
-# Add a colorbar to the figure
-# cbar = fig.colorbar(sc1, ax=axes, orientation="vertical", fraction=0.02, pad=0.04)
-# cbar.set_label("Labels")
-
-# Adjust layout and show the plot
-plt.tight_layout()
-plt.show()
-
-# %%
-reduced_points2, all_labels2
-
-# %%
-import pandas as pd
-
-df = pd.DataFrame(reduced_points2)
-df["label"] = all_labels2  # 将标签添加为新列
-
-# 保存为 CSV 文件
-df.to_csv("checkpoint/reduced_points_with_labels.csv", index=False)

utils/predict.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+
 import models
 from dataset.get_dataset import get_user_dataset
-import torch.nn.functional as F
 
 
 def predict(model_state_dict, signal_file, args):