Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
-import numpy as np
-import io
+import requests
+from io import BytesIO
 
 # Set page configuration
 st.set_page_config(
@@ -20,6 +20,12 @@ 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()
 
@@ -30,15 +36,18 @@ uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png
 st.markdown("### Or try an example:")
 col1, col2 = st.columns(2)
 
+EXAMPLE_1 = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/person1.jpg"
+EXAMPLE_2 = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/person2.jpg"
+
 with col1:
     if st.button("Example 1"):
-        example_img = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/person1.jpg"
-        st.session_state.example_image = example_img
+        st.session_state.example_image = EXAMPLE_1
+        st.session_state.example_loaded = True
 
 with col2:
     if st.button("Example 2"):
-        example_img = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/person2.jpg"
-        st.session_state.example_image = example_img
+        st.session_state.example_image = EXAMPLE_2
+        st.session_state.example_loaded = True
 
 # Process the image and display results
 if uploaded_file is not None:
@@ -55,12 +64,16 @@ if uploaded_file is not None:
         st.progress(float(pred["score"]))
         st.write(f"{pred['label']}: {pred['score']:.2%}")
 
-elif 'example_image' in st.session_state:
+elif 'example_loaded' in st.session_state and st.session_state.example_loaded:
     # Process example image
-    st.image(st.session_state.example_image, caption="Example Image", use_column_width=True)
+    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_column_width=True)
     
     with st.spinner("Analyzing age..."):
-        predictions = pipe(st.session_state.example_image)
+        # Pass the actual PIL Image object to the pipeline
+        predictions = pipe(image)
     
     # Display results
     st.markdown("### Results:")