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 predict_trash(image):
    """Simple prediction function"""
    # Resize to 224x224
    img = cv2.resize(image, (224, 224))
    
    # Convert BGR to RGB (if needed)
    if len(img.shape) == 3 and img.shape[2] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # Normalize
    img = img / 255.0
    
    # Add batch dimension
    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) * 100
    result = class_names[np.argmax(predictions)]
    
    return f"{result} ({confidence:.1f}% confident)"

# Create the simplest possible interface
demo = gr.Interface(
    fn=predict_trash,
    inputs=gr.Image(sources=["upload", "webcam"], type="numpy"),
    outputs="text",
    title="Trash Classifier",
    description="Take a photo or upload image → AI says if it's Organic or Recyclable",
    examples=[
        ["https://images.unsplash.com/photo-1558904541-efa843a96d01?w=400"],  # Banana
        ["https://images.unsplash.com/photo-1602143407151-7111542de6e8?w=400"],  # Bottle
    ]
)

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