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

st.set_page_config(page_title="ViT Image Classifier", page_icon="🖼️")

st.title("🖼️ ViT Image Classification")
st.write("Upload an image to classify it using Google's Vision Transformer model.")

@st.cache_resource
def load_model():
    return pipeline("image-classification", model="google/vit-base-patch16-224")

# Load model
pipe = load_model()

# File uploader
uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Classify
    with st.spinner("Classifying..."):
        predictions = pipe(image)
    
    # Show results
    st.subheader("Predictions:")
    for i, pred in enumerate(predictions):
        st.write(f"{i+1}. **{pred['label']}** - {pred['score']:.3f}")

else:
    st.info("Please upload an image to get started!")