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

# Download the model from Hugging Face Hub
model_path = hf_hub_download(repo_id="kk20krishna/my-cifar10-model", filename="my-cifar10-model.h5")
model = tf.keras.models.load_model(model_path)

# Define class names
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck']

# Sample images (replace with actual image paths or URLs for deployment)
sample_images = [
    "https://huggingface.co/spaces/kk20krishna/CIFAR10/resolve/main/dog.jpg",  # Replace with actual image paths or URLs
    "https://huggingface.co/spaces/kk20krishna/CIFAR10/resolve/main/cat.jpg",
]


def predict_image(image):
    # Preprocess the image
    image = image.resize((32, 32))
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)

    # Make prediction
    prediction = model.predict(image)
    predicted_class = np.argmax(prediction)
    confidence = prediction[0][predicted_class]

    return class_names[predicted_class], f"{confidence:.2f}"


# Create Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Predicted Class"),
        gr.Textbox(label="Confidence")
    ],
    examples=sample_images,
    title="CIFAR-10 Image Classifier",
    description="Upload or choose a sample image to classify.",
    live=False  # Set this to True if you want live feedback
)

# Launch the app
iface.launch()