import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image

# Load your trained model
model = load_model('my_modelled.h5')

# Define the class labels
class_labels = ['DR', 'DME', 'NONE']

def classify_image(img):
    # Preprocess the image for the model
    img = img.resize((256, 256))  # Resize image to match model input
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array / 255.0  # Normalize the image

    # Predict the class
    prediction = model.predict(img_array)
    class_index = np.argmax(prediction)
    class_label = class_labels[class_index]
    
    return class_label

def create_interface():
    # Define the Gradio interface using the updated API
    iface = gr.Interface(
        fn=classify_image,
        inputs=gr.Image(image_mode='RGB'),
        outputs=gr.Label(num_top_classes=3),
        live=True
    )
    return iface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share = True)