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

# Load model
model = tf.keras.models.load_model("model.h5")

# Correct class names in training order
class_names = ['batteries', 'clothes', 'e-waste', 'glass', 'light blubs', 'metal', 'organic', 'paper', 'plastic']

# Match training input size
IMG_SIZE = (224, 224)

def preprocess(image):
    image = image.resize(IMG_SIZE)
    image_array = tf.keras.preprocessing.image.img_to_array(image)
    image_array = image_array / 255.0  # Normalize if this was done in training
    return np.expand_dims(image_array, axis=0)

def predict(image):
    input_tensor = preprocess(image)
    predictions = model.predict(input_tensor)[0]
    top_class = class_names[np.argmax(predictions)]
    confidence = np.max(predictions) * 100
    return f"{top_class} ({confidence:.2f}%)"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an image of a recyclable item to classify it."),
    outputs=gr.Textbox(label="Output"),
    title="Waste Image Classifier"
)

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