import gradio as gr
from transformers import pipeline

# 모델 로드
pipeline = pipeline("image-classification", model="yangy50/garbage-classification")

def classify_image(image):
    results = pipeline(image)
    return {res['label']: res['score'] for res in results}

iface = gr.Interface(
    fn=classify_image, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Label(num_top_classes=3),
    title="쓰레기 분류기 ♻️",
    description="사진을 올리면 어떤 쓰레기인지 알려줍니다."
)

iface.launch()