import gradio as gr
from fastbook import load_learner, PILImage

# Load your model
learn = load_learner('BirdOrForest.pkl')

def predict_image(image):
    img = PILImage.create(image)
    pred, pred_idx, probs = learn.predict(img)
    
    # Convert probabilities to dictionary format (for gr.Label output)
    classes = learn.dls.vocab  # Get class names from the data loader
    return {classes[i]: float(probs[i]) for i in range(len(classes))}

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(),       # Image input
    outputs=gr.Label(num_top_classes=2),  # Label output with probabilities
    examples=[
        ["Examples/1.jpg"], 
        ["Examples/2.jpg"], 
        ["Examples/3.jpg"], 
        ["Examples/4.jpg"]
    ],
    title="Welcome to Bird or Forest Classifier",  # Title of the app
    description="Please upload an image to classify whether it's a Bird or a Forest scene."  # Welcome message
)

# Launch the app without enable_queue
iface.launch()