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

# Load the model with weights
model_path = "my_model.h5"  # Replace with the correct path to your model
model = load_model(model_path)

# Define the image height and width
image_height = 224
image_width = 224  # Adjust to match the input size of your model

# Map the class indices to class names based on your dataset
class_names = {
    0: "Class 0",
    1: "Class 1",
    2: "Class 2",
    3: "Class 3",
    4: "Class 4",
}

# Define the prediction function
def predict_image(img_array):
    # Preprocess the input image
    img_array = img_array.reshape((1, image_height, image_width, 3))  # Reshape to (1, 224, 224, 3)
    #img_array = img_array.astype(np.float32) / 255.0  # Normalize pixel values

    # Perform additional preprocessing if needed, e.g., center-cropping

    # Make predictions using the loaded model
    predictions = model.predict(img_array)

    # Get the predicted class label
    predicted_class_index = np.argmax(predictions)
    predicted_class = class_names[predicted_class_index]

    # Return the predicted class name for Gradio to display in the output
    return predicted_class

# Create Gradio Interface
iface = gr.Interface(
    fn=predict_image,
    inputs="image",  # Gradio will automatically create an image uploader
    outputs="text",  # Display the predicted class name as text
)

# Launch the Gradio Interface
iface.launch(share=True)