Update app.py

app.py

@@ -49,25 +49,30 @@ def predict_disease(image):
     Preprocesses the image, performs prediction, and returns the predicted class with confidence score.
     """
     # Resize and preprocess image
-    img = tf.image.resize(image, [256, 256]) 
+    # Resize and preprocess the image
+    img = tf.image.resize(image, [256, 256])
     img = tf.expand_dims(img, axis=0)  # Add batch dimension
 
     # Predict probabilities
-    prediction = model.predict(img)
+    prediction = model.predict(img)[0]  # Extract first element (batch of size 1)
     labels = ['Early Blight', 'Late Blight', 'Healthy']
-    
-    # Get predicted class and confidence
-    predicted_class = labels[prediction.argmax()]
-    confidence = prediction.max()
-    return f"Prediction: {predicted_class}, Confidence: {confidence:.2%}"
+
+    # Create a dictionary of class probabilities
+    class_probs = {label: float(prob) for label, prob in zip(labels, prediction)}
 
 # ... (Rest of your code) ...
 
 iface = gr.Interface(
     fn=predict_disease,
-    inputs=gr.Image(),  # Remove shape argument
-    outputs="text",
-    title="Potato Disease Classification",
-    description="Upload an image of a potato leaf to classify its disease.",
+    inputs=gr.Image(type="numpy", tool="editor"),  # Allow image editing
+    outputs=gr.Label(num_top_classes=3),          # Display all three probabilities
+    title="Potato Disease Classifier 🍟🥔",
+    description=(
+        "Upload an image of a potato leaf, and I'll classify it as one of the following:\n"
+        "- **Early Blight** 🌱\n"
+        "- **Late Blight** 🌧️\n"
+        "- **Healthy** 💚\n\n"
+        "Check out the probability bars for more details!"
+    ),
 )
-iface.launch(share=True)
+iface.launch(share=True)