hot-dog-gradio / app.py
mischeiwiller's picture
Update app.py
a96ede8 verified
import gradio as gr
from transformers import pipeline
import numpy as np
# Initialize the pipeline
classifier = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
def predict(image):
# Check if the input is a file path or numpy array
if isinstance(image, str):
# If it's a file path, pass it directly to the pipeline
predictions = classifier(image)
else:
# If it's a numpy array, we need to convert it to PIL Image
from PIL import Image
image = Image.fromarray(image.astype('uint8'), 'RGB')
predictions = classifier(image)
# Convert predictions to the format expected by Gradio
return {p["label"]: float(p["score"]) for p in predictions}
# Create the Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(label="Upload hot dog candidate"),
outputs=gr.Label(num_top_classes=2),
title="Hot Dog? Or Not?",
description="Upload an image to see if it's a hot dog or not!",
)
if __name__ == "__main__":
iface.launch()