Create app.py

app.py

@@ -0,0 +1,139 @@
+import streamlit as st
+from transformers import pipeline
+from PIL import Image
+import requests
+from io import BytesIO
+import logging
+import torch
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class AgeClassifier:
+    def __init__(self):
+        try:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.pipe = pipeline("image-classification", model="nateraw/vit-age-classifier", device=self.device)
+            logger.info(f"Model loaded successfully on {self.device}")
+        except Exception as e:
+            logger.error(f"Failed to initialize pipeline: {e}")
+            raise
+
+    def classify_image(self, image):
+        try:
+            return self.pipe(image)
+        except Exception as e:
+            logger.error(f"Classification failed: {e}")
+            return None
+
+    @staticmethod
+    def format_results(results):
+        if not results:
+            return "No valid results"
+        return results
+
+def load_image_from_url(url):
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        image = Image.open(BytesIO(response.content))
+        return image
+    except Exception as e:
+        st.error(f"Error loading image from URL: {e}")
+        return None
+
+def main():
+    st.set_page_config(
+        page_title="Age Classification App",
+        page_icon="👤",
+        layout="wide"
+    )
+
+    st.title("Age Classification App 👤")
+    st.write("Upload an image or provide a URL to classify the age range of people in the image.")
+
+    # Initialize the classifier
+    @st.cache_resource
+    def get_classifier():
+        return AgeClassifier()
+
+    classifier = get_classifier()
+
+    # Create two columns for input methods
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.subheader("Upload Image")
+        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+    with col2:
+        st.subheader("Image URL")
+        image_url = st.text_input("Enter image URL")
+
+    # Process the image
+    image = None
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+    elif image_url:
+        image = load_image_from_url(image_url)
+
+    if image:
+        # Display the image
+        st.image(image, caption="Input Image", use_column_width=True)
+
+        # Add a classify button
+        if st.button("Classify Age"):
+            with st.spinner("Classifying..."):
+                results = classifier.classify_image(image)
+                
+                if results:
+                    # Create a bar chart
+                    st.subheader("Classification Results")
+                    
+                    # Convert results to format suitable for bar chart
+                    labels = [r['label'] for r in results]
+                    scores = [r['score'] * 100 for r in results]
+                    
+                    # Display most likely age range
+                    most_likely = max(results, key=lambda x: x['score'])
+                    st.success(f"Most likely age range: {most_likely['label']} ({most_likely['score']*100:.1f}%)")
+                    
+                    # Create bar chart
+                    chart_data = {
+                        'Age Range': labels,
+                        'Confidence (%)': scores
+                    }
+                    st.bar_chart(chart_data, x='Age Range', y='Confidence (%)')
+                    
+                    # Display detailed results in an expander
+                    with st.expander("See detailed results"):
+                        st.write("Confidence scores for all age ranges:")
+                        for result in results:
+                            st.write(f"{result['label']}: {result['score']*100:.1f}%")
+                else:
+                    st.error("Could not classify the image. Please try another image.")
+
+    # Add information about the model
+    with st.sidebar:
+        st.header("About")
+        st.write("""
+        This app uses the ViT (Vision Transformer) model trained for age classification.
+        
+        The model classifies images into the following age ranges:
+        - 0-2 years
+        - 3-9 years
+        - 10-19 years
+        - 20-29 years
+        - 30-39 years
+        - 40-49 years
+        - 50-59 years
+        - 60-69 years
+        - 70+ years
+        """)
+        
+        st.write("Model: nateraw/vit-age-classifier")
+        st.write(f"Running on: {classifier.device}")
+
+if __name__ == "__main__":
+    main()