from fastai.vision.all import *
import gradio as gr
import os

os.makedirs("uploads", exist_ok=True)

def classify_image(img):
    learn = load_learner("model.pkl")
    _, _, probs = learn.predict(img)
    # Return a dictionary of label->confidence pairs
    return dict(zip(learn.dls.vocab, map(float, probs)))

# Get a list of example images if they exist
examples = []
for filename in ["exa1.jpg", "exa2.jpg"]:
    if os.path.exists(filename):
        examples.append(filename)

# Create and launch the Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    examples=examples,
    title="German Bread Classification"
)

if __name__ == "__main__":
    interface.launch(inline=False)