from huggingface_hub import from_pretrained_fastai
import gradio as gr


categories = ('checkers', 'chess')
repo_id = "moshe742/chess-or-not"
learn = from_pretrained_fastai(repo_id)


def classify_image(image):
    global learn
    pred, idx, probs = learn.predict(image)
    return dict(zip(categories, map(float, probs)))


image = gr.Image(height=192, width=192)
label = gr.Label()
examples = ['chess.jpg', 'checkers.webp', "dono.jpg"]

demo = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
demo.launch()