Update app.py

app.py

@@ -12,9 +12,10 @@ from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_sc
 import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
+from io import BytesIO
 
 # File uploader
-st.title("Model Training with Metrics")
+st.title("Model Training with Metrics and Correlation Heatmap")
 uploaded_file = st.file_uploader("Choose a CSV file", type=["csv"])
 
 if uploaded_file is not None:
@@ -86,13 +87,27 @@ if uploaded_file is not None:
     # Correlation Heatmap
     st.subheader("Correlation Heatmap")
     corr = df.corr()
-    plt.figure(figsize=(10, 6))
-    sns.heatmap(corr, annot=True, cmap='coolwarm', fmt='.2f')
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f", cbar=True)
     st.pyplot(plt)
-
-    # Correlation Metrics
-    st.subheader("Correlation Metrics")
-    st.dataframe(corr)
+    
+    # Save heatmap as PNG
+    buf = BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    st.download_button(
+        label="Download Correlation Heatmap as PNG",
+        data=buf,
+        file_name="correlation_heatmap.png",
+        mime="image/png"
+    )
+    
+    # Highlight highly correlated pairs
+    st.subheader("Highly Correlated Features")
+    high_corr = corr.abs().unstack().sort_values(ascending=False).drop_duplicates()
+    high_corr = high_corr[high_corr >= 0.8]
+    high_corr_df = high_corr[high_corr.index.get_level_values(0) != high_corr.index.get_level_values(1)]
+    st.write(high_corr_df)
 
     target = st.selectbox("Select Target Variable", df.columns)
     features = [col for col in df.columns if col != target]
@@ -129,24 +144,16 @@ if uploaded_file is not None:
         st.subheader("Classification Model Performance Metrics")
         st.dataframe(metrics_df)
 
-        # Download as CSV
-        st.download_button(
-            label="Download Classification Report as CSV",
-            data=metrics_df.to_csv(index=False),
-            file_name="classification_report.csv",
-            mime="text/csv"
-        )
-
-        # Download as PNG
+        # Save metrics as PNG
         fig, ax = plt.subplots()
-        ax.axis('off')
-        table = ax.table(cellText=metrics_df.values, colLabels=metrics_df.columns, loc='center', cellLoc='center')
-        table.auto_set_font_size(False)
-        table.set_fontsize(10)
-        plt.savefig("classification_report.png")
+        sns.barplot(data=metrics_df, x="Model", y="Accuracy", ax=ax)
+        ax.set_title("Classification Model Performance")
+        buf = BytesIO()
+        fig.savefig(buf, format="png")
+        buf.seek(0)
         st.download_button(
             label="Download Classification Report as PNG",
-            data=open("classification_report.png", "rb"),
+            data=buf,
             file_name="classification_report.png",
             mime="image/png"
         )
@@ -179,15 +186,16 @@ if uploaded_file is not None:
         st.subheader("Regression Model Performance Metrics")
         st.dataframe(regression_metrics_df)
 
-        # Download as CSV
+        # Save metrics as PNG
+        fig, ax = plt.subplots()
+        sns.barplot(data=regression_metrics_df, x="Model", y="R² Score", ax=ax)
+        ax.set_title("Regression Model Performance")
+        buf = BytesIO()
+        fig.savefig(buf, format="png")
+        buf.seek(0)
         st.download_button(
-            label="Download Regression Report as CSV",
-            data=regression_metrics_df.to_csv(index=False),
-            file_name="regression_report.csv",
-            mime="text/csv"
+            label="Download Regression Report as PNG",
+            data=buf,
+            file_name="regression_report.png",
+            mime="image/png"
         )
-
-        # Download as PNG
-        fig, ax = plt.subplots()
-        ax.axis('off')
-        table = ax.table(cellText=regression_metrics_df.values