import gradio as gr
from transformers import pipeline

pipeline = pipeline(
    "image-classification", model="fauzanardh/anime-time-beit_classifier"
)


def classify_image(input_img):
    data = pipeline(input_img, top_k=4)
    out = {}
    for d in data:
        out[d["label"]] = d["score"]
    return out


demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(label="Prediction"),
)

demo.launch()