Update app.py

app.py

@@ -3,36 +3,34 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# Load the model
+# Load model
 model = tf.keras.models.load_model("model.h5")
 
-# Define the class labels (update based on your project)
-class_names = ['batteries', 'clothes', 'e-waste', 'glass', 'light blubs', 'metal', 'organic', 'paper', 'plastic']  # Example
+# Define class names — replace with your actual folder names if different
+class_names = ['batteries', 'cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
-# Preprocess image
-def preprocess_image(image):
-    image = image.resize((224, 224))  # Adjust based on your model input
-    image = np.array(image) / 255.0   # Normalize if required
-    image = np.expand_dims(image, axis=0)  # Add batch dimension
-    return image
+# Define image size used in training
+IMG_SIZE = (224, 224)
 
-# Prediction function
-def classify_image(image):
-    image = preprocess_image(image)
-    prediction = model.predict(image)
-    class_index = np.argmax(prediction)
-    class_name = class_names[class_index]
-    confidence = prediction[0][class_index]
-    return f"{class_name} ({confidence*100:.2f}%)"
+def preprocess(image):
+    image = image.resize(IMG_SIZE)
+    image_array = tf.keras.preprocessing.image.img_to_array(image)
+    image_array = image_array / 255.0  # Normalize if you did this during training
+    return np.expand_dims(image_array, axis=0)
 
-# Gradio interface
-interface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Waste Image Classifier",
-    description="Upload an image of a recyclable item to classify it."
+def predict(image):
+    input_tensor = preprocess(image)
+    predictions = model.predict(input_tensor)[0]
+    top_class = class_names[np.argmax(predictions)]
+    confidence = np.max(predictions) * 100
+    return f"{top_class} ({confidence:.2f}%)"
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload an image of a recyclable item to classify it."),
+    outputs=gr.Textbox(label="Output"),
+    title="Waste Image Classifier"
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()