diff --git a/main.py b/main.py
index defe61c48e24afd3c3849dc149fdae613af76c96..d192f4c4f51c89a57b582ff11bcb1b15dd3d7342 100644
--- a/main.py
+++ b/main.py
@@ -37,31 +37,32 @@ def extract_features(trainingDataDir, trainingDataSubDirs):
     create_csv_header()
     # Looping over every file inside the subdirectories for feature extraction
     for trainingDataSubDir in trainingDataSubDirs:
-        for audio_file_name in os.listdir(trainingDataDir/f'{trainingDataSubDir}'):
-            if audio_file_name.endswith(".wav"):
-                audio_file = trainingDataDir/f'{trainingDataSubDir}/{audio_file_name}'
-                print("Extracting Features from Directory "+trainingDataSubDir+" and file "+audio_file.name)
-                y, sr = librosa.load(audio_file, mono=True)
-                mfcc_features = librosa.feature.mfcc(y=y, sr=sr,
-                                                     n_mfcc=(constants.MFCC_RANGE_END - constants.MFCC_RANGE_START))
-                to_append = f'{audio_file.name}'
-                for mfcc_segment in mfcc_features:
-                    to_append += f' {np.mean(mfcc_segment)}'
-                if trainingDataSubDir == constants.CAR:
-                    to_append += f' {constants.LIGHT_WEIGHT}'
-                elif trainingDataSubDir == constants.BUS:
-                    to_append += f' {constants.MEDIUM_WEIGHT}'
-                elif trainingDataSubDir == constants.TRUCK:
-                    to_append += f' {constants.HEAVY_WEIGHT}'
-                elif trainingDataSubDir == constants.MOTORCYCLE:
-                    to_append += f' {constants.TWO_WHEELED}'
-                elif trainingDataSubDir == constants.TRAM:
-                    to_append += f' {constants.RAIL_BOUND}'
-
-                file = open(constants.FEATURES_CSV_NAME, 'a', newline='')
-                with file:
-                    writer = csv.writer(file)
-                    writer.writerow(to_append.split())
+        if os.path.isdir(trainingDataSubDir):
+            for audio_file_name in os.listdir(trainingDataDir/f'{trainingDataSubDir}'):
+                if audio_file_name.endswith(".wav"):
+                    audio_file = trainingDataDir/f'{trainingDataSubDir}/{audio_file_name}'
+                    print("Extracting Features from Directory "+trainingDataSubDir+" and file "+audio_file.name)
+                    y, sr = librosa.load(audio_file, mono=True)
+                    mfcc_features = librosa.feature.mfcc(y=y, sr=sr,
+                                                         n_mfcc=(constants.MFCC_RANGE_END - constants.MFCC_RANGE_START))
+                    to_append = f'{audio_file.name}'
+                    for mfcc_segment in mfcc_features:
+                        to_append += f' {np.mean(mfcc_segment)}'
+                    if trainingDataSubDir == constants.CAR:
+                        to_append += f' {constants.LIGHT_WEIGHT}'
+                    elif trainingDataSubDir == constants.BUS:
+                        to_append += f' {constants.MEDIUM_WEIGHT}'
+                    elif trainingDataSubDir == constants.TRUCK:
+                        to_append += f' {constants.HEAVY_WEIGHT}'
+                    elif trainingDataSubDir == constants.MOTORCYCLE:
+                        to_append += f' {constants.TWO_WHEELED}'
+                    elif trainingDataSubDir == constants.TRAM:
+                        to_append += f' {constants.RAIL_BOUND}'
+
+                    file = open(constants.FEATURES_CSV_NAME, 'a', newline='')
+                    with file:
+                        writer = csv.writer(file)
+                        writer.writerow(to_append.split())
 
 
 def preprocessing_csv_data():
@@ -121,17 +122,20 @@ def train_and_save_model(compiledModel, X_train, y_train, X_test, y_test):
                                          callbacks=[tensorboard_callback])
 
     # Saving the trained model to avoid re-training
-    #print(training_history)
     compiledModel.save(constants.TRAINED_MODEL)
     return training_history
 
 
-def predict(X_test, y_test):
+def predict(X_test, y_test, encoder):
     print("Predictions.....")
     final_predictions = np.argmax(compiled_model.predict(X_test), axis=-1)
     target_names = [constants.LIGHT_WEIGHT, constants.MEDIUM_WEIGHT, constants.HEAVY_WEIGHT, constants.TWO_WHEELED,
                     constants.RAIL_BOUND]
+    test_labels = encoder.inverse_transform(y_test)
+    output_labels = encoder.inverse_transform(final_predictions)
     print(classification_report(y_test, final_predictions, target_names=target_names))
+    df = pd.DataFrame({"Actual Output": test_labels, "Predicted Output": output_labels})
+    df.to_csv(r'prediction_report.csv', index=False)
 
 
 # Changing Directory to Training Dataset Folder
@@ -145,4 +149,4 @@ X_input_features = normalize_data(processed_features_data)
 X_train_data, X_test_data, y_train_data, y_test_data = train_test_data_split(X_input_features, target_audio_labels)
 compiled_model = create_and_compile_model()
 model_training_history = train_and_save_model(compiled_model, X_train_data, y_train_data, X_test_data, y_test_data)
-predict(X_test_data, y_test_data)
+predict(X_test_data, y_test_data, encoder_object)