predictive-maintenance-ai-test1/main.py

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()