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

#function part

def classify_image(image_path):
    # Load the pre-trained image classification model
    classifier = pipeline("image-classification", model="nateraw/vit-age-classifier")
    image = Image.open(image_path)
    predictions = classifier(image)
    return predictions

st.set_page_config(page_title="Age Classifier", page_icon="📷")
st.header("Image Age Classification")

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

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Perform image classification
    st.text('Classifying image...')
    predictions = classify_image(uploaded_file)
    
    # Display the top prediction
    if predictions:
        top_prediction = predictions[0]
        st.write(f"**Predicted Age Group:** {top_prediction['label']}")
        st.write(f"**Confidence:** {top_prediction['score']:.2f}")