import streamlit as st
from transformers import pipeline
from PIL import Image
import matplotlib.pyplot as plt

# Load the pre-trained model
age_classifier = pipeline("image-classification", model="nateraw/vit-age-classifier")

# Function to classify age from an image
def classify_age(image):
    result = age_classifier(image)
    predicted_age = result[0]['label']
    confidence = result[0]['score']
    return predicted_age, confidence

# Streamlit UI
st.title("Age Classification App")
st.write("Upload an image to classify the person's age.")

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

# Process the uploaded image
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")

    # Display uploaded image
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Get prediction
    predicted_age, confidence = classify_age(image)

    # Show results
    st.write(f"### Predicted Age: {predicted_age}")
    st.write(f"**Confidence:** {confidence:.2f}")