import streamlit as st
from transformers import pipeline
from PIL import Image

# 1. 사전학습된 이미지 분류 파이프라인 로드
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

# 2. Streamlit UI 구성
st.title("이미지 분류 데모")
st.write("Hugging Face의 Vision Transformer(ViT) 모델을 사용해 이미지 분류를 시도합니다.")

uploaded_file = st.file_uploader("이미지를 업로드해주세요", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    # 3. 이미지 열기
    image = Image.open(uploaded_file)
    st.image(image, caption="업로드한 이미지")

    # 4. 분류 실행
    results = classifier(image)

    # 5. 결과 표시
    st.write("## 예측 결과")
    for result in results:
        st.write(f"**{result['label']}**: {result['score']:.4f}")