Update app.py

app.py

@@ -1,97 +1,104 @@
 import torch
 import gradio as gr
 from PIL import Image
+import os
 from transformers import AutoImageProcessor, AutoModelForImageClassification
-from selenium import webdriver
-from selenium.webdriver.chrome.options import Options
-import time
-
-# ========== Configuration ==========
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
+from sklearn.metrics import confusion_matrix
 
+# Make sure the model is loaded only once in the deployment environment
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ========== Predict using MobileNetV2 from Hugging Face ==========
+# Assuming model path is provided by Hugging Face Spaces
+model_path = "./models"  # Store model weights inside the space directory
+os.makedirs(model_path, exist_ok=True)
 
-def predict(image):
+# Load the trained MobileNetV2 model from Hugging Face Hub
+def load_model():
     processor = AutoImageProcessor.from_pretrained("linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification")
     model = AutoModelForImageClassification.from_pretrained("linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification").to(device)
-    model.eval()
+    model.eval()  # Set the model to evaluation mode
+    return processor, model
+
+processor, model = load_model()
 
+# Define the class names - These should match the labels in the dataset
+class_names = ['Healthy', 'Disease_1', 'Disease_2', 'Disease_3']  # Update with your actual classes
+
+def predict(image):
     if isinstance(image, str):
         image = Image.open(image).convert("RGB")
-
+    
     inputs = processor(images=image, return_tensors="pt").to(device)
-
+    
     with torch.no_grad():
         logits = model(**inputs).logits
-        predicted_class = logits.argmax(-1).item()
-        class_name = model.config.id2label.get(predicted_class, "Unknown")
-
-    return class_name
-
-# ========== Scrape Disease Info ==========
-
-def get_disease_info(disease_name):
-    chrome_options = Options()
-    chrome_options.add_argument("--headless")  # Run headless (no UI)
-    driver = webdriver.Chrome(options=chrome_options)
+    predicted_class_idx = logits.argmax(-1).item()
+    predicted_class = class_names[predicted_class_idx]
+    
+    return predicted_class
 
-    # Perform a Google search for the disease
-    search_query = f"{disease_name} plant disease causes symptoms treatment"
-    search_url = f"https://www.google.com/search?q={search_query}"
+# Detailed Wikipedia info
+from urllib.parse import quote
+import requests
 
+def get_detailed_wikipedia_info(query):
     try:
-        driver.get(search_url)
-        time.sleep(2)  # Wait for JavaScript to load the results
-
-        # Find the first relevant link
-        results = driver.find_elements_by_xpath("//a[@href]")
-
-        # Find the first valid link that seems to be from a trusted source (like Wikipedia)
-        disease_info_url = None
-        for result in results:
-            href = result.get_attribute("href")
-            if href and "wikipedia" in href:
-                disease_info_url = href
-                break
-
-        if not disease_info_url:
-            return f"No relevant information found for '{disease_name}'."
-
-        # Now, scrape the disease details from the Wikipedia link
-        driver.get(disease_info_url)
-        time.sleep(2)
-
-        # Extract text content from the page (just the first few paragraphs)
-        paragraphs = driver.find_elements_by_tag_name("p")
-        disease_info = []
-
-        for paragraph in paragraphs[:5]:  # Limit to 5 paragraphs for brevity
-            disease_info.append(paragraph.text.strip())
-
-        info_summary = "\n".join(disease_info)
-        return info_summary
-
+        # Step 1: Search for the article
+        search_url = f"https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={quote(query)}&format=json"
+        search_response = requests.get(search_url).json()
+        search_results = search_response.get("query", {}).get("search", [])
+        
+        if not search_results:
+            return query, "No Wikipedia pages found."
+        
+        # Get the top result title
+        top_title = search_results[0]["title"]
+        
+        # Step 2: Get full page content (plaintext extract)
+        extract_url = f"https://en.wikipedia.org/w/api.php?action=query&prop=extracts&exintro=&explaintext=&titles={quote(top_title)}&format=json"
+        extract_response = requests.get(extract_url).json()
+        
+        page_data = extract_response.get("query", {}).get("pages", {})
+        page = next(iter(page_data.values()))
+        extract_text = page.get("extract", "No detailed info found.")
+        
+        # Step 3: Format output with Markdown link
+        page_link = f"https://en.wikipedia.org/wiki/{quote(top_title.replace(' ', '_'))}"
+        info = (
+            f"### 🌿 **{top_title}**\n\n"
+            f"{extract_text.strip()}\n\n"
+            f"🔗 [Click here to read more on Wikipedia]({page_link})"
+        )
+        
+        return top_title, info
+    
     except Exception as e:
-        return f"Could not fetch information for '{disease_name}'. Error: {str(e)}"
-
-    finally:
-        driver.quit()
-
-# ========== Web Interface ==========
+        return query, f"Error fetching info: {str(e)}"
 
+# Wrapped predict function
 def wrapped_predict(image):
+    # Step 1: Use your local ML model to predict the disease name
     disease_name = predict(image)
-    disease_info = get_disease_info(disease_name)
+    
+    # Step 2: Use the disease name to query detailed Wikipedia info
+    disease_api_name, disease_info = get_detailed_wikipedia_info(disease_name)
+    
     return disease_name, disease_info
 
+# Gradio Interface for Web Deployment
 def launch_web():
     interface = gr.Interface(
         fn=wrapped_predict,
         inputs=gr.Image(type="pil", label="Upload Leaf Image"),
-        outputs=[gr.Label(label="Predicted Disease"), gr.Textbox(label="Precautions and Solutions")],
+        outputs=[
+            gr.Label(label="Predicted Disease"),  # For displaying the disease name
+            gr.Textbox(label="Disease Information")  # For displaying the detailed information
+        ],
         title="Plant Disease Detection",
-        description="Upload a leaf image to detect the disease and receive actionable solutions fetched from the web."
+        description="Upload a leaf image to detect the disease."
     )
     interface.launch(share=True, debug=True)