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()