import gradio as gr
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow.keras.models import load_model

class_names = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']


model=load_model("best_mobilenetv2_model.keras")

def classify_image(img):

    img = img.convert("RGB")
    img = img.resize((224, 224))

    img_tensor = tf.convert_to_tensor(np.array(img), dtype=tf.float32)

    img_tensor = tf.expand_dims(img_tensor, axis=0)

    prediction = model.predict(img_tensor)
    predicted_class_index = np.argmax(prediction)
    predicted_class_name = class_names[predicted_class_index]
    confidence = prediction[0][predicted_class_index]

    return f"Predicted: {predicted_class_name} (Confidence: {confidence:.2%})"


iface = gr.Interface(
   fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload Waste Image"),
    outputs=gr.Textbox(label="Prediction"),
    title="♻️ Waste Classifier",
    description="Upload an image of cardboard, plastic, metal, paper, trash, or glass to classify it."
)


# Launch the interface
iface.launch()