import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np

# Use TFSMLayer to load the SavedModel
model = tf.keras.models.load_model("waste_sort_model.keras")

class_names = ["Non-Recyclable", "Recyclable"]

def classify_image(img):
    """Classify uploaded image as recyclable or non-recyclable."""
    img = img.resize((150, 150))  # Resize to match model input size
    img_array = np.array(img) / 255.0  # Normalize pixel values
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension

    predictions = model(img_array)
    predicted_class = class_names[np.argmax(predictions)]
    return f"Prediction: {predicted_class}"

# Define Gradio Interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Waste Classification",
    description="Upload an image of waste to classify as Recyclable or Non-Recyclable.",
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch(share=True)