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

# 1. Load the model with error handling
try:
    model_path = 'plantvillage_efficientnet_b0.keras'
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file {model_path} not found in Space!")
    
    model = tf.keras.models.load_model(model_path, compile=False)
    print("Model loaded successfully!")
except Exception as e:
    print(f"CRITICAL ERROR LOADING MODEL: {e}")
    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:
        return "Error: Model failed to load. Check Space logs."
    
    img = tf.cast(image, tf.float32)
    img = tf.image.resize(img, (224, 224))
    img = tf.keras.applications.efficientnet.preprocess_input(img)
    img = tf.expand_dims(img, axis=0)

    preds = model.predict(img)[0]
    return {classes[i]: float(preds[i]) for i in range(len(classes))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=3),
    title="Plant Disease Detector",
    description="Identify plant leaf diseases."
)

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