Update app.py

app.py

@@ -1,56 +1,42 @@
-import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Disable GPU warnings
 import gradio as gr
 import tensorflow as tf
 import numpy as np
 from PIL import Image
-import traceback
-
-# Load model
-model = tf.keras.models.load_model("MobileNet_model.h5")
-class_names = ["Fake", "Low", "Medium", "High"]
-
-def predict_image(img):
-    try:
-        if img is None:
-            return {"error": "No image provided"}
-            
-        # If img is a file path, open it
-        if isinstance(img, str):
-            img = Image.open(img)
-            
-        # Ensure it's a PIL Image
-        if not isinstance(img, Image.Image):
-            return {"error": f"Expected PIL Image, got {type(img)}"}
-            
-        img = img.resize((128, 128))
-        img_array = np.array(img) / 255.0
-        
-        # Handle grayscale images
-        if len(img_array.shape) == 2:
-            img_array = np.stack((img_array,) * 3, axis=-1)
-        # Handle RGBA images
-        elif img_array.shape[2] == 4:
-            img_array = img_array[:, :, :3]
-            
-        img_array = np.expand_dims(img_array, axis=0)
-        predictions = model.predict(img_array)
-        class_index = np.argmax(predictions, axis=1)[0]
-        confidence_scores = {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}
-        return {"Predicted Class": class_names[class_index], "Confidence Scores": confidence_scores}
-    except Exception as e:
-        return {"error": str(e), "traceback": traceback.format_exc()}
-
-# Create interface with explicit input type
-iface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.JSON(),
-    examples=[["example1.jpg"], ["example2.jpg"]],
-    title="Fire Detection API",
-    description="Upload an image to detect fire presence and intensity"
+
+# 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="Waste Classification API",
+    description="Send an image to classify it into one of four categories: Fake, Low, Medium, or High."
 )
 
-# Launch with debugging enabled
-iface.launch(server_name="0.0.0.0", server_port=7860, share=True, debug=True, show_error=True)
+# Launch API
+if __name__ == "__main__":
+    interface.launch(share=True)