app.py

import gradio as gr
from transformers import pipeline

# 1. Hugging Face에서 모델 불러오기
# 처음 실행 시 모델을 다운로드하느라 시간이 조금 걸릴 수 있습니다.
print("모델을 불러오는 중입니다...")
pipeline_clf = pipeline("image-classification", model="yangy50/garbage-classification")

# 2. 분류 함수 정의 (이미지가 들어오면 결과를 반환)
def classify_image(image):
    # 모델이 예측한 결과(리스트)를 가져옴
    predictions = pipeline_clf(image)
    
    # Gradio가 이해할 수 있는 딕셔너리 형태로 변환 {'라벨': 확률}
    # 예: {'plastic': 0.98, 'metal': 0.02}
    return {p['label']: p['score'] for p in predictions}

# 3. Gradio 인터페이스 만들기
iface = gr.Interface(
    fn=classify_image,             # 실행할 함수
    inputs=gr.Image(type="pil"),   # 입력: 이미지 업로드
    outputs=gr.Label(num_top_classes=3), # 출력: 상위 3개 라벨 표시
    title="쓰레기 분류기 (Garbage Classifier)",
    description="이미지를 올리면 종이, 플라스틱, 유리, 금속 등으로 분류해줍니다."
)

# 4. 앱 실행
if __name__ == "__main__":
    iface.launch()