Update app.py

app.py

@@ -105,7 +105,7 @@ def plot_correlation(df, column):
     plt.xlabel("Index")
     plt.ylabel(column)
     plt.title(f"Correlation Plot of {column}")
-    st.pyplot()
+    return plt.gcf()  # Return the matplotlib figure
 
 
 def main():
@@ -166,14 +166,19 @@ def main():
                 correlation = warehouse_df[warehouse_column].corr(industry_df[industry_column])
                 st.header("Correlation")
                 st.write(f"The correlation between {warehouse_column} in warehouse item stocks and {industry_column} in industry item stocks is: {correlation}")
+                st.write()
+
 
             # Show correlation plot for each dataset
-            if st.button("Show for each dataset"):
+            if st.button("Show for each dataset
+            
                 st.subheader("Correlation Plot for Warehouse Dataset")
-                plot_correlation(warehouse_df, warehouse_column)
+                warehouse_corr_plot = plot_correlation(warehouse_df, warehouse_column)
+                st.pyplot(warehouse_corr_plot)
                 
                 st.subheader("Correlation Plot for Industry Dataset")
-                plot_correlation(industry_df, industry_column)
+                industry_corr_plot = plot_correlation(industry_df, industry_column)
+                st.pyplot(industry_corr_plot)
 
 if __name__ == "__main__":
     main()