import gradio as gr
import tensorflow as tf
import numpy as np

# Load your trained TensorFlow model
model = tf.keras.models.load_model('best_model_weights.h5')  # Load your saved model

# Define a function to make predictions
def classify_image(input_image):
    # Preprocess the input image (resize and normalize)
    input_image = tf.image.resize(input_image, (224, 224))  # Make sure to match your model's input size
    input_image = (input_image / 255.0)  # Normalize to [0, 1]
    input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension

    # Make a prediction using your model
    prediction = model.predict(input_image)

    # Assuming your model outputs probabilities for two classes, you can return the class with the highest probability
    class_index = np.argmax(prediction)
    class_labels = ["Class 0", "Class 1"]  # Replace with your actual class labels
    predicted_class = class_labels[class_index]

    return predicted_class

# Create a Gradio interface
input_interface = gr.inputs.Image()  # Gradio input component for image
output_interface = gr.outputs.Text()  # Gradio output component for text

# Create the Gradio app
app = gr.Interface(
    fn=classify_image,
    inputs=input_interface,
    outputs=output_interface,
    live=True,
    title="Image Classifier",
    description="Classify images using a trained model."
)

# Start the Gradio app
app.launch()