Ace10er's picture
Update app.py
a246748 verified
Raw
History Blame Contribute Delete
4 kB
import torch
import gradio as gr
from PIL import Image
import os
from transformers import AutoImageProcessor, AutoModelForImageClassification
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
from urllib.parse import quote
import requests
# Ensure the model is loaded only once in the deployment environment
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Model loading function
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() # Set the model to evaluation mode
return processor, model
# Load model and processor
processor, model = load_model()
# Define the class names - Update with your actual class names
class_names = ['Healthy', 'Disease_1', 'Disease_2', 'Disease_3'] # Update with your actual classes
# Prediction function
def predict(image):
if isinstance(image, str):
image = Image.open(image).convert("RGB")
# Preprocess the image
inputs = processor(images=image, return_tensors="pt").to(device)
# Inference
with torch.no_grad():
logits = model(**inputs).logits
predicted_class_idx = logits.argmax(-1).item()
predicted_class = class_names[predicted_class_idx]
return predicted_class
# Function to get detailed Wikipedia info
def get_detailed_wikipedia_info(query):
try:
# 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 query, f"Error fetching info: {str(e)}"
# Wrapped predict function
def wrapped_predict(image):
# Step 1: Use the model to predict the disease name
disease_name = predict(image)
# 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"), # 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."
)
interface.launch(share=True, debug=True)
if __name__ == "__main__":
launch_web()