# app.py
import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np

# Load the model once
model = tf.keras.models.load_model("model/garbage_model.h5")
classes = ["Organic", "Recyclable", "Hazardous", "Metal", "Glass", "E-Waste"]

def predict_image(img):
    if isinstance(img, Image.Image):
        image = img.resize((150, 150))
    else:
        image = Image.fromarray(img).resize((150, 150))

    arr = np.array(image) / 255.0
    arr = np.expand_dims(arr, axis=0)
    preds = model.predict(arr)
    idx = int(np.argmax(preds))
    return {classes[i]: float(preds[0][i]) for i in range(len(classes))}

interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="♻️ Garbage Classifier",
    description="Upload an image of garbage to classify it!",
)

if __name__ == "__main__":
    interface.launch()