Update app.py

app.py

@@ -1,19 +1,23 @@
 import pandas as pd
 import numpy as np
-from sklearn.model_selection import train_test_split, GridSearchCV
+from sklearn.model_selection import train_test_split
 from xgboost import XGBClassifier
-from sklearn.metrics import accuracy_score, classification_report
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
 from sklearn.preprocessing import StandardScaler
 from imblearn.over_sampling import SMOTE
 import gradio as gr
+import matplotlib.pyplot as plt
+import seaborn as sns
 import io
+import base64
+from PIL import Image
 import warnings
 warnings.filterwarnings('ignore')
 
 # Function to load and preprocess data
-def load_and_preprocess_data(file_path):
+def load_and_preprocess_data(file):
     try:
-        data = pd.read_csv(file_path)
+        data = pd.read_csv(file.name)
         
         # Convert suits and ranks to numerical values
         suit_order = {'spades': 0, 'hearts': 1, 'clubs': 2, 'diamonds': 3}
@@ -27,9 +31,9 @@ def load_and_preprocess_data(file_path):
         data['Lion Suit Num'] = data['Lion Suit'].map(suit_order)
         data['Lion Rank Num'] = data['Lion Rank'].map(rank_order)
         
-        return data
+        return data, None
     except Exception as e:
-        return f"Error loading data: {str(e)}"
+        return None, f"Error loading data: {str(e)}"
 
 # Feature engineering
 def create_features(data, n_games=3):
@@ -67,17 +71,50 @@ def create_features(data, n_games=3):
                ['suit_mean', 'suit_std', 'rank_mean', 'rank_std'])
     return pd.DataFrame(features, columns=columns)
 
-# Training function
-def train_model(file_path, n_estimators, learning_rate, max_depth, subsample):
-    output = io.StringIO()
+# Function to plot confusion matrix
+def plot_confusion_matrix(y_true, y_pred, title):
+    cm = confusion_matrix(y_true, y_pred)
+    plt.figure(figsize=(6, 4))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
+    plt.title(title)
+    plt.xlabel('Predicted')
+    plt.ylabel('Actual')
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img = Image.open(buf)
+    plt.close()
+    return img
+
+# Function to plot accuracy bar chart
+def plot_accuracy_chart(accuracies):
+    plt.figure(figsize=(8, 5))
+    plt.bar(accuracies.keys(), accuracies.values(), color='skyblue')
+    plt.title('Model Accuracy Comparison')
+    plt.ylabel('Accuracy')
+    plt.xticks(rotation=45)
+    plt.ylim(0, 1)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img = Image.open(buf)
+    plt.close()
+    return img
+
+# Training function with progress tracking
+def train_model(file, n_estimators, learning_rate, max_depth, subsample, progress=gr.Progress()):
+    progress(0, desc="Starting...")
+    results = []
     
     try:
         # Load and preprocess data
-        data = load_and_preprocess_data(file_path)
-        if isinstance(data, str):
-            return data
+        progress(0.1, desc="Loading and preprocessing data...")
+        data, error = load_and_preprocess_data(file)
+        if error:
+            return error, None, None
         
         # Create features
+        progress(0.2, desc="Engineering features...")
         n_games = 3
         features = create_features(data, n_games)
         targets = {
@@ -90,12 +127,18 @@ def train_model(file_path, n_estimators, learning_rate, max_depth, subsample):
         }
         
         # Scale features
+        progress(0.3, desc="Scaling features...")
         scaler = StandardScaler()
         features_scaled = scaler.fit_transform(features)
         features_scaled = pd.DataFrame(features_scaled, columns=features.columns)
         
-        results = []
-        for target_name, target in targets.items():
+        accuracies = {}
+        confusion_matrices = []
+        
+        # Train models
+        for i, (target_name, target) in enumerate(targets.items()):
+            progress(0.4 + (i / len(targets)) * 0.4, desc=f"Training {target_name} model...")
+            
             # Split data
             X_train, X_test, y_train, y_test = train_test_split(
                 features_scaled, target, test_size=0.2, random_state=42
@@ -127,20 +170,30 @@ def train_model(file_path, n_estimators, learning_rate, max_depth, subsample):
             y_pred = model.predict(X_test)
             accuracy = accuracy_score(y_test, y_pred)
             report = classification_report(y_test, y_pred, zero_division=0)
+            accuracies[target_name] = accuracy
             
             results.append(f"**{target_name} Results**\n")
             results.append(f"Accuracy: {accuracy:.2f}\n")
             results.append(f"Classification Report:\n{report}\n")
+            
+            # Generate confusion matrix plot
+            cm_plot = plot_confusion_matrix(y_test, y_pred, f"Confusion Matrix - {target_name}")
+            confusion_matrices.append(cm_plot)
+        
+        progress(0.9, desc="Generating visualizations...")
+        # Generate accuracy bar chart
+        accuracy_plot = plot_accuracy_chart(accuracies)
         
-        return "\n".join(results)
+        progress(1.0, desc="Completed!")
+        return "\n".join(results), accuracy_plot, confusion_matrices
     
     except Exception as e:
-        return f"Error during training: {str(e)}"
+        return f"Error during training: {str(e)}", None, None
 
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Card Game Prediction Model Training")
-    gr.Markdown("Upload the training dataset and configure hyperparameters to train the model.")
+    gr.Markdown("Upload the training dataset and configure hyperparameters to train the model. Track progress and view results with visualizations.")
     
     file_input = gr.File(label="Upload TRAINING_CARD_DATA.csv")
     n_estimators = gr.Slider(50, 300, value=100, step=10, label="Number of Estimators")
@@ -149,12 +202,16 @@ with gr.Blocks() as demo:
     subsample = gr.Slider(0.5, 1.0, value=0.8, step=0.1, label="Subsample")
     
     train_button = gr.Button("Train Model")
-    output = gr.Textbox(label="Training Results")
+    
+    output_text = gr.Textbox(label="Training Results")
+    accuracy_plot = gr.Image(label="Accuracy Comparison")
+    confusion_plots = gr.Gallery(label="Confusion Matrices")
     
     train_button.click(
         fn=train_model,
         inputs=[file_input, n_estimators, learning_rate, max_depth, subsample],
-        outputs=output
+        outputs=[output_text, accuracy_plot, confusion_plots],
+        _js="() => {return {progress: true}}"
     )
 
 demo.launch()