import torch
import gradio as gr
from transformers import pipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

def predict(image):
    classifier = pipeline(task="image-classification")
    preds = classifier(image)
    return {pred["label"]: round(float(pred["score"]), 4) for pred in preds}

def clear_interface():
    interface.inputs[0].clear()
    interface.outputs[0].clear()

description = """
"""

interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.inputs.Image(label="Image to classify", type="pil"),
    ],
    outputs=gr.outputs.Label(),
    title="Image Classifier",
    description=description,
    buttons=[
        gr.Button(label="Clear", callback=clear_interface)
    ]
)

interface.launch()