Update app.py

app.py

@@ -14,21 +14,66 @@ def predict_segmentation(image):
     SIZE_X = 128
     SIZE_Y = 128
 
-    img = cv2.resize(image, (SIZE_Y, SIZE_X))
-    img = np.expand_dims(img, axis=2)
-    img = normalize(img, axis=1)
 
-    # Prepare image for prediction
-    img = np.expand_dims(img, axis=0)
+    train_images = []
 
-    # Predict
-    prediction = model.predict(img)
+    img = cv2.imread(image, 0)       
+    img = cv2.resize(img, (SIZE_Y, SIZE_X))
+    train_images.append(img)
+
+    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
+
+    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)
     predicted_img = np.argmax(prediction, axis=3)[0, :, :]
 
-    return predicted_img
+    # Create an RGB image with a transparent background
+    rgba_img = np.zeros((predicted_img.shape[0], predicted_img.shape[1], 4))
+
+    # Define the color for the segmented area (e.g., red)
+    segmented_color = [1, 0, 0]  # Red color in RGB
+
+    # Set the segmented area to the desired color
+    for i in range(3):
+        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
+
+
+
+
+
+
+    
+
+    # img = cv2.resize(image, (SIZE_Y, SIZE_X))
+    # img = np.expand_dims(img, axis=2)
+    # img = normalize(img, axis=1)
+
+    # # Prepare image for prediction
+    # img = np.expand_dims(img, axis=0)
+
+    # # Predict
+    # prediction = model.predict(img)
+    # predicted_img = np.argmax(prediction, axis=3)[0, :, :]
+
+    return rgba_img
 
-# Load the model
-model = tf.keras.models.load_model("model100.h5", custom_objects={'dice_coef': dice_coef})
 
 # Gradio Interface
 iface = gr.Interface(