Update app.py

app.py

@@ -1,54 +1,32 @@
 import gradio as gr
-import torch
-from transformers import ViTFeatureExtractor, ViTForImageClassification
+from transformers import pipeline
 from PIL import Image
-import requests
 
-# Load a pre-trained Vision Transformer model from Hugging Face
-model_name = "nateraw/vit-base-patch16-224-in21k"  # Replace with the model you've trained or a similar model
-model = ViTForImageClassification.from_pretrained(model_name)
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+# Load the Hugging Face image classification pipeline with EfficientNetB0
+# This model is generic for plant disease, so if you have a specific tobacco disease model, replace it accordingly
+classifier = pipeline("image-classification", model="nateraw/efficientnet-b0")
 
-# Define the disease labels (placeholders)
-labels = {
-    0: "Healthy",
-    1: "Tobacco Mosaic Virus",
-    2: "Brown Spot",
-    3: "Frog Eye Leaf Spot",
-    4: "Other"
-}
-
-# Define a function for disease detection
-def detect_disease(image):
-    # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Run the model
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_class = logits.argmax().item()
+def identify_disease(image):
+    # Use the classifier to predict the disease
+    predictions = classifier(image)
     
-    # Get disease name from label dictionary
-    disease_name = labels.get(predicted_class, "Unknown Disease")
+    # Format the output to include disease name and confidence score
+    results = [{"Disease": pred["label"], "Confidence": f"{pred['score'] * 100:.2f}%"} for pred in predictions]
     
-    return f"Disease Detected: {disease_name}"
-
-# Build Gradio Interface
-title = "Tobacco Leaf Disease Detection"
-description = """
-Upload or take a real-time picture of a tobacco leaf, and the app will detect the disease (if any).
-"""
+    # Return the uploaded image along with the results
+    return image, results
 
-# Create Gradio interface with camera and real-time processing
-iface = gr.Interface(
-    fn=detect_disease,
-    inputs=gr.Image(source="camera", type="pil", tool="editor"),
-    outputs="text",
-    title=title,
-    description=description,
-    live=True  # Enables real-time processing
+# Define Gradio interface
+interface = gr.Interface(
+    fn=identify_disease, 
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=[
+        gr.outputs.Image(type="pil", label="Uploaded Image"),
+        gr.outputs.Dataframe(type="pandas", label="Predictions")
+    ],
+    title="Tobacco Plant Disease Identifier",
+    description="Upload an image of a tobacco plant, and this tool will identify the disease along with the confidence score."
 )
 
-# Launch Gradio app
-iface.launch()
+# Launch the app
+interface.launch()