sangal-aarushi's picture
Update app.py
04eb49f verified
raw
history blame contribute delete
902 Bytes
import gradio as gr
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
# Load model from Hugging Face Hub
model_name = "wambugu71/crop_leaf_diseases_vit"
model = ViTForImageClassification.from_pretrained(model_name)
processor = ViTImageProcessor.from_pretrained(model_name)
# Prediction function
def predict(image):
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predicted_idx = outputs.logits.argmax(-1).item()
predicted_label = model.config.id2label[predicted_idx]
return f"🌱 Predicted Disease: {predicted_label}"
# Create Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs="text",
title="🌾 AI-Powered Crop Disease Detector",
description="Upload a leaf image of Corn, Potato, Rice, or Wheat to get the predicted disease."
)
iface.launch()