Update app.py

app.py

@@ -1,56 +1,48 @@
+import gradio as gr
 from fastapi import FastAPI
-from pydantic import BaseModel
 import joblib
 import numpy as np
 
-# 1. Load the trained model
-# Ensure 'iris_model.pkl' is uploaded to the same folder in Hugging Face
+# 1. Load the model you saved from your notebook
 model = joblib.load("iris_model.pkl")
 
-# 2. Define the class mapping based on your notebook encoding
-# In your notebook: Iris-setosa=1, Iris-versicolor=2, Iris-virginica=3
-class_names = {
-    1: "Iris-setosa", 
-    2: "Iris-versicolor", 
-    3: "Iris-virginica"
-}
+# 2. Define the mapping from your notebook's label encoding
+# Iris-setosa: 1, Iris-versicolor: 2, Iris-virginica: 3
+class_names = {1: "Iris-setosa", 2: "Iris-versicolor", 3: "Iris-virginica"}
 
-# 3. Define the input data format using Pydantic
-class IrisInput(BaseModel):
-    SepalLengthCm: float
-    SepalWidthCm: float
-    PetalLengthCm: float
-    PetalWidthCm: float
-
-# 4. Initialize FastAPI
-app = FastAPI()
-
-# 5. Define the Home Route (Health Check)
-@app.get("/")
-def home():
-    return {"message": "Iris Species Prediction API is Live!"}
-
-# 6. Define the Prediction Route
-@app.post("/predict")
-def predict_species(data: IrisInput):
-    # Extract features from the input object
-    features = np.array([[
-        data.SepalLengthCm, 
-        data.SepalWidthCm, 
-        data.PetalLengthCm, 
-        data.PetalWidthCm
-    ]])
+# 3. Create the prediction function for the UI
+def predict_iris(sepal_l, sepal_w, petal_l, petal_w):
+    # Prepare the input array for the Logistic Regression model
+    features = np.array([[sepal_l, sepal_w, petal_l, petal_w]])
     
-    # Make the prediction using the loaded model
+    # Get the numerical prediction
     prediction = model.predict(features)
-    
-    # The model returns an array (e.g., [1]), so we take the first item
     predicted_class = int(prediction[0])
     
-    # Map the number back to the name
-    species_name = class_names.get(predicted_class, "Unknown")
-    
-    return {
-        "predicted_class": predicted_class,
-        "species_name": species_name
-    }
+    # Return the species name
+    return class_names.get(predicted_class, "Unknown")
+
+# 4. Set up the Gradio Interface
+interface = gr.Interface(
+    fn=predict_iris,
+    inputs=[
+        gr.Slider(4.0, 8.0, label="Sepal Length (cm)"),
+        gr.Slider(2.0, 4.5, label="Sepal Width (cm)"),
+        gr.Slider(1.0, 7.0, label="Petal Length (cm)"),
+        gr.Slider(0.1, 2.5, label="Petal Width (cm)"),
+    ],
+    outputs=gr.Textbox(label="Predicted Species"),
+    title="Iris Species Classifier",
+    description="Slide the values to predict if the flower is Setosa, Versicolor, or Virginica.",
+    theme="soft"
+)
+
+# 5. Initialize FastAPI and mount the Gradio UI
+app = FastAPI()
+
+@app.get("/health")
+def health_check():
+    return {"status": "online"}
+
+# This mounts the UI to the root "/" path
+app = gr.mount_gradio_app(app, interface, path="/")