import gradio as gr
from fastai.vision.all import *
from fastcore.all import *
from fastdownload import download_url
from PIL import Image

learn = load_learner('guitar_classifier_model.pkl')
categories = ('acoustic', 'flying v', 'jazzmaster', 'lespaul', 'sg', 'stratocaster', 'telecaster')

def predict(image):
    # Ensure the image is a PIL image (if not already)
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    # Resize the image to (192, 192)
    image = image.resize((192, 192))
    # Predict
    pred, idx, probs = learn.predict(image)
    return dict(zip(categories, map(float, probs)))

# Remove the 'shape' parameter from gr.Image
image_component = gr.Image(type="pil")
label_component = gr.Label()

# Example images (ensure these files are available)
examples = ["acoustic.jpg", "fly.jpeg", "lespaul.jpeg", "sg.jpg", "strat.jpg", "telecaster.jpg"]

iface = gr.Interface(
    fn=predict,
    inputs=image_component,
    outputs=label_component,
    examples=examples,
    title="Guitar Classifier",
    description="Upload a guitar image to get a prediction.",
)

iface.launch()