OGOGOG's picture
Create app.py
19a6479 verified
import gradio as gr
from transformers import pipeline
from PIL import Image
# Load image classifier (general-purpose)
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
# Define recyclable classes (based on common outputs of the model)
RECYCLABLE_CLASSES = {
"plastic bottle", "water bottle", "can", "glass", "cup",
"paper", "newspaper", "cardboard", "box", "carton", "tin"
}
def classify_trash(image):
results = classifier(image)
top_label = results[0]['label'].lower()
confidence = results[0]['score']
is_recyclable = any(recycle_word in top_label for recycle_word in RECYCLABLE_CLASSES)
label = "♻️ Recyclable" if is_recyclable else "🗑️ Not Recyclable"
return f"{label}\nDetected: {top_label}\nConfidence: {confidence:.2%}"
# Gradio interface
demo = gr.Interface(
fn=classify_trash,
inputs=gr.Image(type="pil"),
outputs=gr.Textbox(label="Classification"),
title="♻️ Trash Classifier: Recyclable or Not?",
description="Upload an image of an object (like a bottle, banana peel, or can) and find out if it is recyclable.",
)
if __name__ == "__main__":
demo.launch()