Update app.py

app.py

@@ -1,5 +1,7 @@
 import streamlit as st
 import pandas as pd
+import matplotlib.pyplot as plt
+import io
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import LabelEncoder
 from sklearn.ensemble import RandomForestClassifier
@@ -10,7 +12,6 @@ from sklearn.tree import DecisionTreeClassifier
 from sklearn.naive_bayes import GaussianNB
 from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
 from tabulate import tabulate
-import io
 
 # Function to convert DataFrame to Excel format
 def to_excel(df):
@@ -20,6 +21,31 @@ def to_excel(df):
     output.seek(0)
     return output
 
+# Function to save table as PNG with bold headings
+def save_table_as_png(df):
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax.axis('tight')
+    ax.axis('off')
+
+    # Create a table from the DataFrame
+    table = ax.table(cellText=df.values, colLabels=df.columns, loc='center', cellLoc='center')
+
+    # Set the font size and bold the header row
+    table.auto_set_font_size(False)
+    table.set_fontsize(10)
+    table.scale(1.2, 1.2)
+
+    # Bold the column headers
+    for (i, j) in zip(range(len(df.columns)), table[0]):
+        table[0, j].set_text_props(weight='bold')  # Make column headers bold
+
+    # Save the table as a PNG image
+    img_path = "/tmp/model_report.png"
+    plt.savefig(img_path, format="png", bbox_inches="tight")
+    plt.close(fig)
+
+    return img_path
+
 # File uploader
 st.title("Model Training with Metrics")
 uploaded_file = st.file_uploader("Choose a CSV file", type=["csv"])
@@ -110,7 +136,7 @@ if uploaded_file is not None:
 
             # Option to download the model performance metrics (Results Table)
             st.download_button(
-                label="Download Model Report",
+                label="Download Model Report (Excel)",
                 data=to_excel(metrics_df),  # The metrics dataframe
                 file_name="model_report.xlsx",
                 mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
@@ -118,8 +144,18 @@ if uploaded_file is not None:
 
             # Option to download the cleaned dataset
             st.download_button(
-                label="Download Cleaned Dataset",
+                label="Download Cleaned Dataset (Excel)",
                 data=to_excel(df),  # The cleaned dataset is 'df'
                 file_name="cleaned_dataset.xlsx",
                 mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
             )
+
+            # Option to download the report as PNG
+            img_path = save_table_as_png(metrics_df)
+            with open(img_path, "rb") as file:
+                st.download_button(
+                    label="Download Model Report (PNG)",
+                    data=file,
+                    file_name="model_report.png",
+                    mime="image/png"
+                )