Update app.py

app.py

@@ -3,60 +3,73 @@ import pandas as pd
 from utils.data_cleaning import handle_missing_values, remove_outliers_iqr, cap_extreme_values
 from utils.visualizations import plot_correlation_heatmap, save_plot_as_png
 from utils.model_training import train_all_models
+import io
 
-# File uploader
+# Streamlit App Title
 st.title("Model Training with Metrics and Correlation Heatmap")
 
 # File uploader with unique keys
-uploaded_file_1 = st.file_uploader("Choose a CSV file for data", type=["csv"], key="file_uploader_1")
-uploaded_file_2 = st.file_uploader("Choose another CSV file for reference", type=["csv"], key="file_uploader_2")
+uploaded_file = st.file_uploader("Upload a CSV file for data analysis", type=["csv"], key="file_uploader_1")
 
-if uploaded_file_1 is not None:
-    st.write("File 1 uploaded successfully!")
-
-if uploaded_file_2 is not None:
-    st.write("File 2 uploaded successfully!")
-    # Show the dataset
-    st.write("Dataset:")
+if uploaded_file is not None:
+    # Read the uploaded file into a DataFrame
+    df = pd.read_csv(uploaded_file)
+    st.write("Dataset Uploaded Successfully!")
     st.dataframe(df)
 
-    # Clean data: Missing values, outliers, and extreme values
-    df = handle_missing_values(df)
-    df = remove_outliers_iqr(df)
-    df = cap_extreme_values(df)
+    # Clean Data: Missing values, outliers, and extreme values
+    st.subheader("Data Cleaning")
+    df_cleaned = handle_missing_values(df)
+    df_cleaned = remove_outliers_iqr(df_cleaned)
+    df_cleaned = cap_extreme_values(df_cleaned)
 
     st.write("Cleaned Dataset:")
-    st.dataframe(df)
+    st.dataframe(df_cleaned)
 
-    # Add clean data download option
+    # Add download option for the cleaned dataset
     st.subheader("Download Cleaned Dataset")
     st.download_button(
         label="Download Cleaned Dataset (CSV)",
-        data=df.to_csv(index=False),
+        data=df_cleaned.to_csv(index=False),
         file_name="cleaned_dataset.csv",
         mime="text/csv"
     )
 
     # Correlation Heatmap
     st.subheader("Correlation Heatmap")
-    corr_plot = plot_correlation_heatmap(df)
+    corr_plot = plot_correlation_heatmap(df_cleaned)
     st.pyplot(corr_plot)
 
     # Save heatmap as PNG
-    heatmap_buf = save_plot_as_png(corr_plot)
+    heatmap_buffer = io.BytesIO()
+    corr_plot.savefig(heatmap_buffer, format='png')
+    heatmap_buffer.seek(0)
+
+    # Download Button for Heatmap
     st.download_button(
         label="Download Correlation Heatmap as PNG",
-        data=heatmap_buf,
+        data=heatmap_buffer,
         file_name="correlation_heatmap.png",
         mime="image/png"
     )
 
-    # Target and features selection
-    target = st.selectbox("Select Target Variable", df.columns)
-    features = [col for col in df.columns if col != target]
-    X = df[features]
-    y = df[target]
+    # Target and Feature Selection
+    st.subheader("Select Target and Features")
+    target = st.selectbox("Select Target Variable", df_cleaned.columns)
+    features = [col for col in df_cleaned.columns if col != target]
+    X = df_cleaned[features]
+    y = df_cleaned[target]
 
-    # Train and evaluate models
+    # Train and Evaluate Models
+    st.subheader("Model Training and Evaluation")
     model_results = train_all_models(X, y)  # Train all models based on data type
+    st.write("Model Training Results:")
     st.dataframe(model_results)
+
+    # Add download option for model results
+    st.download_button(
+        label="Download Model Results (CSV)",
+        data=model_results.to_csv(),
+        file_name="model_results.csv",
+        mime="text/csv"
+    )