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()