Spaces:
Sleeping
Sleeping
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() |