Update app.py

app.py

@@ -1,55 +1,79 @@
 import gradio as gr
+from transformers import pipeline
+
+# Load a pre-trained image classifier (small + fast for demo)
+# You can replace this with a waste-specific model if available
+image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
+
+# Categories with more examples
+compostable = ["vegetable", "vegetables", "tomato", "onion", "potato", "carrot", 
+               "fruit", "apple", "banana", "orange", "mango", "food", "leaves", "cardboard", "paper"]
+recyclable = ["plastic", "bottle", "can", "glass", "metal", "aluminum", "tin", "carton"]
+harmful = ["syringe", "battery", "medical", "medicine", "chemical", "paint", "electronics", "toxic"]
+
+def classify_text(item):
+    """Classify based on text input"""
+    item = item.lower()
+    if any(word in item for word in compostable):
+        return "✅ Compostable Waste (Green Bin)"
+    elif any(word in item for word in recyclable):
+        return "♻️ Recyclable Waste (Blue Bin)"
+    elif any(word in item for word in harmful):
+        return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
+    else:
+        return "🚮 Unknown Waste (Grey Bin)"
+
+def classify_image(image_path):
+    """Classify based on uploaded or webcam image"""
+    preds = image_classifier(image_path, top_k=5)
+    labels = [p["label"].lower() for p in preds]
+
+    # Match predictions with categories
+    if any(word in labels for word in compostable):
+        return "✅ Compostable Waste (Green Bin)"
+    elif any(word in labels for word in recyclable):
+        return "♻️ Recyclable Waste (Blue Bin)"
+    elif any(word in labels for word in harmful):
+        return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
+    else:
+        return "🚮 Unknown Waste (Grey Bin)"
 
-# Categories
-compostable = ["vegetable", "fruit", "food", "leaves"]
-recyclable = ["plastic", "paper", "glass", "metal", "bottle", "can"]
-harmful = ["syringe", "battery", "medical", "chemical"]
-
-def classify_waste(input_type, image_path, text_input):
-    """
-    Classify waste based on either an uploaded image or text description.
-    """
-    # Choose which input to use
-    if input_type == "Image Upload" and image_path:
-        # Demo: classify based on filename
-        filename = image_path.lower()
-        if any(word in filename for word in compostable):
-            return "✅ Compostable Waste (Green Bin)"
-        elif any(word in filename for word in recyclable):
-            return "♻️ Recyclable Waste (Blue Bin)"
-        elif any(word in filename for word in harmful):
-            return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
-        else:
-            return "🚮 General Waste (Grey Bin)"
+def classify_waste(input_type, image_input, text_input):
+    """Main function to choose input type"""
+    if input_type == "Image Upload / Webcam" and image_input is not None:
+        return classify_image(image_input)
     elif input_type == "Text Description" and text_input:
-        item = text_input.lower()
-        if item in compostable:
-            return "✅ Compostable Waste (Green Bin)"
-        elif item in recyclable:
-            return "♻️ Recyclable Waste (Blue Bin)"
-        elif item in harmful:
-            return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
-        else:
-            return "🚮 General Waste (Grey Bin)"
+        return classify_text(text_input)
     else:
         return "❌ Please provide input for the selected option."
 
 # Gradio UI
-input_type = gr.Radio(["Image Upload", "Text Description"], label="Select input type")
-
-image_input = gr.Image(type="filepath", label="Upload a photo of your waste item")
-text_input = gr.Textbox(label="Type the waste item description")
-
-iface = gr.Interface(
-    fn=classify_waste,
-    inputs=[input_type, image_input, text_input],
-    outputs="text",
-    title="♻️ Eco-Friendly Waste Classifier",
-    description="Upload a photo OR type a description of a waste item to see which bin it belongs to!"
-)
-
+with gr.Blocks() as demo:
+    gr.Markdown("# ♻️ Eco-Friendly Waste Classifier")
+    gr.Markdown("Upload a photo (or use webcam) **OR** type a description to see which bin the item belongs to!")
+
+    input_type = gr.Radio(
+        ["Image Upload / Webcam", "Text Description"], 
+        label="Select input type", value="Text Description"
+    )
+
+    with gr.Row():
+        image_input = gr.Image(type="filepath", label="Upload or capture photo", sources=["upload", "webcam"])
+        text_input = gr.Textbox(label="Type the waste item description")
+
+    output = gr.Textbox(label="Classification Result")
+
+    # When user changes
+    classify_btn = gr.Button("Classify")
+    classify_btn.click(
+        fn=classify_waste, 
+        inputs=[input_type, image_input, text_input], 
+        outputs=output
+    )
+
+# Launch
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()