app.py

#!/usr/bin/env python3
"""
Interactive Decision Tree Classifier for the Iris dataset using scikit-learn + Gradio.

Features
- Visualizes the trained decision tree (matplotlib.plot_tree).
- Lets users tweak hyperparameters (criterion, max_depth, min_samples_split).
- Shows accuracy and a full classification report.
- Modular code organization for clarity and reuse.
"""

from __future__ import annotations

import io
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import accuracy_score, classification_report


# -----------------------------
# Data & Config
# -----------------------------

@dataclass
class TrainConfig:
    criterion: str = "gini"                # "gini", "entropy", or "log_loss"
    max_depth: Optional[int] = None        # None means unlimited depth
    min_samples_split: int = 2             # integer >= 2
    test_size: float = 0.25                # test split fraction
    random_state: int = 42                 # reproducibility


def load_iris() -> Tuple[pd.DataFrame, pd.Series, list[str]]:
    """Load the Iris dataset and return X (DataFrame), y (Series), and class names."""
    iris = datasets.load_iris()
    X = pd.DataFrame(iris.data, columns=iris.feature_names)
    y = pd.Series(iris.target, name="target")
    class_names = iris.target_names.tolist()
    return X, y, class_names


# -----------------------------
# Model Pipeline
# -----------------------------

def build_model(cfg: TrainConfig) -> DecisionTreeClassifier:
    """Instantiate a DecisionTreeClassifier from a config."""
    model = DecisionTreeClassifier(
        criterion=cfg.criterion,
        max_depth=cfg.max_depth,
        min_samples_split=cfg.min_samples_split,
        random_state=cfg.random_state,
    )
    return model


def train_and_evaluate(
    cfg: TrainConfig,
) -> Tuple[DecisionTreeClassifier, float, str, pd.DataFrame, pd.Series, list[str]]:
    """
    Train a decision tree and compute accuracy + classification report.
    Returns:
        model, accuracy, report (str), X_test (DataFrame), y_test (Series), class_names (list)
    """
    X, y, class_names = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=cfg.test_size, random_state=cfg.random_state, stratify=y
    )
    model = build_model(cfg)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=class_names)
    return model, acc, report, X_test, y_test, class_names


# -----------------------------
# Visualization
# -----------------------------

def render_tree_image(
    model: DecisionTreeClassifier,
    feature_names: list[str],
    class_names: list[str],
    dpi: int = 120,
) -> np.ndarray:
    """
    Render a matplotlib plot_tree to a PNG image array suitable for display in Gradio.
    """
    fig, ax = plt.subplots(figsize=(12, 8), dpi=dpi)
    plot_tree(
        model,
        feature_names=feature_names,
        class_names=class_names,
        filled=True,
        rounded=True,
        impurity=True,
        proportion=True,
        fontsize=8,
        ax=ax,
    )
    buf = io.BytesIO()
    fig.tight_layout()
    fig.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)

    # Convert buffer to numpy array for gr.Image
    import PIL.Image as Image
    img = Image.open(buf)
    return np.array(img)


# -----------------------------
# Gradio App
# -----------------------------

def inference(
    criterion: str,
    max_depth_enabled: bool,
    max_depth_val: int,
    min_samples_split: int,
    test_size: float,
    random_state: int,
) -> tuple[np.ndarray, str, str]:
    """
    Gradio handler: trains, evaluates, and returns (tree_image, accuracy_text, classification_report).
    """
    cfg = TrainConfig(
        criterion=criterion,
        max_depth=(max_depth_val if max_depth_enabled else None),
        min_samples_split=int(min_samples_split),
        test_size=float(test_size),
        random_state=int(random_state),
    )
    model, acc, report, X_test, y_test, class_names = train_and_evaluate(cfg)
    tree_img = render_tree_image(model, feature_names=X_test.columns.tolist(), class_names=class_names)
    acc_text = f"Accuracy: {acc:.4f}\n\nTest Size: {cfg.test_size} | Random State: {cfg.random_state}\nParams -> criterion={cfg.criterion}, max_depth={cfg.max_depth}, min_samples_split={cfg.min_samples_split}"
    return tree_img, acc_text, report


def build_interface():
    import gradio as gr

    with gr.Blocks(title="Iris Decision Tree (scikit-learn)", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            # 🌸 Iris Decision Tree Classifier
            Tweak hyperparameters and see how the decision tree changes. View accuracy and a full classification report.
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                criterion = gr.Dropdown(
                    label="Criterion",
                    choices=["gini", "entropy", "log_loss"],
                    value="gini",
                    info="Split quality metric"
                )
                max_depth_enabled = gr.Checkbox(
                    label="Limit max_depth?",
                    value=False
                )
                max_depth_val = gr.Slider(
                    label="max_depth (if enabled)",
                    minimum=1,
                    maximum=10,
                    step=1,
                    value=3
                )
                min_samples_split = gr.Slider(
                    label="min_samples_split",
                    minimum=2,
                    maximum=20,
                    step=1,
                    value=2
                )
                test_size = gr.Slider(
                    label="test_size (fraction for test)",
                    minimum=0.1,
                    maximum=0.5,
                    step=0.05,
                    value=0.25
                )
                random_state = gr.Slider(
                    label="random_state",
                    minimum=0,
                    maximum=9999,
                    step=1,
                    value=42
                )
                run_btn = gr.Button("Train & Evaluate", variant="primary")

            with gr.Column(scale=2):
                tree_img = gr.Image(label="Decision Tree", interactive=False)
                accuracy_box = gr.Textbox(label="Accuracy & Parameters", interactive=False, lines=4)
                report_box = gr.Textbox(label="Classification Report", interactive=False, lines=12)

        # Wire events
        inputs = [criterion, max_depth_enabled, max_depth_val, min_samples_split, test_size, random_state]
        outputs = [tree_img, accuracy_box, report_box]

        # Run once on load
        demo.load(fn=inference, inputs=inputs, outputs=outputs)

        # Run on button click
        run_btn.click(fn=inference, inputs=inputs, outputs=outputs)

        gr.Markdown(
            "Tip: Uncheck **Limit max_depth?** to let the tree grow fully (may overfit)."
        )

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)