Update app.py

app.py

@@ -1,41 +1,24 @@
+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
 
-# 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
+# Load model
+model = tf.keras.models.load_model("MobileNet_model.h5")
+class_names = ["Fake", "Low", "Medium", "High"]
 
-# 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}
+def predict_image(img):
+    img = img.resize((128, 128))
+    img_array = np.array(img) / 255.0
+    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}
 
-# 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."
-)
+iface = gr.Interface(fn=predict_image, inputs="image", outputs="json")
+iface.launch(server_name="0.0.0.0", server_port=7860)  # Fix for Hugging Face
 
-# Launch API
-if __name__ == "__main__":
-    interface.launch(share=True)