File size: 3,727 Bytes
6f5f080
a707f72
4372096
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a707f72
 
4372096
a707f72
 
4372096
 
 
a707f72
 
 
 
 
 
 
 
 
 
 
4372096
 
 
a707f72
4372096
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a707f72
 
 
ce2f93e
a707f72
083d24e
ce2f93e
6f5f080
4372096
6dd53f5
4372096
a707f72
4372096
a707f72
 
 
4372096
 
 
a707f72
 
 
 
 
 
 
 
 
 
4372096
 
a707f72
 
 
4372096
a707f72
4372096
c33843a
4372096
 
04d7307
8e684ad
d0770f2
6dd53f5
7345cdf
 
e327f54
c0795c6
 
ce2f93e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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)