diff --git a/code/src/process_stft.py b/code/src/process_stft.py
index 39ee54a..79ab276 100644
--- a/code/src/process_stft.py
+++ b/code/src/process_stft.py
@@ -28,7 +28,7 @@ Fs = 1024
 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):
+def compute_stft(vibration_data, Fs=Fs, window_size=window_size, hop_size=hop_size):
     frequencies, times, Zxx = stft(
         vibration_data, 
         fs=Fs, 
@@ -39,7 +39,7 @@ def compute_stft(vibration_data):
     stft_magnitude = np.abs(Zxx)
     return stft_magnitude.T  # Transpose to have frequencies as columns
 
-def process_damage_case(damage_num):
+def process_damage_case(damage_num, Fs=Fs, window_size=window_size, hop_size=hop_size, output_dirs=output_dirs):
     damage_folder = os.path.join(damage_base_path, f'DAMAGE_{damage_num}')
     if damage_num == 0:
         # Number of test runs per damage case
@@ -83,7 +83,7 @@ def process_damage_case(damage_num):
             vibration_data = df.iloc[:, 1].values
             
             # Perform STFT
-            stft_magnitude = compute_stft(vibration_data)
+            stft_magnitude = compute_stft(vibration_data, Fs=Fs, window_size=window_size, hop_size=hop_size)
             
             # Convert STFT result to DataFrame
             df_stft = pd.DataFrame(