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

# Load the model
model = tf.keras.models.load_model('model.keras')

# Define the class labels
class_labels = {
    0: "Buildings",
    1: "Forest",
    2: "Glacier",
    3: "Mountain",
    4: "Sea",
    5: "Street"
}

# Prediction function
def classify_image(img):
    # Resize the image to the input size expected by your model
    img = img.resize((150, 150))  # Replace 150 with your model's input size

    # Convert the image to a numpy array and preprocess it
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array / 255.0  # Normalize if your model expects normalized inputs

    # Make a prediction
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions, axis=1)

    # Get the class label from the predicted class index
    predicted_label = class_labels.get(predicted_class[0], "Unknown")

    # Return the predicted label
    return f"Predicted class: {predicted_label}"

# Gradio interface
interface = gr.Interface(
    fn=classify_image,  # Function to call
    inputs=gr.Image(type="pil"),  # Input type (image)
    outputs="text",  # Output type (text)
    title="CNN Image Classification",
    description="Upload an image, and the model will classify it into one of the following classes: Buildings, Forest, Glacier, Mountain, Sea, Street."
)

# Launch the interface
interface.launch()