import gradio as gr
import pathlib
import random
import tensorflow as tf
from PIL import Image
from timeit import default_timer as timer
from keras.models import load_model
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Load the model
model = load_model('MyResNet101Model_final.keras')

# Define label mappings based on your dataset
label2id = {'Hispa': 0, 'Bercak Coklat': 1, 'Blast Daun': 2, 'Sehat': 3}
class_names = list(label2id.keys())

# Function to predict an image
def predict(img):
    start = timer()

    img = img.resize((224, 224))  # Ensure size matches the model's expected input
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)  # Create batch dimension

    # Normalize image
    img_array /= 255.0

    # Prediction
    predictions = model.predict(img_array)
    pred_prob = tf.nn.softmax(predictions[0])
    pred_dict = {class_names[i]: float(pred_prob[i]) for i in range(len(class_names))}

    pred_time = round(timer() - start, 5)

    return pred_dict, pred_time

# Example images for demonstration
example_paths = list(pathlib.Path('examples').glob("*/*.jpg"))
example_list = [[str(filepath)] for filepath in random.sample(example_paths, k=4)]

# Set up the Gradio interface
title = 'Klasifikasi Penyakit Daun Padi'
description = 'Upload gambar daun padi untuk mengklasifikasikan penyakitnya.'

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type='pil'),
    outputs=[
        gr.Label(num_top_classes=4, label='Prediksi'),
        gr.Number(label="Waktu prediksi (detik)")
    ],
    description=description,
    title=title,
    allow_flagging='never',
    examples=example_list
)

if __name__ == "__main__":
    demo.launch(debug=True)