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

# Load crop disease model (ViT)
model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")

# Load fallback model (ResNet50 for general image classification)
fallback_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
fallback_feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")

# Define class labels with treatment advice
class_labels = {
    1: {"label": "Stage Corn Common Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
    2: {"label": "Stage Corn Gray Leaf Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
    3: {"label": "Stage Safe Corn Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
    4: {"label": "Stage Corn Northern Leaf Blight", "treatment": "Remove and destroy infected plant debris, apply fungicides, and rotate crops."},
    5: {"label": "Stage Rice Brown Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
    6: {"label": "Stage Safe Rice Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
    7: {"label": "Stage Rice Leaf Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
    8: {"label": "Stage Rice Neck Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
    9: {"label": "Stage Sugarcane Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."},
    10: {"label": "Stage Safe Sugarcane Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
    11: {"label": "Stage Sugarcane Red Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
    12: {"label": "Stage Wheat Brown Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
    13: {"label": "Stage Safe Wheat Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
    14: {"label": "Stage Wheat Yellow Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
}

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

# Confidence threshold for ViT model
CONFIDENCE_THRESHOLD = 0.5

# Inference function with fuzzy confidence
def predict(image):
    # First, use the crop disease model (ViT)
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        confidences = torch.softmax(logits, dim=-1)
        predicted_class_idx = logits.argmax(-1).item()
        confidence = confidences[0, predicted_class_idx].item()

    # If confidence is below the threshold, use the fallback model
    if confidence < CONFIDENCE_THRESHOLD:
        inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs_fallback = fallback_model(**inputs_fallback)
            logits_fallback = outputs_fallback.logits
            confidences_fallback = torch.softmax(logits_fallback, dim=-1)
            predicted_class_idx_fallback = logits_fallback.argmax(-1).item()
            fallback_confidence = confidences_fallback[0, predicted_class_idx_fallback].item()

        # Get the fallback prediction label
        fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]

        return (
            f"Low confidence in ViT model ({confidence * 100:.2f}%).\n"
            f"ResNet-50 predicts: {fallback_label} ({fallback_confidence * 100:.2f}%).\n\n"
            "If this does not match your input, please try another image."
        )

    # If confidence is above the threshold, return the ViT prediction and treatment advice
    predicted_label = labels_list[predicted_class_idx]
    treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
    return (
        f"Disease: {predicted_label} ({confidence * 100:.2f}%)\n\n"
        f"Treatment Advice: {treatment_advice}"
    )

# Create Gradio Interface
interface = gr.Interface(
    fn=predict, 
    inputs="image",
    outputs="text",
    title="Crop Disease Detection",
    description="Upload an image of a crop plant to detect diseases. If confidence is low, ResNet-50 will classify the image."
)

if __name__ == "__main__":
    interface.launch()