code/src/process_stft.py

@@ -6,7 +6,7 @@ import glob
 import multiprocessing  # Added import for multiprocessing
 
 # Define the base directory where DAMAGE_X folders are located
-damage_base_path = 'D:/thesis/data/converted/raw_B'
+damage_base_path = 'D:/thesis/data/converted/raw'
 
 # Define output directories for each sensor
 output_dirs = {
@@ -25,10 +25,7 @@ window = hann(window_size)
 Fs = 1024
 
 # Number of damage cases (adjust as needed)
-num_damage_cases = 6  # Change to 30 if you have 30 damage cases
-
-# Number of test runs per damage case
-num_test_runs = 5
+num_damage_cases = 0  # Change to 30 if you have 30 damage cases
 
 # Function to perform STFT and return magnitude
 def compute_stft(vibration_data):
@@ -44,7 +41,11 @@ def compute_stft(vibration_data):
 
 def process_damage_case(damage_num):
     damage_folder = os.path.join(damage_base_path, f'DAMAGE_{damage_num}')
-    
+    if damage_num == 0:
+        # Number of test runs per damage case
+        num_test_runs = 120
+    else:
+        num_test_runs = 5
     # Check if the damage folder exists
     if not os.path.isdir(damage_folder):
         print(f"Folder {damage_folder} does not exist. Skipping...")
@@ -79,7 +80,6 @@ def process_damage_case(damage_num):
                 print(f"Unexpected number of columns in {file_path}. Expected 2, got {df.shape[1]}. Skipping...")
                 continue
             
-            # Extract vibration data (assuming the second column is sensor data)
             vibration_data = df.iloc[:, 1].values
             
             # Perform STFT
@@ -90,9 +90,19 @@ def process_damage_case(damage_num):
                 stft_magnitude, 
                 columns=[f"Freq_{freq:.2f}" for freq in np.linspace(0, Fs/2, stft_magnitude.shape[1])]
             )
+            # only inlcude 21 samples vector features for first 45 num_test_runs else include 22 samples vector features
+            if damage_num == 0:
+                print(f"Processing damage_num = 0, test_num = {test_num}")
+                if test_num <= 45:
+                    df_stft = df_stft.iloc[:22, :]
+                    print(f"Reduced df_stft shape (21 samples): {df_stft.shape}")
+                else:
+                    df_stft = df_stft.iloc[:21, :]
+                    print(f"Reduced df_stft shape (22 samples): {df_stft.shape}")
             
             # Append to the aggregated list
             aggregated_stft.append(df_stft)
+            print(sum(df.shape[0] for df in aggregated_stft))
         
         # Concatenate all STFT DataFrames vertically
         if aggregated_stft: