import gradio as gr
import numpy as np
import cv2
import tensorflow as tf

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

def classify_waste(image):
    # Preprocess image
    img = cv2.resize(image, (224, 224))
    img = img / 255.0
    img = np.expand_dims(img, axis=0)
    
    # Predict
    predictions = model.predict(img, verbose=0)[0]
    
    # Get results
    class_names = ["🗑️ ORGANIC", "♻️ RECYCLABLE"]
    confidence = np.max(predictions)
    predicted_class = class_names[np.argmax(predictions)]
    
    # Format output
    result = f"""
    ## {predicted_class}
    
    **Confidence:** {confidence:.2%}
    
    **Breakdown:**
    - Organic: {predictions[0]:.2%}
    - Recyclable: {predictions[1]:.2%}
    
    **What to do:**
    { "Compost or dispose in green bin" if predicted_class == "🗑️ ORGANIC" else "Clean and place in recycling bin" }
    """
    
    return result

# Create interface
demo = gr.Interface(
    fn=classify_waste,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Markdown(),
    title="🗑️ AI Trash Sorter",
    description="Upload a photo of waste to classify as Organic or Recyclable",
    examples=[
        ["example_organic.jpg"],
        ["example_recyclable.jpg"]
    ]
)

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