import gradio as gr
import torch
from transformers import ViTForImageClassification, ViTFeatureExtractor
from PIL import Image

# Load model and feature extractor
model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")

# Define class labels and treatment advice
class_labels = {
    0:"Corn___Common_Rust": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants.",
    1:"Corn___Gray_Leaf_Spot": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed.",
    2:"Corn___Healthy": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests.",
    3:"Corn___Northern_Leaf_Blight": "Remove and destroy infected plant debris, apply fungicides, and rotate crops.",
    4:"Rice___Brown_Spot": "Use resistant varieties, improve field drainage, and apply fungicides if necessary.",
    5:"Rice___Healthy": "Maintain proper irrigation, fertilization, and pest control measures.",
    6:"Rice___Leaf_Blast": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management.",
    7:"Rice___Neck_Blast": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear.",
    8:"Wheat___Brown_Rust": "Apply fungicides and practice crop rotation with non-host crops.",
    9:"Wheat___Healthy": "Continue with good management practices, including proper fertilization and weed control.",
    10:"Wheat___Yellow_Rust": "Use resistant varieties, apply fungicides, and rotate crops.",
    11:"Sugarcane__Red_Rot": "Plant resistant varieties and ensure good drainage.",
    12:"Sugarcane__Healthy": "Maintain healthy soil conditions and proper irrigation.",
    13:"Sugarcane__Bacterial Blight": "Use disease-free planting material, practice crop rotation, and destroy infected plants."
}

# Mapping label indices to class labels
labels_list = list(class_labels.keys())

# Inference function
def predict(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class_idx = outputs.logits.argmax(-1).item()
    predicted_label = labels_list[predicted_class_idx]
    treatment_advice = class_labels[predicted_label]
    
    return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"

# Create Gradio Interface
interface = gr.Interface(fn=predict, inputs="image", outputs="text")
interface.launch()