UPDATE app.py

app.py

@@ -1,26 +1,76 @@
-
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
+import os
+
+# Configure Streamlit for Hugging Face Spaces
+st.set_page_config(
+    page_title="ViT Image Classifier",
+    page_icon="🖼️",
+    layout="centered"
+)
 
-st.set_page_config(page_title="ViT Image Classifier")
-st.title("ViT Image Classification")
+st.title("🖼️ ViT Image Classification")
+st.markdown("Upload an image to classify it using Google's Vision Transformer model.")
 
 @st.cache_resource
 def load_model():
-    return pipeline("image-classification", model="google/vit-base-patch16-224")
+    """Load the ViT model with error handling."""
+    try:
+        return pipeline("image-classification", model="google/vit-base-patch16-224")
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None
+
+# Initialize the model
+with st.spinner("Loading model..."):
+    pipe = load_model()
 
-pipe = load_model()
+if pipe is None:
+    st.stop()
 
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+# File uploader
+uploaded_file = st.file_uploader(
+    "Choose an image file", 
+    type=["jpg", "jpeg", "png", "bmp", "tiff"],
+    help="Upload an image in JPG, PNG, BMP, or TIFF format"
+)
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    with st.spinner("Classifying..."):
-        preds = pipe(image)
+    try:
+        # Display the uploaded image
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        # Perform classification
+        with st.spinner("Classifying image..."):
+            preds = pipe(image)
+        
+        # Display results
+        st.subheader("🎯 Classification Results")
         
-    st.subheader("Predictions")
-    for i, pred in enumerate(preds):
-        st.write(f"{i+1}. {pred['label']} ({pred['score']:.3f})")
+        # Create columns for better layout
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.metric("Top Prediction", preds[0]['label'])
+        
+        with col2:
+            st.metric("Confidence", f"{preds[0]['score']:.1%}")
+        
+        # Show all predictions
+        st.subheader("📊 All Predictions")
+        for i, pred in enumerate(preds):
+            confidence = pred['score']
+            st.progress(confidence)
+            st.write(f"**{i+1}. {pred['label']}** - {confidence:.1%}")
+            
+    except Exception as e:
+        st.error(f"Error processing image: {str(e)}")
+
+else:
+    st.info("👆 Please upload an image to get started!")
+
+# Add footer
+st.markdown("---")
+st.markdown("Built with Streamlit and 🤗 Transformers")