import gradio as gr
from PIL import Image
from transformers import pipeline

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


def classify_image(input_img: Image.Image) -> dict:
    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.queue()
demo.launch()