decision-tree / app.py
yashkavaiya's picture
Create app.py
4e0eaa3 verified
import gradio as gr
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.tree import DecisionTreeClassifier, plot_tree
from io import BytesIO
from PIL import Image
# Load available datasets
dataset_names = ["iris", "wine", "breast_cancer", "digits"]
datasets_dict = {
"iris": datasets.load_iris(),
"wine": datasets.load_wine(),
"breast_cancer": datasets.load_breast_cancer(),
"digits": datasets.load_digits(),
}
# Define the function to visualize the decision tree
def visualize_decision_tree(dataset_name, max_depth, min_samples_split, min_samples_leaf, max_features, criterion, splitter, max_leaf_nodes, random_state):
dataset = datasets_dict[dataset_name]
X, y = dataset.data, dataset.target
clf = DecisionTreeClassifier(max_depth=max_depth, min_samples_split=min_samples_split, min_samples_leaf=min_samples_leaf, max_features=max_features, criterion=criterion, splitter=splitter, max_leaf_nodes=max_leaf_nodes, random_state=random_state)
clf.fit(X, y)
fig, ax = plt.subplots(figsize=(10, 8))
plot_tree(clf, feature_names=dataset.feature_names, class_names=dataset.target_names, filled=True, ax=ax)
buf = BytesIO()
fig.savefig(buf, format='png')
buf.seek(0)
image_data = buf.getvalue()
image = Image.open(BytesIO(image_data))
return image
# Define the hyperparameters and their ranges
max_depth_range = [None, 2, 3, 4, 5, 6, 7, 8, 9, 10]
min_samples_split_range = [2, 3, 4, 5, 6, 7, 8, 9, 10]
min_samples_leaf_range = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
max_features_range = [None, 'sqrt', 'log2', 0.1, 0.2, 0.3, 0.4, 0.5]
criterion_range = ['gini', 'entropy']
splitter_range = ['best', 'random']
max_leaf_nodes_range = [None, 2, 3, 4, 5, 6, 7, 8, 9, 10]
random_state_range = [None, 42, 100, 200, 300, 400, 500, 600]
# Create the Gradio interface
dataset_dropdown = gr.components.Dropdown(choices=dataset_names, label="Dataset", value="iris")
max_depth_dropdown = gr.components.Dropdown(choices=max_depth_range, label="Max Depth", value=None)
min_samples_split_dropdown = gr.components.Dropdown(choices=min_samples_split_range, label="Min Samples Split", value=2)
min_samples_leaf_dropdown = gr.components.Dropdown(choices=min_samples_leaf_range, label="Min Samples Leaf", value=1)
max_features_dropdown = gr.components.Dropdown(choices=max_features_range, label="Max Features", value=None)
criterion_dropdown = gr.components.Dropdown(choices=criterion_range, label="Criterion", value="gini")
splitter_dropdown = gr.components.Dropdown(choices=splitter_range, label="Splitter", value="best")
max_leaf_nodes_dropdown = gr.components.Dropdown(choices=max_leaf_nodes_range, label="Max Leaf Nodes", value=None)
random_state_dropdown = gr.components.Dropdown(choices=random_state_range, label="Random State", value=None)
iface = gr.Interface(
fn=visualize_decision_tree,
inputs=[dataset_dropdown, max_depth_dropdown, min_samples_split_dropdown, min_samples_leaf_dropdown, max_features_dropdown, criterion_dropdown, splitter_dropdown, max_leaf_nodes_dropdown, random_state_dropdown],
outputs=gr.Image(type="pil"),
title="Decision Tree Visualization",
description="Visualize a decision tree classifier on various datasets by adjusting hyperparameters."
)
iface.launch()