diff --git a/code/src/ml/inference.py b/code/src/inference.py
similarity index 97%
rename from code/src/ml/inference.py
rename to code/src/inference.py
index 4ed103a..9d59f32 100644
--- a/code/src/ml/inference.py
+++ b/code/src/inference.py
@@ -1,11 +1,12 @@
 from joblib import load
 import pandas as pd
-from src.data_preprocessing import *
-from src.process_stft import compute_stft
+from data_preprocessing import *
+from process_stft import compute_stft
 from typing import List, Tuple
 from sklearn.base import BaseEstimator
 import json
 
+
 def probability_damage(pred: Tuple[np.ndarray, np.ndarray], model_classes: BaseEstimator, percentage=False) -> Dict[str, int]:
     """
     Process the prediction output to return unique labels and their counts.
@@ -132,16 +133,16 @@ def heatmap(result, damage_classes: list[int] = [1, 2, 3, 4, 5, 6]):
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
     import json
-    from scipy.interpolate import UnivariateSpline
+    # from scipy.interpolate import UnivariateSpline
     
 
     result = inference(
         "D:/thesis/models/Sensor A/SVM with StandardScaler and PCA.joblib",
         "D:/thesis/models/Sensor B/SVM with StandardScaler and PCA.joblib",
-        "D:/thesis/data/dataset_B/zzzBD19.TXT"
+        "D:/thesis/data/dataset_B/zzzBU.TXT"
     )
 
-    # heatmap(result)
+    heatmap(result)
     # Convert all keys to strings before dumping to JSON
     # result_with_string_keys = convert_keys_to_strings(result)
     # print(json.dumps(result_with_string_keys, indent=4))