app.py

import gradio as gr
import numpy as np
import cv2
import tensorflow as tf
import torch
from PIL import Image

# ============== HF Transformers / ViT Model ==============
from transformers import ViTImageProcessor, ViTForImageClassification

# ----------- 1. Load the ViT model & processor ------------
vit_processor = ViTImageProcessor.from_pretrained('wambugu1738/crop_leaf_diseases_vit')
vit_model = ViTForImageClassification.from_pretrained(
    'wambugu1738/crop_leaf_diseases_vit',
    ignore_mismatched_sizes=True
)

vit_label_treatment = {
    "Corn___Common_rust": "Use recommended fungicides and ensure crop rotation.",
    "Corn___Cercospora_leaf_spot": "Apply foliar fungicides; ensure good field sanitation.",
    "Potato___Early_blight": "Apply preventive fungicides; remove infected debris.",
    "Potato___Late_blight": "Use certified seed tubers; fungicide sprays when conditions favor disease.",
    "Rice___Leaf_blight": "Use resistant rice varieties, maintain field hygiene.",
    "Rice___Brown_spot": "For brown spot of rice, employ integrated disease management practices. Use disease-resistant rice varieties and ensure balanced fertilization to prevent excessive nitrogen application. Maintain optimal irrigation to avoid water stress, and at the first signs of infection, apply fungicides such as carbendazim or validamycin following label recommendations. Remove and destroy infected plant debris to reduce inoculum, and consider crop rotation to mitigate disease recurrence.",
    "Rice___Brown_Spot": "For brown spot of rice, employ integrated disease management practices. Use disease-resistant rice varieties and ensure balanced fertilization to prevent excessive nitrogen application. Maintain optimal irrigation to avoid water stress, and at the first signs of infection, apply fungicides such as carbendazim or validamycin following label recommendations. Remove and destroy infected plant debris to reduce inoculum, and consider crop rotation to mitigate disease recurrence.",
    "Wheat___Leaf_rust": "Plant resistant wheat varieties, apply foliar fungicides if severe.",
    # Fallback
    "Unknown": "No specific treatment available."
}

def classify_image_vit(image):
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image.astype('uint8'), 'RGB')
    inputs = vit_processor(images=image, return_tensors="pt")
    outputs = vit_model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()

    # Predicted label
    predicted_label = vit_model.config.id2label.get(predicted_class_idx, "Unknown")
    treatment_text = vit_label_treatment.get(predicted_label, "No specific treatment available.")
    return predicted_label, treatment_text


# ============== TensorFlow Model (plant_model_v5-beta.h5) ==============
# Load the model
keras_model = tf.keras.models.load_model('plant_model_v5-beta.h5')

# Define the class names
class_names = {
    0: 'Apple___Apple_scab',
    1: 'Apple___Black_rot',
    2: 'Apple___Cedar_apple_rust',
    3: 'Apple___healthy',
    4: 'Not a plant',
    5: 'Blueberry___healthy',
    6: 'Cherry___Powdery_mildew',
    7: 'Cherry___healthy',
    8: 'Corn___Cercospora_leaf_spot Gray_leaf_spot',
    9: 'Corn___Common_rust',
    10: 'Corn___Northern_Leaf_Blight',
    11: 'Corn___healthy',
    12: 'Grape___Black_rot',
    13: 'Grape___Esca_(Black_Measles)',
    14: 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)',
    15: 'Grape___healthy',
    16: 'Orange___Haunglongbing_(Citrus_greening)',
    17: 'Peach___Bacterial_spot',
    18: 'Peach___healthy',
    19: 'Pepper,_bell___Bacterial_spot',
    20: 'Pepper,_bell___healthy',
    21: 'Potato___Early_blight',
    22: 'Potato___Late_blight',
    23: 'Potato___healthy',
    24: 'Raspberry___healthy',
    25: 'Soybean___healthy',
    26: 'Squash___Powdery_mildew',
    27: 'Strawberry___Leaf_scorch',
    28: 'Strawberry___healthy',
    29: 'Tomato___Bacterial_spot',
    30: 'Tomato___Early_blight',
    31: 'Tomato___Late_blight',
    32: 'Tomato___Leaf_Mold',
    33: 'Tomato___Septoria_leaf_spot',
    34: 'Tomato___Spider_mites Two-spotted_spider_mite',
    35: 'Tomato___Target_Spot',
    36: 'Tomato___Tomato_Yellow_Leaf_Curl_Virus',
    37: 'Tomato___Tomato_mosaic_virus',
    38: 'Tomato___healthy'
}

# Example dictionary of "treatments" for some classes
keras_treatments = {
    'Apple___Apple_scab': "Remove fallen leaves and prune infected branches. Apply fungicides containing captan or myclobutanil.",
    'Apple___Black_rot': "Prune out dead branches. Spray copper-based fungicide during early fruit development.",
    'Apple___Cedar_apple_rust': "Remove nearby juniper trees. Apply fungicides before bud break.",
    'Apple___healthy': "No action required. The plant is healthy.",
    'Blueberry___healthy': "No action required. The plant is healthy.",
    'Cherry___Powdery_mildew': "Apply sulfur-based fungicide. Ensure good air circulation around the plant.",
    'Cherry___healthy': "No action required. The plant is healthy.",
    'Corn___Cercospora_leaf_spot Gray_leaf_spot': "Rotate crops to avoid build-up of pathogens. Use resistant hybrids and apply foliar fungicides.",
    'Corn___Common_rust': "Plant rust-resistant hybrids. Apply fungicides at the first sign of rust.",
    'Corn___Northern_Leaf_Blight': "Use resistant varieties and apply fungicides when lesions are observed.",
    'Corn___healthy': "No action required. The plant is healthy.",
    'Grape___Black_rot': "Remove and destroy infected leaves and fruits. Apply fungicides containing myclobutanil or captan.",
    'Grape___Esca_(Black_Measles)': "Prune and destroy infected wood. Apply fungicides during the growing season.",
    'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)': "Maintain good air circulation. Spray protective fungicides like mancozeb.",
    'Grape___healthy': "No action required. The plant is healthy.",
    'Orange___Haunglongbing_(Citrus_greening)': "Remove and destroy infected trees. Control psyllid vectors with insecticides.",
    'Peach___Bacterial_spot': "Apply copper-based bactericides. Use resistant varieties and avoid overhead irrigation.",
    'Peach___healthy': "No action required. The plant is healthy.",
    'Pepper,_bell___Bacterial_spot': "Apply copper-based sprays. Use certified seeds and avoid overhead irrigation.",
    'Pepper,_bell___healthy': "No action required. The plant is healthy.",
    'Potato___Early_blight': "Use certified seeds and apply preventative fungicides like chlorothalonil.",
    'Potato___Late_blight': "Plant disease-free tubers and use fungicides containing metalaxyl.",
    'Potato___healthy': "No action required. The plant is healthy.",
    'Raspberry___healthy': "No action required. The plant is healthy.",
    'Soybean___healthy': "No action required. The plant is healthy.",
    'Squash___Powdery_mildew': "Use sulfur-based fungicides and ensure good ventilation.",
    'Strawberry___Leaf_scorch': "Remove infected leaves. Apply fungicides containing myclobutanil.",
    'Strawberry___healthy': "No action required. The plant is healthy.",
    'Tomato___Bacterial_spot': "Apply copper-based sprays. Avoid overhead watering.",
    'Tomato___Early_blight': "Prune infected leaves and apply fungicides containing chlorothalonil or mancozeb.",
    'Tomato___Late_blight': "Remove infected plants. Apply fungicides containing chlorothalonil or metalaxyl.",
    'Tomato___Leaf_Mold': "Ensure good ventilation and apply fungicides like mancozeb.",
    'Tomato___Septoria_leaf_spot': "Remove infected leaves and apply fungicides containing chlorothalonil.",
    'Tomato___Spider_mites Two-spotted_spider_mite': "Spray insecticidal soap or neem oil. Maintain humidity levels.",
    'Tomato___Target_Spot': "Use resistant varieties. Apply fungicides containing chlorothalonil.",
    'Tomato___Tomato_Yellow_Leaf_Curl_Virus': "Remove infected plants. Use resistant varieties and control whitefly vectors.",
    'Tomato___Tomato_mosaic_virus': "Remove infected plants and disinfect tools. Use resistant seed varieties.",
    'Tomato___healthy': "No action required. The plant is healthy.",
    'Unknown': "No specific treatment available."
}

def edge_and_cut(img, threshold1, threshold2):
    emb_img = img.copy()
    edges = cv2.Canny(img, threshold1, threshold2)
    edge_coors = []
    for i in range(edges.shape[0]):
        for j in range(edges.shape[1]):
            if edges[i][j] != 0:
                edge_coors.append((i, j))

    if len(edge_coors) == 0:
        return emb_img

    row_min = edge_coors[np.argsort([coor[0] for coor in edge_coors])[0]][0]
    row_max = edge_coors[np.argsort([coor[0] for coor in edge_coors])[-1]][0]
    col_min = edge_coors[np.argsort([coor[1] for coor in edge_coors])[0]][1]
    col_max = edge_coors[np.argsort([coor[1] for coor in edge_coors])[-1]][1]
    new_img = img[row_min:row_max, col_min:col_max]

    # Simple bounding box in white
    emb_color = np.array([255], dtype=np.uint8)
    emb_img[row_min-10:row_min+10, col_min:col_max] = emb_color
    emb_img[row_max-10:row_max+10, col_min:col_max] = emb_color
    emb_img[row_min:row_max, col_min-10:col_min+10] = emb_color
    emb_img[row_min:row_max, col_max-10:col_max+10] = emb_color

    return emb_img

def classify_and_visualize_keras(image):
    # Preprocess the image
    img_array = tf.image.resize(image, [256, 256])
    img_array = tf.expand_dims(img_array, 0) / 255.0

    # Make a prediction
    prediction = keras_model.predict(img_array)
    predicted_class_idx = tf.argmax(prediction[0], axis=-1).numpy()
    confidence = np.max(prediction[0])
    
    # Obtain the predicted label
    predicted_label = class_names.get(predicted_class_idx, "Unknown")

    if confidence < 0.60:
        class_name = "Uncertain / Not in dataset"
        bounded_image = image
        treatment_text = "No treatment recommendation (uncertain prediction)."
    else:
        class_name = predicted_label
        bounded_image = edge_and_cut(image, 200, 400)
        treatment_text = keras_treatments.get(predicted_label, "No specific treatment available.")
    
    return class_name, float(confidence), bounded_image, treatment_text


# ============== Combined Gradio App ==============
def main_model_selector(model_choice, image):
    """
    Dispatch function based on user choice of model:
      - 'ViT (Corn, Potato, Rice, Wheat)' -> use classify_image_vit
      - 'Keras (Apple, Blueberry, Cherry, etc.)' -> use classify_and_visualize_keras
    """
    if image is None:
        return "No image provided.", None, None, None
    
    if model_choice == "ViT (Corn, Potato, Rice, Wheat)":
        # Return: label, treatment
        predicted_label, treatment_text = classify_image_vit(image)
        # For consistency with the Keras model outputs, 
        # we'll keep placeholders for confidence & bounding box
        return predicted_label, None, image, treatment_text
    
    elif model_choice == "Keras (Apple, Blueberry, Cherry, etc.)":
        # Return: class_name, confidence, bounded_image, treatment_text
        class_name, confidence, bounded_image, treatment_text = classify_and_visualize_keras(image)
        return class_name, confidence, bounded_image, treatment_text
    
    else:
        return "Invalid model choice.", None, None, None


# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# **Plant Disease Detection**")
    gr.Markdown(
        "Select which model you want to use, then upload an image to see the prediction, "
        "confidence (if applicable), bounding box (if applicable), and a suggested treatment."
    )

    with gr.Row():
        model_choice = gr.Radio(
            choices=["ViT (Corn, Potato, Rice, Wheat)", "Keras (Apple, Blueberry, Cherry, etc.)"],
            value="Keras (Apple, Blueberry, Cherry, etc.)",
            label="Select Model"
        )

    with gr.Row():
        inp_image = gr.Image(type="numpy", label="Upload Leaf Image")

    # Outputs
    with gr.Row():
        out_label = gr.Textbox(label="Predicted Class")
        out_confidence = gr.Textbox(label="Confidence (If Available)")
    out_bounded_image = gr.Image(label="Visualization (If Available)")
    out_treatment = gr.Textbox(label="Treatment Recommendation")

    # Button
    btn = gr.Button("Classify")

    # Function binding
    btn.click(
        fn=main_model_selector, 
        inputs=[model_choice, inp_image], 
        outputs=[out_label, out_confidence, out_bounded_image, out_treatment]
    )

    # Provide some example images
    gr.Examples(
        examples=[
            ["Keras (Apple, Blueberry, Cherry, etc.)", "corn.jpg"],
            ["Keras (Apple, Blueberry, Cherry, etc.)", "grot.jpg"],
            ["Keras (Apple, Blueberry, Cherry, etc.)", "Potato___Early_blight.jpg"],
            ["Keras (Apple, Blueberry, Cherry, etc.)", "Tomato___Target_Spot.jpg"],
            ["ViT (Corn, Potato, Rice, Wheat)", "corn.jpg"],
        ],
        inputs=[model_choice, inp_image]
    )

demo.launch(share=True)