Antoine245's picture
Update app.py
22bc2c2
raw
history blame
738 Bytes
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()