Spaces:
Sleeping
Sleeping
| 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() | |