Update app.py

app.py

@@ -1,39 +1,48 @@
 import gradio as gr
 import tensorflow as tf
+import tensorflow_hub as hub
 import numpy as np
 
-# Image size used in training
-IMG_SIZE = (224, 224)
-
-# Final label mapping
-index_to_label = {
-    0: "paper",
-    1: "plastic",
-    2: "batteries",
-    3: "metal",
-    4: "glass",
-    5: "clothes",
-    6: "organic",
-    7: "light bulbs",
-    8: "e-waste"
-}
-
-# Load trained model
-model = tf.keras.models.load_model("model.h5")
+# Load your trained model
+model = tf.keras.models.load_model("model.h5", custom_objects={'KerasLayer': hub.KerasLayer})
+
+# Define your waste class labels
+class_names = ['batteries', 'clothes', 'e-waste', 'glass', 'light blubs',
+               'metal', 'organic', 'paper', 'plastic']
+IMG_SIZE = 224
 
 # Image preprocessing function
-def classify_image(img):
-    img = tf.image.resize(img, IMG_SIZE)
-    img = tf.expand_dims(img, 0)
-    img = tf.cast(img, tf.float32) / 255.0
-    preds = model.predict(img)[0]
-    return {index_to_label[i]: float(preds[i]) for i in range(len(preds))}
-
-# Gradio Interface
-gr.Interface(
-    fn=classify_image,
+def preprocess_image(image):
+    image = tf.convert_to_tensor(image, dtype=tf.float32)
+    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE)) / 255.0
+    return tf.expand_dims(image, axis=0)
+
+# Prediction function
+def predict_waste(image):
+    processed = preprocess_image(image)
+    preds = model.predict(processed)
+    label = class_names[np.argmax(preds)]
+    confidence = float(np.max(preds))
+    return f"{label} ({confidence:.2f})"
+
+# Example images – update these paths with actual examples in your Colab space
+examples = [
+    ["examples/Zeitungen-und-Zeitschriften.png"],
+    ["examples/hazardous-toxic-electronic-waste-mixture-260nw-1593448036.jpg"],
+    ["examples/composting_medium-size.jpg"],
+    ["examples/zero-waste-recycling-wooden-thing-wooden-background-zero-waste-recycling-wooden-thing-wooden-background-flat-lay-158056259.jpg"],
+    ["examples/plastic waste.png"]
+     ]
+
+
+# Launch Gradio Interface
+iface = gr.Interface(
+    fn=predict_waste,
     inputs=gr.Image(type="numpy", label="Upload a Waste Image"),
-    outputs=gr.Label(num_top_classes=3),
-    title="♻️ Smart Waste Classifier",
-    description="Upload an image of waste to classify it into 9 categories (paper, plastic, batteries, metal, glass, clothes, organic, light bulbs, e-waste). Built with TensorFlow + MobileNetV2"
-).launch()
+    outputs="text",
+    title="Waste Classification Model 🌱",
+    description="Upload an image of a waste item to classify it as one of 9 categories.",
+    examples=examples
+)
+
+iface.launch()