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

# Load model
def load_model():
    model_path = "resnet50_cifar10_model.h5"
    if not os.path.exists(model_path):
        url = "https://drive.google.com/uc?id=13KgM2DddlsscFQx4uoYK0lesSE6-DAo3"
        gdown.download(url, model_path, quiet=False)
    model = tf.keras.models.load_model(model_path)
    return model

model = load_model()

class_names = ['Airplane', 'Automobile', 'Bird', 'Cat', 'Deer',
               'Dog', 'Frog', 'Horse', 'Ship', 'Truck']

# Prediction function
def predict_cifar10(image):
    image = image.convert("RGB")
    img = image.resize((32, 32))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    prediction = model.predict(img_array)
    predicted_label = class_names[np.argmax(prediction)]
    confidence = float(np.max(prediction)) * 100

    return {predicted_label: confidence}

# Gradio Interface
iface = gr.Interface(
    fn=predict_cifar10,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="🚀 CIFAR-10 Image Classifier using ResNet50",
    description="Upload an image, and the model will classify it into one of the 10 CIFAR-10 classes."
)

iface.launch()