import gradio as gr
import numpy as np
import tensorflow as tf
from PIL import Image

model = tf.keras.models.load_model(
    "waste_classifier_final.h5",
    compile=False
)

labels = [
    "Cardboard",
    "Glass",
    "Metal",
    "Paper",
    "Plastic",
    "Trash"
]

def classify_waste(image):
    img = image.resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    predictions = model.predict(img_array)[0]
    return {labels[i]: float(predictions[i]) for i in range(len(labels))}

demo = gr.Interface(
    fn=classify_waste,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Waste Classifier",
    description="Upload a waste image to classify it."
)

demo.launch()