import os
import numpy as np
import tensorflow as tf
import gradio as gr

# Load model
model_path = 'plantvillage_efficientnet_b0.keras'

if os.path.exists(model_path):
    try:
        model = tf.keras.models.load_model(model_path, compile=False)
        print("Model loaded successfully!")
    except Exception as e:
        print(f"Error loading model: {e}")
        model = None
else:
    print(f"File not found: {model_path}")
    model = None

classes = [
    'Pepper__bell___Bacterial_spot', 'Pepper__bell___healthy', 'Potato___Early_blight',
    'Potato___Late_blight', 'Potato___healthy', 'Tomato_Bacterial_spot',
    'Tomato_Early_blight', 'Tomato_Late_blight', 'Tomato_Leaf_Mold',
    'Tomato_Septoria_leaf_spot', 'Tomato_Spider_mites_Two_spotted_spider_mite',
    'Tomato__Target_Spot', 'Tomato__Tomato_YellowLeaf__Curl_Virus',
    'Tomato__Tomato_mosaic_virus', 'Tomato_healthy'
]

def predict(image):
    if model is None:
        raise gr.Error("Model is not loaded. Check model file path in your Space.")
    
    if image is None:
        return {}

    try:
        # Preprocess input image
        img = tf.image.resize(image, (224, 224))
        img = tf.cast(img, tf.float32)
        img = tf.keras.applications.efficientnet.preprocess_input(img)
        img = tf.expand_dims(img, axis=0)

        # Run inference
        preds = model.predict(img, verbose=0)[0]
        
        # Format predictions dict
        return {classes[i]: float(preds[i]) for i in range(len(classes))}

    except Exception as err:
        print(f"Prediction Error: {err}")
        raise gr.Error(f"Prediction error: {str(err)}")

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Leaf Image"),
    outputs=gr.Label(num_top_classes=3, label="Predictions"),
    title="Plant Disease Detector",
    description="Upload a crop leaf image to identify diseases."
)

if __name__ == '__main__':
    demo.launch()