Update app.py

app.py

@@ -1,18 +1,21 @@
 import os
 
-# ✅ Use temp directory for safe model caching (fixes permission issues)
+# ✅ Use temp dir for safe model caching in Spaces/Docker
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 
 import streamlit as st
 from PIL import Image
 import random
 import torch
-from transformers import AutoImageProcessor, SiglipForImageClassification
+from transformers import AutoImageProcessor, SiglipForImageClassification, logging
 
-# --- Constants ---
+# Optional: show more debug info if something fails
+logging.set_verbosity_error()
+
+# Constants
 MODEL_NAME = "prithivMLmods/Recycling-Net-11"
 
-# Sustainability Tips
+# Daily sustainability tips
 TIPS = [
     "Rinse containers before recycling to avoid contamination.",
     "Avoid using plastic bags for recyclables – use bins or boxes.",
@@ -26,7 +29,7 @@ TIPS = [
     "Avoid single-use plastics whenever possible.",
 ]
 
-# Government recycling resources
+# Government recycling links
 GOVERNMENT_LINKS = {
     "Pakistan": "https://environment.gov.pk/",
     "India": "https://www.cpcb.nic.in/",
@@ -38,15 +41,20 @@ GOVERNMENT_LINKS = {
     "Germany": "https://www.bmu.de/en/topics/water-waste-soil/waste-management",
 }
 
-# --- Load Model ---
-@st.cache_resource(show_spinner=False)
+# Load model and processor
+@st.cache_resource(show_spinner="🔄 Loading AI model...")
 def load_model():
-    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-    model = SiglipForImageClassification.from_pretrained(MODEL_NAME)
-    model.eval()
-    return processor, model
-
-# --- Predict Function ---
+    try:
+        processor = AutoImageProcessor.from_pretrained(MODEL_NAME, revision="main")
+        model = SiglipForImageClassification.from_pretrained(MODEL_NAME, revision="main")
+        model.eval()
+        return processor, model
+    except Exception as e:
+        st.error("❌ Failed to load the model. Please check the model name or your connection.")
+        st.exception(e)
+        raise e
+
+# Prediction function
 def predict(image: Image.Image, processor, model):
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
@@ -58,9 +66,9 @@ def predict(image: Image.Image, processor, model):
     confidence = conf.item()
     return class_name, confidence
 
-# --- Suggestion Function ---
+# Recycling tip per label
 def get_suggestion(label: str) -> str:
-    tips = {
+    suggestions = {
         "aluminium": "Rinse and recycle aluminum cans. They are infinitely recyclable.",
         "batteries": "Do not throw in the trash. Use proper e-waste collection centers.",
         "cardboard": "Flatten and keep dry. Avoid greasy pizza boxes.",
@@ -72,17 +80,16 @@ def get_suggestion(label: str) -> str:
         "soft plastics": "Often require store drop-off. Don’t mix with other recyclables.",
         "takeaway cups": "Check local rules. Many are lined and not recyclable curbside.",
     }
-    return tips.get(label, "Please check your local rules for proper disposal of this item.")
+    return suggestions.get(label, "Please check your local rules for proper disposal of this item.")
 
-# --- Streamlit App ---
+# Main app
 def main():
     st.set_page_config(page_title="♻️ Recycling Helper AI", layout="centered")
     st.title("♻️ Recycling Helper AI")
     st.subheader("An AI-powered app to identify recyclable materials and promote sustainability.")
-
     st.markdown("---")
 
-    # Sidebar content
+    # Sidebar
     with st.sidebar:
         st.header("📘 About This App")
         st.markdown(
@@ -102,10 +109,10 @@ def main():
         tip = random.choice(TIPS)
         st.success(tip)
 
-    # Load the model
+    # Load model
     processor, model = load_model()
 
-    # File upload
+    # Upload image
     st.markdown("### 📤 Upload Waste Image")
     uploaded_file = st.file_uploader("Upload an image of a recyclable item", type=["png", "jpg", "jpeg"])
 
@@ -114,24 +121,22 @@ def main():
             image = Image.open(uploaded_file).convert("RGB")
             st.image(image, caption="Uploaded Image", use_column_width=True)
 
-            with st.spinner("Analyzing with AI model..."):
-                class_name, confidence = predict(image, processor, model)
-
-            st.success(f"**Predicted Material:** `{class_name}`  \n**Confidence:** `{confidence:.2%}`")
+            with st.spinner("🔍 Classifying image..."):
+                label, confidence = predict(image, processor, model)
 
-            suggestion = get_suggestion(class_name)
-            st.info(f"**Tip:** {suggestion}")
+            st.success(f"**Predicted Material:** `{label}`  \n**Confidence:** `{confidence:.2%}`")
+            st.info(f"**Disposal Tip:** {get_suggestion(label)}")
 
         except Exception as e:
-            st.error(f"Something went wrong during prediction: {e}")
+            st.error("An error occurred during prediction.")
+            st.exception(e)
 
-    # Optional: Show class labels
-    with st.expander("🔍 Show Model Classes"):
+    with st.expander("🔍 Show All Recognizable Materials"):
         st.write(model.config.id2label)
 
     st.markdown("---")
     st.caption("Made with 💚 for a sustainable future | Hackathon 2025")
 
-# Run app
+# Run
 if __name__ == "__main__":
     main()