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

# LOAD MODEL (.h5) ✅
model = tf.keras.models.load_model("lite_model.h5", compile=False)

# LABELS
class_names = [
    "Didgeridoo", "Tambourine", "Xylophone", "acordian", "alphorn",
    "bagpipes", "banjo", "bongo drum", "casaba", "castanets",
    "clarinet", "clavichord", "concertina", "drums", "dulcimer",
    "flute", "guiro", "guitar", "harmonica", "harp",
    "marakas", "ocarina", "piano", "saxaphone", "sitar",
    "steel drum", "trombone", "trumpet", "tuba", "violin"
]

IMG_SIZE = (224, 224)

def preprocess_image(image):
    image = image.convert("RGB")
    image = image.resize(IMG_SIZE)
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image

def predict(image):
    img = preprocess_image(image)
    preds = model.predict(img)[0]
    return {class_names[i]: float(preds[i]) for i in range(len(class_names))}

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="🎵 Musical Instrument Classifier",
    description="Upload an image to predict the instrument",
)

interface.launch()