add normalized confusion matrices

src/plots.py

@@ -8,13 +8,16 @@ from . import config
 
 
 def plot_confusion_matrix(y_true, y_pred, classes, model_name):
-    """Generates and saves a confusion matrix heatmap."""
+    """Generates and saves both Raw and Normalized confusion matrix heatmaps."""
+
+    # 1. Calculate Raw Matrix
     cm = confusion_matrix(y_true, y_pred)
 
+    # Plot Raw Counts
     plt.figure(figsize=(10, 8))
     sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                 xticklabels=classes, yticklabels=classes)
-    plt.title(f"{model_name} Confusion Matrix")
+    plt.title(f"{model_name} Confusion Matrix (Counts)")
     plt.ylabel('True Label')
     plt.xlabel('Predicted Label')
     plt.tight_layout()
@@ -23,6 +26,26 @@ def plot_confusion_matrix(y_true, y_pred, classes, model_name):
     plt.savefig(os.path.join(config.MODEL_DIR, filename))
     plt.close()
 
+    # 2. Calculate Normalized Matrix
+    # Divide each row element by the sum of that row (True Label count)
+    cm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+    # Replace NaN with 0 (safe guard for empty classes)
+    cm_norm = np.nan_to_num(cm_norm)
+
+    # Plot Normalized Percentages
+    plt.figure(figsize=(10, 8))
+    # Use fmt='.2f' to show 2 decimal places (e.g., 0.95)
+    sns.heatmap(cm_norm, annot=True, fmt='.2f', cmap='Greens',
+                xticklabels=classes, yticklabels=classes)
+    plt.title(f"{model_name} Confusion Matrix (Normalized)")
+    plt.ylabel('True Label')
+    plt.xlabel('Predicted Label')
+    plt.tight_layout()
+
+    filename_norm = f"{model_name.lower()}_confusion_matrix_normalized.png"
+    plt.savefig(os.path.join(config.MODEL_DIR, filename_norm))
+    plt.close()
+
 
 def plot_multiclass_roc(model, X_test, y_test, classes, model_name):
     """Generates and saves a Multi-class ROC Curve."""