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)
    preds = [{"score": round(pred["score"], 4), "label": pred["label"]} for pred in preds]
    return preds

description = """
"""

def format_output(output):
    # Format the JSON output for better display
    formatted_output = ""
    for i, pred in enumerate(output):
        formatted_output += f"{i}: {{\n"
        formatted_output += f"  score: {pred['score']},\n"
        formatted_output += f"  label: {pred['label']}\n"
        formatted_output += "}\n"
    return formatted_output

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.inputs.Image(label="Image to classify", type="pil"),
    ],
    outputs=gr.outputs.JSON(),
    title="Image Classifier",
    description=description
)

# Set the post-processing hook to format the output
iface.interface_output_postprocessing = format_output

# Launch the interface
iface.launch()