Update app.py

app.py

@@ -49,6 +49,12 @@ with col2:
         st.session_state.example_image = EXAMPLE_2
         st.session_state.example_loaded = True
 
+# 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
@@ -57,12 +63,15 @@ if uploaded_file is not None:
     
     with st.spinner("Analyzing age..."):
         predictions = pipe(image)
+        top_pred = get_top_prediction(predictions)
         
-    # Display results
-    st.markdown("### Results:")
-    for pred in predictions:
-        st.progress(float(pred["score"]))
-        st.write(f"{pred['label']}: {pred['score']:.2%}")
+    # 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
@@ -74,18 +83,21 @@ elif 'example_loaded' in st.session_state and st.session_state.example_loaded:
     with st.spinner("Analyzing age..."):
         # Pass the actual PIL Image object to the pipeline
         predictions = pipe(image)
+        top_pred = get_top_prediction(predictions)
     
-    # Display results
-    st.markdown("### Results:")
-    for pred in predictions:
-        st.progress(float(pred["score"]))
-        st.write(f"{pred['label']}: {pred['score']:.2%}")
+    # 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 model returns confidence scores 
-for each possible age category.
+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.
 """)