from keras.models import load_model
from keras.preprocessing import image
import numpy as np
import os
import gradio as gr
loaded_model = load_model('diabetic_retinopathy_model.h5')
import gradio as gr
import numpy as np
from tensorflow.keras.preprocessing import image

# Class mapping
class_mapping = {
    0: 'No DR',
    1: 'Mild',
    2: 'Moderate',
    3: 'Severe',
    4: 'Proliferative DR'
}

# URL of the fixed example image to display
example_image_url = "1.jpg"  # Replace with the actual URL

def predict_diabetic_retinopathy(test_image, loaded_model, height=512, width=512):
    # Always return the example image
    try:
        if test_image is None:
            return "No image uploaded. Please upload an image.", example_image_url

        # Ensure the image is in the correct format
        img = image.img_to_array(test_image)
        
        # Resize the image while maintaining the aspect ratio
        img = np.array(image.smart_resize(img, (height, width)))
        
        img_array = np.expand_dims(img, axis=0)
        img_array /= 255.0  # Normalize the image array
        
        # Make predictions
        predictions = loaded_model.predict(img_array)

        # Convert predictions to the corresponding class
        predicted_class = np.argmax(predictions)

        # Return the predicted class and the example image URL
        return f"**Predicted Diabetic Retinopathy Stage:** {class_mapping[predicted_class]}", example_image_url
    
    except Exception as e:
        return f"An error occurred: {str(e)}", example_image_url

# Create the Gradio interface with fixed example images
example_images = [
    "No_DR.png",
    "Mild.png",
    "Moderate.png",
    "Proliferate_DR.png"
]

iface = gr.Interface(
    fn=lambda img: predict_diabetic_retinopathy(img, loaded_model),
    inputs=gr.Image(type="numpy", label="Upload Retina Image"),
    outputs=[gr.Markdown(label="Prediction Result"), gr.Image(value=example_image_url, label="Example Image")],
    title="Diabetic Retinopathy Prediction",
    description="Upload an image of the retina to predict the stage of diabetic retinopathy.",
    theme="default",
    examples=example_images
)

# Launch the interface
iface.launch()