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.keras")


# Classes in training order
class_names = ['batteries', 'clothes', 'e-waste', 'glass', 'light blubs', 'metal', 'organic', 'paper', 'plastic']
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
    return np.expand_dims(image_array, axis=0)

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

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an image of a recyclable item"),
    outputs=gr.Textbox(label="Prediction"),
    title="♻️ Waste Image Classifier",
    description="Classify recyclable waste images into 9 categories using a MobileNetV2-based model."
)

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