Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-import cv2  # Add this import statement
+import cv2
 from keras.utils import normalize
 
 def dice_coef(y_true, y_pred):
@@ -14,28 +14,21 @@ def predict_segmentation(image):
     SIZE_X = 128
     SIZE_Y = 128
 
-    train_images = []
-   
+    # The image is already a NumPy array, resize and convert to grayscale if needed
     img = cv2.resize(image, (SIZE_Y, SIZE_X))
-    train_images.append(img)
+    if len(img.shape) == 3 and img.shape[2] == 3:  # If the image is RGB
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # Convert to grayscale
 
-    train_images = np.array(train_images)
-    train_images = np.expand_dims(train_images, axis=3)
-    train_images = normalize(train_images, axis=1)
-    X_test = train_images
+    img = np.expand_dims(img, axis=2)  # Add the channel dimension
+    img = normalize(img, axis=1)
+    X_test = np.expand_dims(img, axis=0)  # Add the batch dimension
 
     custom_objects = {'dice_coef': dice_coef}
     with tf.keras.utils.custom_object_scope(custom_objects):
         model = tf.keras.models.load_model("model100.h5")
-        
-    
-    # Normalize the test image
-    test_img = X_test[0]
-    test_img_norm = test_img[:, :, 0][:, :, None]
-    test_img_input = np.expand_dims(test_img_norm, 0)
 
     # Get the prediction
-    prediction = model.predict(test_img_input)
+    prediction = model.predict(X_test)
     predicted_img = np.argmax(prediction, axis=3)[0, :, :]
 
     # Create an RGB image with a transparent background
@@ -49,12 +42,10 @@ def predict_segmentation(image):
         rgba_img[:, :, i] = np.where(predicted_img > 0, segmented_color[i], 0)
 
     # Create an alpha channel: 1 where there is segmentation, 0 otherwise
-    alpha_channel = np.where(predicted_img > 0, 1, 0)
-    rgba_img[:, :, 3] = alpha_channel
+    rgba_img[:, :, 3] = np.where(predicted_img > 0, 1, 0)
 
     return rgba_img
 
-
 # Gradio Interface
 iface = gr.Interface(
     fn=predict_segmentation,
@@ -63,4 +54,4 @@ iface = gr.Interface(
     live=False
 )
 
-iface.launch()
+iface.launch(share=True)