Update app.py

app.py

@@ -578,9 +578,9 @@ def calculate_phishing_score(text):
     return probability
 
 def generate_confusion_matrix_plot(_cm):
-    """Generate confusion matrix plot"""
+    """Generate confusion matrix plot - optimized for performance"""
     plt.style.use('dark_background')
-    fig, ax = plt.subplots(figsize=(5, 4), facecolor='#1a1a1a')
+    fig, ax = plt.subplots(figsize=(5, 4), facecolor='#1a1a1a', dpi=80)
     ax.set_facecolor('#1a1a1a')
     
     sns.heatmap(
@@ -591,27 +591,34 @@ def generate_confusion_matrix_plot(_cm):
         cmap="YlOrBr",
         cbar=True,
         square=True,
-        annot_kws={"size": 16, "weight": "bold", "color": "#0f0f0f"},
-        linewidths=2,
+        annot_kws={"size": 14, "weight": "bold", "color": "#0f0f0f"},
+        linewidths=1,
         linecolor='#0f0f0f',
-        cbar_kws={'label': 'Count', 'shrink': 0.8}
+        cbar_kws={'label': 'Count', 'shrink': 0.8},
+        vmin=0,
+        vmax=_cm.max()
     )
     
-    ax.set_xlabel("Predicted", fontsize=11, fontweight='bold', color='#FFD700')
-    ax.set_ylabel("Actual", fontsize=11, fontweight='bold', color='#FFD700')
-    ax.set_xticklabels(["Safe", "Phishing"], fontsize=10, color='#e5e7eb')
-    ax.set_yticklabels(["Safe", "Phishing"], fontsize=10, rotation=0, color='#e5e7eb')
-    ax.set_title("Confusion Matrix", fontsize=13, fontweight='bold', pad=12, color='#FFD700')
+    ax.set_xlabel("Predicted", fontsize=10, fontweight='bold', color='#FFD700')
+    ax.set_ylabel("Actual", fontsize=10, fontweight='bold', color='#FFD700')
+    ax.set_xticklabels(["Safe", "Phishing"], fontsize=9, color='#e5e7eb')
+    ax.set_yticklabels(["Safe", "Phishing"], fontsize=9, rotation=0, color='#e5e7eb')
+    ax.set_title("Confusion Matrix", fontsize=12, fontweight='bold', pad=10, color='#FFD700')
     
-    cbar = ax.collections[0].colorbar
-    cbar.ax.yaxis.set_tick_params(color='#e5e7eb')
-    plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#e5e7eb')
+    try:
+        cbar = ax.collections[0].colorbar
+        if cbar:
+            cbar.ax.yaxis.set_tick_params(color='#e5e7eb')
+            plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#e5e7eb')
+    except:
+        pass
     
     plt.tight_layout()
     buf = io.BytesIO()
-    plt.savefig(buf, format='png', facecolor='#1a1a1a', dpi=100)
+    plt.savefig(buf, format='png', facecolor='#1a1a1a', dpi=80, bbox_inches='tight')
     buf.seek(0)
     plt.close(fig)
+    plt.close('all')
     
     return buf