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

# Load your trained model
model = load_model("waste_sort_model.h5")
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.predict(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()