import os from collections import defaultdict import matplotlib.pyplot as plt import numpy as np from scipy import signal # Set backend to torch because tensorflow's backend doesn't work on my gpu os.environ["KERAS_BACKEND"] = "torch" import tensorflow as tf from tensorflow.keras import mixed_precision from dataset.HUSTMotorDatasetParser import MUSTMotorDatasetParser from dataset.datasetParser import DatasetParser # Enable FP16 mixed_precision.set_global_policy('mixed_float16') parser: DatasetParser = MUSTMotorDatasetParser() N_STEPS = 2000 MAX_RECORDS_PER_FILE = 160000 TRAIN_RATIO = 0.7 VALIDATION_RATIO = 0.15 TARGET_FREQUENCY = "ALL" NOISE_STDDEV = 0.05 RAW_SAMPLING_RATE = parser.get_raw_sampling_rate() TARGET_SAMPLES_PER_REV = N_STEPS CLASS_NAMES = parser.get_dataset_motor_problems() CLASS_INDEX = {class_name: index for index, class_name in enumerate(CLASS_NAMES)} def extract_rotation_windows(raw_series: np.ndarray, freq_hz: int, stride_ratio: float = 0.10, target_samples: int = 2000) -> np.ndarray: """ Slices raw contiguous signal into 1-revolution windows with a sliding stride and resamples each window to exactly target_samples via Fourier interpolation. """ raw_samples_per_rev = int(round(RAW_SAMPLING_RATE / freq_hz)) stride_samples = max(1, int(raw_samples_per_rev * stride_ratio)) windows = [] num_samples = len(raw_series) for start_idx in range(0, num_samples - raw_samples_per_rev + 1, stride_samples): end_idx = start_idx + raw_samples_per_rev rev_window = raw_series[start_idx:end_idx] # Shape: (raw_samples_per_rev, 3) # Resample to exactly target_samples (2000) rev_2000 = signal.resample(rev_window, target_samples, axis=0) windows.append(rev_2000) if not windows: return np.empty((0, target_samples, raw_series.shape[1]), dtype=np.float32) return np.array(windows, dtype=np.float32) def parse_freq_hz(freq_str: str) -> int: """Converts frequency strings like '30HZ' to integer 30.""" return int(freq_str.replace("HZ", "")) def load_data( records, target_frequency="ALL", stride_ratio=0.10, train_ratio=0.70, valid_ratio=0.15, noise_std=0.05 ): series_by_file = defaultdict(list) # 1. Group records by file and extract Vibration Channels ONLY (X, Y, Z) for record in records: if target_frequency != "ALL" and record["frequency"] != target_frequency: continue source_key = (record["problem_type"], record["frequency"]) # series_by_file[source_key].append((record["x"], record["y"], record["z"], (record["x"]**2 + record["y"]**2 + record["z"]**2))) series_by_file[source_key].append((record["x"], record["y"], record["z"])) train_inputs, train_labels = [], [] valid_inputs, valid_labels = [], [] test_inputs, test_labels = [], [] # 2. Process each recording file individually for (problem_type, frequency), series_values in series_by_file.items(): raw_series = np.asarray(series_values, dtype=np.float32) label = CLASS_INDEX[problem_type] freq_hz = parse_freq_hz(frequency) # PREVENT DATA LEAKAGE: Chronologically split raw series BEFORE generating windows n_samples = len(raw_series) train_end = int(n_samples * train_ratio) valid_end = int(n_samples * (train_ratio + valid_ratio)) raw_train = raw_series[:train_end] raw_valid = raw_series[train_end:valid_end] raw_test = raw_series[valid_end:] # Extract sliding windows independently on non-overlapping raw splits x_tr = extract_rotation_windows(raw_train, freq_hz, stride_ratio=stride_ratio, target_samples=TARGET_SAMPLES_PER_REV) x_val = extract_rotation_windows(raw_valid, freq_hz, stride_ratio=stride_ratio, target_samples=TARGET_SAMPLES_PER_REV) x_te = extract_rotation_windows(raw_test, freq_hz, stride_ratio=stride_ratio, target_samples=TARGET_SAMPLES_PER_REV) if len(x_tr) > 0: train_inputs.append(x_tr) train_labels.append(np.full(len(x_tr), label, dtype=np.int64)) if len(x_val) > 0: valid_inputs.append(x_val) valid_labels.append(np.full(len(x_val), label, dtype=np.int64)) if len(x_te) > 0: test_inputs.append(x_te) test_labels.append(np.full(len(x_te), label, dtype=np.int64)) if not train_inputs: raise ValueError(f"No records matching target_frequency='{target_frequency}' found.") # 3. Concatenate segments across all files X_train = np.concatenate(train_inputs, axis=0) Y_train = np.concatenate(train_labels, axis=0) X_valid = np.concatenate(valid_inputs, axis=0) Y_valid = np.concatenate(valid_labels, axis=0) X_test = np.concatenate(test_inputs, axis=0) Y_test = np.concatenate(test_labels, axis=0) # 4. Standardize channels based on training set stats channel_mean = X_train.mean(axis=(0, 1), keepdims=True) channel_std = X_train.std(axis=(0, 1), keepdims=True) channel_std[channel_std == 0] = 1.0 X_train = (X_train - channel_mean) / channel_std X_valid = (X_valid - channel_mean) / channel_std X_test = (X_test - channel_mean) / channel_std # 5. DATASET DOUBLING VIA NOISE AUGMENTATION if noise_std > 0: rng = np.random.default_rng(seed=42) noise = rng.normal(loc=0.0, scale=noise_std, size=X_train.shape).astype(np.float32) X_train_noisy = X_train + noise # Concatenate original clean signals + noisy copies X_train = np.concatenate([X_train, X_train_noisy], axis=0) Y_train = np.concatenate([Y_train, Y_train], axis=0) # 6. Shuffle training samples after augmentation rng = np.random.default_rng(seed=42) train_indices = rng.permutation(len(X_train)) X_train, Y_train = X_train[train_indices], Y_train[train_indices] return X_train, Y_train, X_valid, Y_valid, X_test, Y_test records = parser.parse_dataset() print(f"Total records loaded: {len(records)}") X_train, Y_train, X_valid, Y_valid, X_test, Y_test = load_data( records, target_frequency=TARGET_FREQUENCY, noise_std=NOISE_STDDEV, stride_ratio=0.70 ) print(f"X_train shape: {X_train.shape}, Y_train shape: {Y_train.shape} | X_valid shape: {X_valid.shape}, Y_valid shape: {Y_valid.shape} | X_test shape: {X_test.shape}, Y_test shape: {Y_test.shape}") model = tf.keras.Sequential([ tf.keras.Input(shape=(N_STEPS, X_train.shape[2])), tf.keras.layers.GRU(32, return_sequences=False), tf.keras.layers.Dropout(0.4), tf.keras.layers.Dense(len(CLASS_NAMES), activation="softmax"), ]) optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3, clipnorm=1.0) model.compile( loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"], ) train_dataset = ( tf.data.Dataset.from_tensor_slices((X_train, Y_train)) .shuffle(buffer_size=len(X_train)) .batch(256) .prefetch(tf.data.AUTOTUNE) ) valid_dataset = ( tf.data.Dataset.from_tensor_slices((X_valid, Y_valid)) .batch(256) .prefetch(tf.data.AUTOTUNE) ) early_stopping_callback = tf.keras.callbacks.EarlyStopping( monitor="val_loss", patience=50, mode="min", restore_best_weights=True, ) lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau( monitor='val_loss', factor=0.5, patience=12, min_lr=1e-4, verbose=1 ), history = model.fit( train_dataset, epochs=500, validation_data=valid_dataset, callbacks=[early_stopping_callback, lr_scheduler], shuffle=False, ) # Verify on test set test_dataset = tf.data.Dataset.from_tensor_slices((X_test, Y_test)).batch(256) test_loss, test_accuracy = model.evaluate(test_dataset, verbose=0) print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}") # Verification step: Ensure output matrix dimensions match class count predicted_classes = np.argmax(model.predict(X_test, verbose=0), axis=1) confusion_matrix = np.zeros((len(CLASS_NAMES), len(CLASS_NAMES)), dtype=np.int64) for actual_class, predicted_class in zip(Y_test, predicted_classes): confusion_matrix[actual_class, predicted_class] += 1 print("Confusion matrix (rows=actual, columns=predicted):") print(confusion_matrix) plt.figure(figsize=(12, 5)) plt.plot(history.history["accuracy"], label="Train accuracy") plt.plot(history.history["val_accuracy"], label="Validation accuracy") plt.xlabel("Epoch") plt.ylabel("Accuracy") plt.legend() plt.show() plt.figure(figsize=(7, 6)) plt.imshow(confusion_matrix, cmap="Blues") plt.colorbar(label="Number of windows") plt.xticks(range(len(CLASS_NAMES)), CLASS_NAMES) plt.yticks(range(len(CLASS_NAMES)), CLASS_NAMES) plt.xlabel("Predicted class") plt.ylabel("Actual class") plt.title("Fault classification confusion matrix") plt.show()