import gradio as gr
from transformers import pipeline
from PIL import Image

# Initialize Hugging Face pipeline
pipe = pipeline("image-classification", model="yangy50/garbage-classification")

def classify_image(image):
    # Ensure image is in RGB mode
    image = image.convert("RGB")
    results = pipe(image)
    return {result['label']: result['score'] for result in results}

demo = gr.Interface(fn=classify_image, inputs="image", outputs="label")
demo.launch()