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

# Mapping label indices to class labels
labels_list = [class_labels[i]["label"] for i in range(1, 15)]

# 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]
    
    # Find corresponding treatment
    treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
    
    return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"

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