Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import pipeline | |
| from PIL import Image | |
| import requests | |
| from io import BytesIO | |
| # Set page configuration | |
| st.set_page_config( | |
| page_title="Age Group Classification", | |
| page_icon="👶👧👩👵", | |
| layout="centered" | |
| ) | |
| # Add title and description | |
| st.title("Age Group Classification") | |
| st.markdown("Upload an image of a person to classify their approximate age group using the `nateraw/vit-age-classifier` model.") | |
| def load_model(): | |
| """Load the age classification model and cache it.""" | |
| return pipeline("image-classification", model="nateraw/vit-age-classifier") | |
| def load_image_from_url(url): | |
| """Load an image from a URL.""" | |
| response = requests.get(url) | |
| img = Image.open(BytesIO(response.content)) | |
| return img | |
| # Load the model | |
| pipe = load_model() | |
| # Create file uploader | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| # Function to get top prediction | |
| def get_top_prediction(predictions): | |
| # Get the prediction with highest confidence | |
| top_prediction = max(predictions, key=lambda x: x['score']) | |
| return top_prediction | |
| # Process the image and display results | |
| if uploaded_file is not None: | |
| # Process uploaded image | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_container_width=True) | |
| with st.spinner("Analyzing age..."): | |
| predictions = pipe(image) | |
| top_pred = get_top_prediction(predictions) | |
| # Display result | |
| st.markdown("### Result:") | |
| st.metric( | |
| label="Predicted Age Range", | |
| value=top_pred['label'], | |
| delta=f"Confidence: {top_pred['score']:.2%}" | |
| ) | |
| elif 'example_loaded' in st.session_state and st.session_state.example_loaded: | |
| # Process example image | |
| with st.spinner("Loading example image..."): | |
| # Download and load the image properly | |
| image = load_image_from_url(st.session_state.example_image) | |
| st.image(image, caption="Example Image", use_container_width=True) | |
| with st.spinner("Analyzing age..."): | |
| # Pass the actual PIL Image object to the pipeline | |
| predictions = pipe(image) | |
| top_pred = get_top_prediction(predictions) | |
| # Display result | |
| st.markdown("### Result:") | |
| st.metric( | |
| label="Predicted Age Range", | |
| value=top_pred['label'], | |
| delta=f"Confidence: {top_pred['score']:.2%}" | |
| ) | |
| # Add information about the model | |
| st.markdown("---") | |
| st.markdown("### About the Model") | |
| st.markdown(""" | |
| This app uses the `nateraw/vit-age-classifier` model from Hugging Face, which classifies | |
| images into age groups like "0-2", "3-9", "10-19", etc. The app displays only the most | |
| likely age range prediction with its confidence score. | |
| """) |