import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np

# Load the pre-trained Fashion MNIST model
model_path = "kia_fashion_mnist_keras_model.h5"
model = tf.keras.models.load_model(model_path)

# Labels for Fashion MNIST
labels = [
    'T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
    'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'
]

def predict_fashion(image):
    # Convert image to grayscale if not already and resize
    image = Image.fromarray(image).convert('L').resize((28, 28))
    # Normalize the image
    image = np.array(image) / 255.0
    # Reshape for model input
    image = image.reshape(1, 28, 28, 1)

    # Make a prediction
    predictions = model.predict(image)
    prediction = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)

    # Prepare the output
    result = f"Predicted category: {labels[prediction]} with confidence: {confidence:.2f}"
    return result

# Create Gradio interface
input_image = gr.Image(image_mode='L')
output_label = gr.Label()
interface = gr.Interface(fn=predict_fashion,
                         inputs=input_image,
                         outputs=output_label,
                         examples=["images/0.png", "images/1.png", "images/2.png", "images/3.png"],
                         title="Fashion MNIST Classifier",
                         description="Drag and drop an image or select an example below to predict the fashion category.")

# Launch the interface
interface.launch()