Update app.py

app.py

@@ -3,40 +3,23 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# Load the trained model
-model = tf.keras.models.load_model("MobileNet_model.h5")  # Ensure the model file is uploaded in the same directory
-
-# Define class names
-class_names = ["Fake", "Low", "Medium", "High"]  # Modify if needed
-
-# Image Preprocessing Function
-img_size = (128, 128)  # Ensure it matches the input size used during training
-
-def preprocess_image(image):
-    image = image.resize(img_size)  # Resize image
-    image = np.array(image) / 255.0  # Normalize as done in ImageDataGenerator (rescale=1./255)
-    image = np.expand_dims(image, axis=0)  # Add batch dimension
-    return image
-
-# API Endpoint for Prediction
-def predict(image):
-    image = preprocess_image(image)
-    predictions = model.predict(image)
-    predicted_class = np.argmax(predictions, axis=1)[0]  # Get predicted class index
-    confidence_scores = {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}  # Get probability scores
-    
-    return {"Predicted Class": class_names[predicted_class], "Confidence Scores": confidence_scores}
-
-# Gradio API Interface
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),  # Accept image as input
-    outputs=gr.JSON(),  # Return JSON response
-    title="Fire Detection API",
-    description="Send an image to classify it into one of four categories: Fake, Low, Medium, or High."
-)
-
-# Launch API
-if __name__ == "__main__":
-    interface.launch(share=True)
-
+# Load model
+model = tf.keras.models.load_model("Mobilenet_model.h5")
+
+# Define class labels
+class_names = ["Organic", "Recyclable", "Hazardous"]
+
+def classify_image(image):
+    image = image.resize((128, 128))  # Resize to match input size
+    img_array = np.array(image) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+    predictions = model.predict(img_array)
+    class_index = np.argmax(predictions[0])
+    return class_names[class_index]
+
+gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Waste Classification Model",
+).launch()