Update app.py

app.py

@@ -1,16 +1,43 @@
 import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
-from sklearn.ensemble import RandomForestRegressor
-from sklearn.metrics import mean_squared_error
-from mpl_toolkits.mplot3d import Axes3D
+from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
+from sklearn.metrics import mean_squared_error, accuracy_score
 
 # ------------------------------------------------
-# Random Forest 2D
+# Helpers
 # ------------------------------------------------
-def rf_2d(n_points, noise, n_estimators, max_depth):
+
+def generate_2d_regression(n_points, noise):
     X = np.linspace(0, 10, n_points).reshape(-1, 1)
     y = 2.5 * X.flatten() + 5 + np.random.randn(n_points) * noise
+    return X, y
+
+
+def generate_3d_regression(n_points, noise):
+    x1 = np.linspace(0, 10, n_points)
+    x2 = np.linspace(0, 10, n_points)
+    X1, X2 = np.meshgrid(x1, x2)
+    Z = 3 * X1 + 2 * X2 + 10 + np.random.randn(*X1.shape) * noise
+
+    X_flat = np.column_stack((X1.ravel(), X2.ravel()))
+    Z_flat = Z.ravel()
+    return X1, X2, X_flat, Z, Z_flat
+
+
+def generate_classification(n_points, noise):
+    X = np.random.randn(n_points, 2)
+    y = (X[:, 0]**2 + X[:, 1] > 0.5).astype(int)
+    X += np.random.randn(*X.shape) * noise * 0.1
+    return X, y
+
+
+# ------------------------------------------------
+# 2D Random Forest Regression + Trees + Importance
+# ------------------------------------------------
+
+def rf_2d(n_points, noise, n_estimators, max_depth, show_trees):
+    X, y = generate_2d_regression(n_points, noise)
 
     model = RandomForestRegressor(
         n_estimators=n_estimators,
@@ -23,26 +50,31 @@ def rf_2d(n_points, noise, n_estimators, max_depth):
     mse = mean_squared_error(y, y_pred)
 
     fig, ax = plt.subplots(figsize=(5, 4))
-    ax.scatter(X, y, s=20, color="orange", label="Data")
-    ax.plot(X, y_pred, color="blue", linewidth=2, label="RF Prediction")
+    ax.scatter(X, y, s=20, label="Data")
+    ax.plot(X, y_pred, linewidth=2, label="RF Avg")
+
+    # Individual trees
+    if show_trees:
+        for tree in model.estimators_[:5]:
+            ax.plot(X, tree.predict(X), linewidth=1, alpha=0.5)
+
     ax.set_title("2D Random Forest Regression")
     ax.legend()
 
-    return fig, f"MSE: {mse:.4f}"
+    # Feature importance
+    imp_fig, imp_ax = plt.subplots(figsize=(4, 3))
+    imp_ax.bar(["x"], model.feature_importances_)
+    imp_ax.set_title("Feature Importance")
+
+    return fig, imp_fig, f"MSE: {mse:.4f}"
 
 
 # ------------------------------------------------
-# Random Forest 3D
+# 3D Random Forest Regression + Rotation
 # ------------------------------------------------
-def rf_3d(n_points, noise, n_estimators, max_depth):
-    x1 = np.linspace(0, 10, n_points)
-    x2 = np.linspace(0, 10, n_points)
-    X1, X2 = np.meshgrid(x1, x2)
-
-    Z = 3 * X1 + 2 * X2 + 10 + np.random.randn(*X1.shape) * noise
 
-    X_flat = np.column_stack((X1.ravel(), X2.ravel()))
-    Z_flat = Z.ravel()
+def rf_3d(n_points, noise, n_estimators, max_depth, rotate):
+    X1, X2, X_flat, Z, Z_flat = generate_3d_regression(n_points, noise)
 
     model = RandomForestRegressor(
         n_estimators=n_estimators,
@@ -54,66 +86,95 @@ def rf_3d(n_points, noise, n_estimators, max_depth):
     Z_pred = model.predict(X_flat).reshape(X1.shape)
     mse = mean_squared_error(Z_flat, Z_pred.ravel())
 
-    fig = plt.figure(figsize=(5, 4))
-    ax = fig.add_subplot(111, projection="3d")
-
-    idx = np.random.choice(len(Z_flat), 400, replace=False)
-    ax.scatter(
-        X_flat[idx, 0],
-        X_flat[idx, 1],
-        Z_flat[idx],
-        s=8,
-        alpha=0.3,
-        color="orange"
-    )
+    frames = []
+
+    angles = range(0, 360, 10) if rotate else [45]
+
+    for angle in angles:
+        fig = plt.figure(figsize=(5, 4))
+        ax = fig.add_subplot(111, projection="3d")
+
+        idx = np.random.choice(len(Z_flat), min(400, len(Z_flat)), replace=False)
+        ax.scatter(X_flat[idx, 0], X_flat[idx, 1], Z_flat[idx], s=8, alpha=0.3)
 
-    ax.plot_surface(X1, X2, Z_pred, alpha=0.7, color="blue")
-    ax.set_title("3D Random Forest Surface")
+        ax.plot_surface(X1, X2, Z_pred, alpha=0.7)
+        ax.view_init(elev=25, azim=angle)
+        ax.set_title("3D Random Forest Surface")
 
-    return fig, f"MSE: {mse:.4f}"
+        frames.append(fig)
+
+    return frames, f"MSE: {mse:.4f}"
 
 
 # ------------------------------------------------
-# Gradio UI
+# Classification Version
 # ------------------------------------------------
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🌲 Random Forest Regression Visualizer  
-        Interactive **2D & 3D Random Forest** playground  
-        """
+
+def rf_classification(n_points, noise, n_estimators, max_depth):
+    X, y = generate_classification(n_points, noise)
+
+    model = RandomForestClassifier(
+        n_estimators=n_estimators,
+        max_depth=max_depth,
+        random_state=42
     )
+    model.fit(X, y)
+
+    y_pred = model.predict(X)
+    acc = accuracy_score(y, y_pred)
 
-    with gr.Row():
-        mode = gr.Radio(
-            ["2D Regression", "3D Regression"],
-            value="2D Regression",
-            label="Mode"
-        )
+    fig, ax = plt.subplots(figsize=(5, 4))
+    scatter = ax.scatter(X[:, 0], X[:, 1], c=y_pred, s=20)
+    ax.set_title("Random Forest Classification")
+
+    # Feature importance
+    imp_fig, imp_ax = plt.subplots(figsize=(4, 3))
+    imp_ax.bar(["x1", "x2"], model.feature_importances_)
+    imp_ax.set_title("Feature Importance")
 
-    with gr.Row():
-        n_points = gr.Slider(20, 200, value=80, step=10, label="Data Points")
-        noise = gr.Slider(0.0, 5.0, value=1.0, label="Noise Level")
+    return fig, imp_fig, f"Accuracy: {acc:.4f}"
 
-    with gr.Row():
-        n_estimators = gr.Slider(10, 200, value=50, step=10, label="Trees")
-        max_depth = gr.Slider(2, 20, value=8, step=1, label="Max Depth")
 
-    run_btn = gr.Button("🌲 Train Random Forest")
+# ------------------------------------------------
+# Gradio UI (Auto‑run on input change)
+# ------------------------------------------------
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🌲 Random Forest Visualizer (Full Interactive)")
+
+    mode = gr.Radio(
+        ["2D Regression", "3D Regression", "Classification"],
+        value="2D Regression",
+        label="Mode"
+    )
+
+    n_points = gr.Slider(50, 200, value=100, step=10, label="Data Points")
+    noise = gr.Slider(0.0, 5.0, value=1.0, label="Noise")
+    n_estimators = gr.Slider(10, 150, value=50, step=10, label="Trees")
+    max_depth = gr.Slider(2, 20, value=8, step=1, label="Max Depth")
+
+    show_trees = gr.Checkbox(label="Show Individual Trees (2D)", value=False)
+    rotate = gr.Checkbox(label="Rotate 3D", value=False)
 
     plot = gr.Plot()
+    plot2 = gr.Plot()
     metric = gr.Markdown()
 
-    def run(mode, n_points, noise, n_estimators, max_depth):
+    def run(mode, n_points, noise, n_estimators, max_depth, show_trees, rotate):
         if mode == "2D Regression":
-            return rf_2d(n_points, noise, n_estimators, max_depth)
+            return (*rf_2d(n_points, noise, n_estimators, max_depth, show_trees))
+        elif mode == "3D Regression":
+            frames, text = rf_3d(n_points, noise, n_estimators, max_depth, rotate)
+            return frames[0], None, text
         else:
-            return rf_3d(n_points, noise, n_estimators, max_depth)
+            return (*rf_classification(n_points, noise, n_estimators, max_depth))
+
+    inputs = [mode, n_points, noise, n_estimators, max_depth, show_trees, rotate]
+
+    for inp in inputs:
+        inp.change(run, inputs=inputs, outputs=[plot, plot2, metric])
+
+    demo.load(run, inputs=inputs, outputs=[plot, plot2, metric])
 
-    run_btn.click(
-        run,
-        inputs=[mode, n_points, noise, n_estimators, max_depth],
-        outputs=[plot, metric]
-    )
 
 demo.launch()