# Import necessary libraries
import streamlit as st
from transformers import pipeline
from PIL import Image

# Function to classify age from an image
def classify_age(image_path):
    # Load the age classification model
    age_classifier = pipeline("image-classification", model="nateraw/vit-age-classifier")
    # Predict age
    result = age_classifier(image_path)
    return result

# Main part of the app
st.set_page_config(page_title="Age Classifier", page_icon="👶")
st.header("Age Classification using nateraw/vit-age-classifier")
uploaded_file = st.file_uploader("Upload an image of a person...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Save the uploaded file temporarily
    with open(uploaded_file.name, "wb") as file:
        file.write(uploaded_file.getvalue())
    
    # Display the uploaded image
    st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
    
    # Classify age
    st.text('Classifying age...')
    age_result = classify_age(uploaded_file.name)
    
    # Display the result
    st.write("Age Classification Result:")
    for res in age_result:
        st.write(f"{res['label']}: {res['score']:.2f}")