Update app.py

app.py

@@ -38,24 +38,17 @@ def predict(image):
     preprocessed_image = preprocess_image(image)
     prediction = run_inference(interpreter, preprocessed_image)
     
-    # Convert the prediction to float32 for softmax calculation
-    prediction = prediction.astype(np.float32)
-    
-    # Apply softmax to get probabilities
-    probabilities = tf.nn.softmax(prediction[0]).numpy()
-    
-    # Get class probabilities for Flood and NoFlood
-    flood_prob = probabilities[0]
-    no_flood_prob = probabilities[1]
-    
-    # Return the predicted label along with probabilities
-    return {"Flood": float(flood_prob), "NoFlood": float(no_flood_prob)}
+    # Get the predicted class by comparing the output values
+    if prediction[0][0] > prediction[0][1]:
+        return "Flood"
+    else:
+        return "No Flood"
 
 # Create Gradio interface
 interface = gr.Interface(
     fn=predict,  # Function that runs the model
     inputs=gr.Image(type="pil"),  # Input is an image, processed as PIL image
-    outputs=gr.Label(num_top_classes=2),  # Outputs a label with the predicted classes
+    outputs="text",  # Output is text indicating Flood or NoFlood
 )
 
 # Launch the app