Update app.py

app.py

@@ -11,6 +11,7 @@ from utils.visualizations import (
 )
 from utils.model_training import train_all_models
 import matplotlib.pyplot as plt
+import seaborn as sns
 import io
 
 # New Function: Combined Histogram and Bar Plot Comparison
@@ -80,14 +81,13 @@ if uploaded_file is not None:
         # Correlation Heatmap
         st.subheader("Correlation Heatmap")
         st.write("Visualizing correlations between numeric features...")
-        heatmap_plot = plot_correlation_heatmap(df_cleaned)
-        st.pyplot(heatmap_plot)
+        heatmap_fig = plot_correlation_heatmap(df_cleaned)
+        st.pyplot(heatmap_fig)
 
         # Save and download heatmap as PNG
         heatmap_buffer = io.BytesIO()
-        fig = plt.gcf()  # Get the current figure object
-        fig.savefig(heatmap_buffer, format="png")  # Save the figure to the buffer
-        heatmap_buffer.seek(0)
+        heatmap_fig.savefig(heatmap_buffer, format="png")  # Save the figure to the buffer
+        heatmap_buffer.seek(0)  # Reset buffer position to the beginning
 
         st.download_button(
             label="Download Correlation Heatmap (PNG)",
@@ -172,7 +172,6 @@ if uploaded_file is not None:
                 file_name="model_results.csv",
                 mime="text/csv"
             )
-
     except Exception as e:
         st.error(f"An error occurred: {e}")
 else: