import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np

# Load your trained model
model = tf.keras.models.load_model("model.keras", custom_objects={"KerasLayer": hub.KerasLayer})

# Define your waste class labels
class_names = ['batteries', 'clothes', 'e-waste', 'glass', 'light blubs',
               'metal', 'organic', 'paper', 'plastic']
IMG_SIZE = 224

# Image preprocessing function
def preprocess_image(image):
    image = tf.convert_to_tensor(image, dtype=tf.float32)
    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE)) / 255.0
    return tf.expand_dims(image, axis=0)

# Prediction function
def predict_waste(image):
    processed = preprocess_image(image)
    preds = model.predict(processed)
    label = class_names[np.argmax(preds)]
    confidence = float(np.max(preds))
    return f"{label} ({confidence:.2f})"


# Launch Gradio Interface
import gradio as gr

def predict_waste(image):
    # Your prediction logic here
    return "Example category: Plastic"

examples = [
    ["examples/download (1).jpg"],
    ["examples/download (13) (1).jpg"],
    ["examples/download (1) (1).jpg"],
    ["examples/download (14).jpg"]
]

with gr.Blocks(title="Waste Classification Model 🌱") as demo:
    gr.Markdown("""
    # ♻️ Waste Classification with AI  
    Upload or drag an image of a waste item below. Our model will classify it into one of 9 categories such as plastic, glass, paper, or metal.

    - ✅ Powered by MobileNetV2  
    - 🌍 Helping smart cities manage waste efficiently  
    - 📷 Works with real-world waste images  
    """)

    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="numpy", label="📤 Upload a Waste Image")
            submit_btn = gr.Button("🔍 Classify Waste")
        with gr.Column(scale=1):
            prediction_output = gr.Textbox(label="🧠 Predicted Category", interactive=False)

    submit_btn.click(fn=predict_waste, inputs=image_input, outputs=prediction_output)

    gr.Examples(
        examples=examples,
        inputs=image_input,
        label="🖼 Example Waste Images",
        examples_per_page=4
    )

    gr.Markdown("""
    ---
    🔬 *Model accuracy may vary based on image clarity and background noise. For best results, upload clear, well-lit images.*  
    """)

demo.launch()