Update app.py

app.py

@@ -77,23 +77,25 @@ if uploaded_file is not None:
             mime="text/csv"
         )
 
-        # Correlation Heatmap
-        st.subheader("Correlation Heatmap")
-        st.write("Visualizing correlations between numeric features...")
-        heatmap_plot = plot_correlation_heatmap(df_cleaned)
-        st.pyplot(heatmap_plot)
-
-        # Save and download heatmap as PNG
-        heatmap_buffer = io.BytesIO()
-        heatmap_plot.savefig(heatmap_buffer, format="png")
-        heatmap_buffer.seek(0)
+     # Correlation Heatmap
+st.subheader("Correlation Heatmap")
+st.write("Visualizing correlations between numeric features...")
+heatmap_plot = plot_correlation_heatmap(df_cleaned)
+st.pyplot(heatmap_plot)
+
+# Save and download heatmap as PNG
+heatmap_buffer = io.BytesIO()
+fig = plt.gcf()  # Get the current figure
+fig.savefig(heatmap_buffer, format="png")  # Save the figure to the buffer
+heatmap_buffer.seek(0)
+
+st.download_button(
+    label="Download Correlation Heatmap (PNG)",
+    data=heatmap_buffer,
+    file_name="correlation_heatmap.png",
+    mime="image/png"
+)
 
-        st.download_button(
-            label="Download Correlation Heatmap (PNG)",
-            data=heatmap_buffer,
-            file_name="correlation_heatmap.png",
-            mime="image/png"
-        )
 
         # Additional Visualizations
         st.subheader("Additional Visualizations")