import gradio as gr
from transformers import pipeline
from PIL import Image

# Load the waste classification model
classifier = pipeline(
    "image-classification",
    model="watersplash/waste-classification"
)

# Recycling instructions & fun facts
waste_instructions = {
    "plastic": "♻️ Place in plastic recycling bin. Avoid single-use plastics where possible.",
    "glass": "🍾 Clean and recycle in glass bin. Broken glass should be wrapped safely before disposal.",
    "paper": "📄 Recycle clean paper. Avoid recycling wet or oily paper.",
    "metal": "🥫 Recycle cans after rinsing. Large scrap metal should go to a recycling center.",
    "cardboard": "📦 Flatten boxes before recycling.",
    "organic": "🌱 Compost food waste and garden scraps.",
    "trash": "🗑️ Not recyclable. Dispose properly."
}

fun_facts = {
    "plastic": "It can take up to 500 years for plastic to decompose!",
    "glass": "Glass is 100% recyclable and can be recycled endlessly without loss of quality.",
    "paper": "Recycling 1 ton of paper saves 17 trees and 7,000 gallons of water.",
    "metal": "Aluminum cans can be recycled and back on shelves in as little as 60 days.",
    "cardboard": "Cardboard can be recycled up to 5–7 times before the fibers weaken.",
    "organic": "Composting reduces landfill waste and returns nutrients to the soil.",
    "trash": "Some waste cannot be recycled — reduce and reuse where possible."
}

# Function to process uploaded image
def identify_waste(image):
    # Resize image to prevent memory issues
    image = image.resize((224, 224))

    # Run classification
    results = classifier(image)
    label = results[0]["label"].lower()
    
    if label in waste_instructions:
        return f"🗂️ Detected Waste Type: {label.capitalize()}\n\n✅ {waste_instructions[label]}\n\n💡 Fun Fact: {fun_facts[label]}"
    else:
        return f"🗂️ Detected Waste Type: {label}\n\n⚠️ No recycling info available. Please dispose responsibly."

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# ♻️ Smart Waste Identifier")
    gr.Markdown("Upload an image to identify the type of waste and get recycling instructions.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Waste Image")
            btn = gr.Button("Identify Waste")
        with gr.Column():
            output_text = gr.Textbox(label="Result", lines=6)

    btn.click(fn=identify_waste, inputs=image_input, outputs=output_text)

    gr.Markdown("---")
    gr.Markdown("### ℹ️ This app classifies input images into categories and provides recycling tips or fun facts.")

# Launch the app (no enable_queue, since older Gradio version)
demo.launch()