Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| from PIL import Image | |
| # ------------------------------- | |
| # Load pre-trained image classifier | |
| # ------------------------------- | |
| # Small and fast model for demo | |
| try: | |
| image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224") | |
| except Exception as e: | |
| print("⚠️ Could not load image classifier:", e) | |
| image_classifier = None | |
| # ------------------------------- | |
| # Waste categories | |
| # ------------------------------- | |
| 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"] | |
| # ------------------------------- | |
| # Text classification | |
| # ------------------------------- | |
| def classify_text(item): | |
| 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)" | |
| # ------------------------------- | |
| # Image classification | |
| # ------------------------------- | |
| def classify_image(image_path): | |
| if image_classifier is None: | |
| return "⚠️ Image classifier not available." | |
| try: | |
| img = Image.open(image_path) | |
| preds = image_classifier(img, top_k=5) | |
| labels = [p["label"].lower() for p in preds] | |
| 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)" | |
| except Exception as e: | |
| return f"⚠️ Error classifying image: {e}" | |
| # ------------------------------- | |
| # Main classifier | |
| # ------------------------------- | |
| def classify_waste(input_type, image_input, text_input): | |
| 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: | |
| return classify_text(text_input) | |
| else: | |
| return "❌ Please provide input for the selected option." | |
| # ------------------------------- | |
| # Gradio UI | |
| # ------------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# ♻️ EcoSort: Smart 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") | |
| classify_btn = gr.Button("Classify") | |
| classify_btn.click( | |
| fn=classify_waste, | |
| inputs=[input_type, image_input, text_input], | |
| outputs=output | |
| ) | |
| # ------------------------------- | |
| # Launch | |
| # ------------------------------- | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |