Update app.py

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
-import re
-import pandas as pd
-import numpy as np
 from sklearn.datasets import load_iris, load_wine, load_breast_cancer
 from sklearn.model_selection import train_test_split
 from sklearn.tree import DecisionTreeClassifier, plot_tree
 import matplotlib.pyplot as plt
 import pandas as pd
+from PIL import Image
 import io
 
-# Load datasets into a dictionary for easy selection
+# Datasets
 DATASETS = {
     "Iris": load_iris(),
     "Wine": load_wine(),
@@ -17,15 +15,12 @@ DATASETS = {
 }
 
 def train_and_plot(dataset_name, criterion, splitter, max_depth):
-    # Load chosen dataset
     data = DATASETS[dataset_name]
     X = pd.DataFrame(data.data, columns=data.feature_names)
     y = data.target
 
-    # Split into train/test
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
-    # Create Decision Tree model
     clf = DecisionTreeClassifier(
         criterion=criterion,
         splitter=splitter,
@@ -33,24 +28,21 @@ def train_and_plot(dataset_name, criterion, splitter, max_depth):
         random_state=42
     )
     clf.fit(X_train, y_train)
-
-    # Calculate accuracy
     accuracy = clf.score(X_test, y_test)
 
-    # Plot decision tree
+    # Plot tree
     fig, ax = plt.subplots(figsize=(12, 8))
     plot_tree(clf, feature_names=data.feature_names, class_names=data.target_names, filled=True, ax=ax)
     plt.tight_layout()
 
-    # Save plot to a file-like object
     buf = io.BytesIO()
     plt.savefig(buf, format="png")
     buf.seek(0)
+    img = Image.open(buf)  # Convert to PIL
     plt.close(fig)
 
-    return f"Model Accuracy: {accuracy:.2%}", buf
+    return f"Model Accuracy: {accuracy:.2%}", img
 
-# Create Gradio interface
 demo = gr.Interface(
     fn=train_and_plot,
     inputs=[
@@ -61,7 +53,7 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Textbox(label="Accuracy"),
-        gr.Image(type="file", label="Decision Tree Plot")
+        gr.Image(type="pil", label="Decision Tree Plot")
     ],
     title="Interactive Decision Tree Classifier",
     description="Adjust the parameters and see how the Decision Tree changes."