from fastai.vision.all import load_learner, PILImage
import gradio as gr

# Load your trained model
learn = load_learner("model.pkl")

# Define prediction function
def classify_image(img):
    pred_class, pred_idx, probs = learn.predict(PILImage.create(img))
    return {str(learn.dls.vocab[i]): float(probs[i]) for i in range(len(probs))}

# Gradio Blocks interface without custom CSS
with gr.Blocks() as demo:
    gr.Markdown("# Bike Classifier")
    with gr.Row():
        image = gr.Image(type="pil", label="Upload a Bike Image")
        label = gr.Label(num_top_classes=3, label="Predicted Classes")
    submit_btn = gr.Button("Predict")
    submit_btn.click(fn=classify_image, inputs=image, outputs=label)

# Launch app
if __name__ == "__main__":
    demo.launch()