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

# Load the classifier model
@st.cache_resource
def load_classifier():
    return pipeline("image-classification", model="nateraw/vit-age-classifier")

classifier = load_classifier()

# Streamlit UI
st.title("Age Classifier App 🧑👵")
st.write("Upload an image to predict the age category.")

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

if uploaded_file is not None:
    # Load and display the image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Classify the age
    st.write("Classifying...")
    results = classifier(image)

    # Get the label with the highest score
    if results:
        best_result = max(results, key=lambda x: x["score"])
        st.subheader("Predicted Age Category:")
        st.write(f"🟢 **{best_result['label']}** (Confidence: {best_result['score']:.4f})")