import gradio as gr
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np

# Load the model
model = load_model("shap efficient B0 multi.h5")

# Optionally: Recompile the model to avoid the warning (if needed for evaluation)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# Preprocessing function
def preprocess_image(image):
    # Resize image to the model's input size
    image = image.resize((400, 400))  # Adjust based on your model's input size
    # Normalize the image (if required by your model)
    image = np.array(image) / 255.0  # Normalize pixel values to [0, 1]
    return np.expand_dims(image, axis=0)  # Add batch dimension

# Prediction function
def predict(image):
    # Preprocess the input image
    preprocessed = preprocess_image(image)
    # Make predictions
    predictions = model.predict(preprocessed)
    
    # Define the class labels (you can update these with your actual class names)
    class_labels = ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative DR']
    
    # Create a dictionary with class labels as keys and probabilities as values
    prediction_dict = {class_labels[i]: float(predictions[0][i]) for i in range(len(class_labels))}
    
    # Get the predicted class with the highest probability
    predicted_class = class_labels[np.argmax(predictions, axis=-1)[0]]
    return predicted_class, prediction_dict

# Gradio interface
interface = gr.Interface(
    fn=predict,  # The function to call for predictions
    inputs=gr.Image(type="pil"),  # The input type: image
    outputs=[gr.Text(), gr.Label()]  # The output type: predicted class and probability
)

# Launch the Gradio app
interface.launch(share=True)  # Share the app with a public link