Update app.py

app.py

@@ -4,12 +4,15 @@ import os
 from classifier import GarbageClassifier
 from config import Config
 
-# Detect whether is in Huggingface Space
+# Check if running in Hugging Face Spaces environment
 try:
     import spaces
+
     HF_SPACES = True
+    print("Running in Hugging Face Spaces environment")
 except ImportError:
     HF_SPACES = False
+    print("Running in local environment")
 
 # Initialize classifier
 config = Config()
@@ -21,9 +24,9 @@ classifier.load_model()
 print("Model loaded successfully!")
 
 
-def classify_garbage(image):
+def classify_garbage_impl(image):
     """
-    Classify garbage in uploaded image
+    Actual classification implementation
     """
     if image is None:
         return "Please upload an image", "No image provided"
@@ -35,6 +38,15 @@ def classify_garbage(image):
         return "Error", f"Classification failed: {str(e)}"
 
 
+# Apply GPU decorator based on environment
+if HF_SPACES:
+    classify_garbage = spaces.GPU(classify_garbage_impl)
+    print("GPU decorator applied for Hugging Face Spaces")
+else:
+    classify_garbage = classify_garbage_impl
+    print("Running without GPU decorator")
+
+
 def get_example_images():
     """Get example images if they exist"""
     example_dir = "test_images"
@@ -73,11 +85,14 @@ with gr.Blocks(title="Garbage Classification System") as demo:
 
     # Category information
     with gr.Accordion("📋 Garbage Categories Information", open=False):
-        category_info = classifier.get_categories_info()
-        for category, description in category_info.items():
-            gr.Markdown(f"**{category}**: {description}")
-
-    # Examples
+        try:
+            category_info = classifier.get_categories_info()
+            for category, description in category_info.items():
+                gr.Markdown(f"**{category}**: {description}")
+        except Exception as e:
+            gr.Markdown(f"Categories information not available: {str(e)}")
+
+    # Examples section
     examples = get_example_images()
     if examples:
         gr.Examples(examples=examples, inputs=image_input, label="Example Images")
@@ -97,4 +112,8 @@ with gr.Blocks(title="Garbage Classification System") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    # Set appropriate launch parameters based on environment
+    if HF_SPACES:
+        demo.launch()
+    else:
+        demo.launch(share=True)