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()