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


pipeline = pipeline(task="image-classification", model="nateraw/vit-age-classifier")

def predict(image):
    predictions = pipeline(image)
    return {p["label"]: p["score"] for p in predictions}

def main():
    st.title("Age Classification From Image")
    
    with st.form("my_form"):
        uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

        if uploaded_file is not None:
        # Display the uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption="Your uploaded Image", use_column_width=True)
        clicked = st.form_submit_button("Press to predict")
        if clicked:
            results = predict(image)
            k = [] 
            v = []
            for key, value in results.items():
                value = round(value*100,2)
                v.append(value)
                k.append(key)
            vp = [str(item) + '%' for item in v]
            result = k[0]
            st.success('The predicted age is {}'.format(result))
            df = pd.DataFrame({'Prediction': k,'Accuracy':vp})
            st.dataframe(df,hide_index=True)

if __name__ == "__main__":
    main()