confusion tree add

src/main.py

@@ -1,11 +1,12 @@
 import pandas as pd
 from sklearn.model_selection import train_test_split
+from sklearn.metrics import confusion_matrix
 
 # Import custom modules
 from data import load_conn_data, preprocess_data, feature_scaling
 from evaluation import evaluate_model_comprehensive
 from models import get_all_models
-from visualization import print_model_comparison, plot_feature_importance, print_performance_metrics
+from visualization import print_model_comparison, plot_feature_importance, print_performance_metrics, plot_confusion_matrix
 
 def main():
     # Load data
@@ -37,6 +38,25 @@ def main():
         )
         print_performance_metrics(results, y_test, label_encoder)
         all_results[model_name] = results
+
+        # --- CONFUSION MATRIX ---
+        # Prefer a prediction returned in results, otherwise use the trained model to predict
+        y_pred = results.get('y_pred', None)
+        if y_pred is None:
+            model = results.get('model', None)
+            if model is not None:
+                try:
+                    y_pred = model.predict(X_test_scaled)
+                except Exception:
+                    y_pred = None
+
+        if y_pred is not None:
+            cm = confusion_matrix(y_test, y_pred)
+            if label_encoder is not None and hasattr(label_encoder, "classes_"):
+                class_names = label_encoder.classes_.tolist()
+            else:
+                class_names = [str(c) for c in sorted(set(list(y_test) + list(y_pred)))]
+            plot_confusion_matrix(cm, class_names, model_name, show=True)
 
     # Visualization
     metrics = ['accuracy', 'precision', 'recall', 'f1']
@@ -45,6 +65,7 @@ def main():
     feature_names = X.columns.tolist()
     plot_feature_importance(all_results, feature_names, show=True)
 
+
     return all_results, comparison_df, label_encoder
 
 if __name__ == "__main__":

src/visualization/__init__.py

@@ -1,4 +1,4 @@
 from .console import print_model_comparison, print_performance_metrics
-from .plots import plot_feature_importance
+from .plots import plot_feature_importance, plot_confusion_matrix
 
-__all__ = ['print_model_comparison','plot_feature_importance', 'print_performance_metrics']
+__all__ = ['print_model_comparison','plot_feature_importance', 'print_performance_metrics', 'plot_confusion_matrix']

src/visualization/plots.py

@@ -15,4 +15,24 @@ def plot_feature_importance(all_results, feature_names, show=True):
         plt.xticks(range(len(importances)), [feature_names[i] for i in indices], rotation=45)
         plt.tight_layout()
         if show:
-            plt.show()
+            plt.show()
+def plot_confusion_matrix(cm, class_names, model_name='', show=True):
+    plt.figure(figsize=(8, 6))
+    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    plt.title(f'Confusion Matrix for {model_name}')
+    plt.colorbar()
+    tick_marks = np.arange(len(class_names))
+    plt.xticks(tick_marks, class_names, rotation=45)
+    plt.yticks(tick_marks, class_names)
+
+    thresh = cm.max() / 2.
+    for i, j in np.ndindex(cm.shape):
+        plt.text(j, i, format(cm[i, j], 'd'),
+                 horizontalalignment="center",
+                 color="white" if cm[i, j] > thresh else "black")
+
+    plt.ylabel('True label')
+    plt.xlabel('Predicted label')
+    plt.tight_layout()
+    if show:
+        plt.show()