Update app.py

app.py

@@ -1,19 +1,13 @@
 import streamlit as st
 import pandas as pd
-from utils.data_cleaning import preprocess_data, remove_outliers_iqr, cap_extreme_values, convert_string_to_numeric
-from utils.visualizations import (
-    plot_correlation_heatmap,
-    plot_histogram,
-    plot_box_plot,
-    plot_pair_plot,
-    plot_scatter_plot,
-    plot_bar_plot,
-)
-from utils.model_training import train_all_models
-import matplotlib.pyplot as plt
 import seaborn as sns
+import matplotlib.pyplot as plt
 import io
 
+# Import custom functions from your utils
+from utils.data_cleaning import preprocess_data, remove_outliers_iqr, cap_extreme_values, convert_string_to_numeric
+from utils.model_training import train_all_models
+
 # New Function: Combined Histogram and Bar Plot Comparison
 def combined_histogram_barplot(df):
     """
@@ -42,6 +36,71 @@ def combined_histogram_barplot(df):
     plt.tight_layout()
     return fig
 
+# Plotting Functions
+def plot_correlation_heatmap(df):
+    """
+    Plot a correlation heatmap for the numeric columns in the dataframe.
+    """
+    corr = df.corr()
+    fig = plt.figure(figsize=(10, 8))  # Create a new figure object
+    heatmap = sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f", linewidths=0.5)
+    plt.title("Correlation Heatmap")
+    return fig  # Return the figure object
+
+def save_figure_as_png(fig):
+    """
+    Save the given figure as a PNG file to a BytesIO buffer.
+    """
+    buffer = io.BytesIO()
+    fig.savefig(buffer, format="png")  # Save the figure to the buffer
+    buffer.seek(0)  # Reset the buffer's position to the beginning
+    return buffer
+
+def plot_histogram(df, column):
+    """
+    Plot a histogram for a specific column in the dataframe.
+    """
+    plt.figure(figsize=(8, 6))
+    sns.histplot(df[column], kde=True, bins=30, color="skyblue")
+    plt.title(f"Histogram of {column}")
+    plt.xlabel(column)
+    plt.ylabel("Frequency")
+    return plt.gcf()
+
+def plot_box_plot(df, column):
+    """
+    Plot a box plot for a specific column in the dataframe.
+    """
+    plt.figure(figsize=(8, 6))
+    sns.boxplot(x=df[column])
+    plt.title(f"Box Plot of {column}")
+    return plt.gcf()
+
+def plot_pair_plot(df):
+    """
+    Plot a pair plot for numeric columns in the dataframe.
+    """
+    numeric_columns = df.select_dtypes(include=['float64', 'int64']).columns
+    return sns.pairplot(df[numeric_columns])
+
+def plot_scatter_plot(df, x_col, y_col):
+    """
+    Plot a scatter plot between two numeric columns.
+    """
+    plt.figure(figsize=(8, 6))
+    sns.scatterplot(x=df[x_col], y=df[y_col], color="green")
+    plt.title(f"Scatter Plot between {x_col} and {y_col}")
+    return plt.gcf()
+
+def plot_bar_plot(df, column):
+    """
+    Plot a bar plot for a categorical column.
+    """
+    plt.figure(figsize=(8, 6))
+    sns.countplot(x=df[column])
+    plt.title(f"Bar Plot of {column}")
+    return plt.gcf()
+
 # Streamlit App Title
 st.title("Data Analysis, Model Training, and Visualization")
 
@@ -82,12 +141,10 @@ if uploaded_file is not None:
         st.subheader("Correlation Heatmap")
         st.write("Visualizing correlations between numeric features...")
         heatmap_fig = plot_correlation_heatmap(df_cleaned)
-        st.pyplot(heatmap_fig)
+        st.pyplot(heatmap_fig)  # Display the heatmap using Streamlit
 
         # Save and download heatmap as PNG
-        heatmap_buffer = io.BytesIO()
-        heatmap_fig.savefig(heatmap_buffer, format="png")  # Save the figure to the buffer
-        heatmap_buffer.seek(0)  # Reset buffer position to the beginning
+        heatmap_buffer = save_figure_as_png(heatmap_fig)  # Save the figure to buffer
 
         st.download_button(
             label="Download Correlation Heatmap (PNG)",