trash-sorter-ai / app.py
tahzaya's picture
Upload app.py with huggingface_hub
de394d4 verified
Raw
History Blame Contribute Delete
1.32 kB
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()