import gradio as gr
import numpy as np
import huggingface_hub
import onnxruntime as ort
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.inception_v3 import preprocess_input
from PIL import Image

# Define the model loading function
def load_model():
    # Download the model from Hugging Face Hub
    path = huggingface_hub.hf_hub_download(repo_id="srijonashraf/maize-leaf-disease-detection", filename="best_model.keras")
    
    # Set session options
    options = ort.SessionOptions()
    options.intra_op_num_threads = 8
    options.inter_op_num_threads = 8
    
    # Create InferenceSession with ONNX Runtime
    session = ort.InferenceSession(path, sess_options=options, providers=["CPUExecutionProvider", "CUDAExecutionProvider"])
    return session

# Load the model
model = load_model()

# Define the image size your model expects
IMG_SIZE = (224, 224)

# Define class names
class_names = [
    'Corn___Common_Rust',
    'Corn___Gray_Leaf_Spot',
    'Corn___Healthy',
    'Corn___Northern_Leaf_Blight',
    'Corn___Northern_Leaf_Spot',
    'Corn___Phaeosphaeria_Leaf_Spot'
]

# Define prediction function
def predict(image):
    img = Image.fromarray(np.uint8(image)).resize(IMG_SIZE)
    img_array = preprocess_input(np.expand_dims(img, axis=0).astype(np.float32))
    
    # Run inference
    predictions = model.run(None, {"input": img_array})[0]
    predicted_class = np.argmax(predictions)
    confidence = np.max(predictions)

    if confidence <= 0.8:
        return "Unknown Object"
    else:
        return {class_names[predicted_class]: float(confidence)}

# Create Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=6),
    title="Maize Leaf Disease Detection",
    description="Upload an image of a maize leaf to classify its disease."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()