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

# Define the pipeline
@st.cache_resource
def load_pipeline():
    return pipeline("image-classification", model="yangy50/garbage-classification")

pipe = load_pipeline()

# Streamlit UI
st.title("Garbage Classification App")
st.write("Upload an image to classify it as a type of garbage.")

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

if uploaded_file is not None:
    # Load image
    image = Image.open(uploaded_file)
    
    # Display image
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Run inference
    results = pipe(image)
    
    # Get top prediction
    top_prediction = max(results, key=lambda x: x["score"])
    
    # Display result
    st.write(f"**Predicted Class:** {top_prediction['label']}")
    st.write(f"**Confidence:** {top_prediction['score']:.2f}")