ageclassifier / app.py
AndrewLi403's picture
Create app.py
e242f20 verified
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}")