Update utils/visualizations.py

utils/visualizations.py

@@ -1,127 +1,57 @@
 import seaborn as sns
 import matplotlib.pyplot as plt
-import pandas as pd
-import streamlit as st
 
-# Correlation Heatmap
 def plot_correlation_heatmap(df):
     """
-    Plot a correlation heatmap for numeric columns in the dataframe.
+    Plot a correlation heatmap for the numeric columns in the dataframe.
     """
-    # Select only numeric columns
-    numeric_df = df.select_dtypes(include=['float64', 'int64'])
-    
-    # Compute the correlation matrix
-    corr = numeric_df.corr()
-
-    # Plot the heatmap
+    corr = df.corr()
     plt.figure(figsize=(10, 8))
-    sns.heatmap(corr, annot=True, cmap='coolwarm', fmt='.2f', linewidths=0.5)
+    heatmap = sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f", linewidths=0.5)
     plt.title("Correlation Heatmap")
-    
-    return plt
-
-# Save plot as PNG image
-def save_plot_as_png(plot, filename="plot.png"):
-    """
-    Save a given plot as a PNG file.
-    """
-    plot.savefig(filename, format='png')
-    return filename
+    return heatmap
 
-# Distribution Plot (Histogram)
 def plot_histogram(df, column):
     """
-    Plot a histogram for a specified column in the dataframe.
+    Plot a histogram for a specific column in the dataframe.
     """
     plt.figure(figsize=(8, 6))
-    sns.histplot(df[column], kde=True, color='blue', bins=20)
-    plt.title(f"Distribution of {column}")
+    sns.histplot(df[column], kde=True, bins=30, color="skyblue")
+    plt.title(f"Histogram of {column}")
     plt.xlabel(column)
     plt.ylabel("Frequency")
-    return plt
+    return plt.gcf()
 
-# Box Plot (For Outliers)
 def plot_box_plot(df, column):
     """
-    Plot a box plot for a specified column to visualize outliers.
+    Plot a box plot for a specific column in the dataframe.
     """
     plt.figure(figsize=(8, 6))
-    sns.boxplot(x=df[column], color='orange')
+    sns.boxplot(x=df[column])
     plt.title(f"Box Plot of {column}")
-    plt.xlabel(column)
-    return plt
+    return plt.gcf()
 
-# Pair Plot (For Visualizing Relationships Between Features)
 def plot_pair_plot(df):
     """
-    Plot a pair plot to visualize relationships between numeric columns in the dataframe.
+    Plot a pair plot for numeric columns in the dataframe.
     """
-    numeric_df = df.select_dtypes(include=['float64', 'int64'])
-    pair_plot = sns.pairplot(numeric_df)
-    pair_plot.fig.set_size_inches(10, 8)
-    return pair_plot
+    numeric_columns = df.select_dtypes(include=['float64', 'int64']).columns
+    return sns.pairplot(df[numeric_columns])
 
-# Scatter Plot (For Visualizing Relationship Between Two Features)
-def plot_scatter_plot(df, x_column, y_column):
+def plot_scatter_plot(df, x_col, y_col):
     """
-    Plot a scatter plot to visualize the relationship between two features.
+    Plot a scatter plot between two numeric columns.
     """
     plt.figure(figsize=(8, 6))
-    sns.scatterplot(x=df[x_column], y=df[y_column], color='green')
-    plt.title(f"Scatter Plot between {x_column} and {y_column}")
-    plt.xlabel(x_column)
-    plt.ylabel(y_column)
-    return plt
+    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()
 
-# Bar Plot (For Comparing Categorical Data)
 def plot_bar_plot(df, column):
     """
     Plot a bar plot for a categorical column.
     """
     plt.figure(figsize=(8, 6))
-    sns.countplot(x=df[column], palette='viridis')
+    sns.countplot(x=df[column])
     plt.title(f"Bar Plot of {column}")
-    plt.xlabel(column)
-    plt.ylabel("Count")
-    return plt
-
-    st.subheader("Correlation Heatmap")
-    if st.button("Generate Correlation Heatmap"):
-        heatmap_plot = plot_correlation_heatmap(df)
-        st.pyplot(heatmap_plot)
-
-    st.subheader("Histogram")
-    selected_column_hist = st.selectbox("Select Column for Histogram", df.select_dtypes(include=['float64', 'int64']).columns)
-    if st.button("Generate Histogram"):
-        hist_plot = plot_histogram(df, selected_column_hist)
-        st.pyplot(hist_plot)
-
-    st.subheader("Box Plot")
-    selected_column_box = st.selectbox("Select Column for Box Plot", df.select_dtypes(include=['float64', 'int64']).columns, key="box")
-    if st.button("Generate Box Plot"):
-        box_plot = plot_box_plot(df, selected_column_box)
-        st.pyplot(box_plot)
-
-    st.subheader("Pair Plot")
-    if st.button("Generate Pair Plot"):
-        pair_plot = plot_pair_plot(df)
-        st.pyplot(pair_plot)
-
-    st.subheader("Scatter Plot")
-    numeric_columns = df.select_dtypes(include=['float64', 'int64']).columns
-    x_column = st.selectbox("Select X-Axis Column", numeric_columns, key="scatter_x")
-    y_column = st.selectbox("Select Y-Axis Column", numeric_columns, key="scatter_y")
-    if st.button("Generate Scatter Plot"):
-        scatter_plot = plot_scatter_plot(df, x_column, y_column)
-        st.pyplot(scatter_plot)
-
-    st.subheader("Bar Plot")
-    categorical_columns = df.select_dtypes(include=['object']).columns
-    if not categorical_columns.empty:
-        selected_column_bar = st.selectbox("Select Column for Bar Plot", categorical_columns)
-        if st.button("Generate Bar Plot"):
-            bar_plot = plot_bar_plot(df, selected_column_bar)
-            st.pyplot(bar_plot)
-    else:
-        st.info("No categorical columns available for bar plot.")
+    return plt.gcf()