Antoine245's picture
Update app.py
329e771
raw
history blame
673 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)
preds = [{"score": round(pred["score"], 4), "label": pred["label"]} for pred in preds]
return [pred["label"] for pred in preds] # Return a list of labels only
description = """
"""
gr.Interface(
fn=predict,
inputs=[
gr.inputs.Image(label="Image to classify", type="pil"),
],
outputs=gr.outputs.Label(), # Use Label output instead of JSON
title="Image Classifier",
description=description
).launch()