Update app.py

app.py

@@ -1,35 +1,37 @@
-import cv2
 import gradio as gr
+import cv2
 import numpy as np
 
-# Define a function for marker-based segmentation using the OpenCV watershed algorithm
-def watershed_segmentation(input_image, markers_image):
-    # Convert input images to numpy arrays
-    input_image = np.array(input_image)
-    markers_image = np.array(markers_image)
+# Function to perform image segmentation using OpenCV Watershed Algorithm
+def watershed_segmentation(input_image, scribble_image):
+    # Load the input image and scribble image
+    image = cv2.cvtColor(input_image.astype('uint8'), cv2.COLOR_RGBA2BGR)
+    scribble = cv2.cvtColor(scribble_image.astype('uint8'), cv2.COLOR_RGBA2GRAY)
+
+    # Convert scribble to markers (0 for background, 1 for unknown, 2 for foreground)
+    markers = np.zeros_like(scribble)
+    markers[scribble == 0] = 0
+    markers[scribble == 255] = 1
+    markers[scribble == 128] = 2
+
+    # Apply watershed algorithm
+    cv2.watershed(image, markers)
+
+    # Create a segmented mask
+    segmented_mask = np.zeros_like(image, dtype=np.uint8)
+    segmented_mask[markers == 2] = [0, 0, 255]  # Red color for segmented regions
 
-    # Load the marker image (binary image with markers)
-    markers = cv2.cvtColor(markers_image, cv2.COLOR_BGR2GRAY)
-    
-    # Apply the watershed algorithm
-    image_copy = input_image.copy()
-    cv2.watershed(image_copy, markers)
-    
-    # Apply color mapping to the segmented regions
-    segmented_image = cv2.applyColorMap(markers, cv2.COLORMAP_JET)
-    
-    return segmented_image
+    return segmented_mask
 
-# Define Gradio interfaces for the input and output
-input_image = gr.inputs.Image(type="numpy")
-markers_image = gr.inputs.Image(type="numpy")
-output_image = gr.outputs.Image()
+# Gradio interface
+input_image = gr.inputs.Image(type='pil', label="Upload an image")
+scribble_image = gr.inputs.Image(type='pil', label="Scribble on the image")
+output_image = gr.outputs.Image(type='pil', label="Segmented Image")
 
-# Create a Gradio app
 gr.Interface(
     fn=watershed_segmentation,
-    inputs=[input_image, markers_image],
+    inputs=[input_image, scribble_image],
     outputs=output_image,
-    title="Marker-Based Segmentation with Watershed Algorithm",
-    description="Upload an image and a marker image to perform marker-based segmentation.",
-).launch()
+    title="Image Segmentation using Watershed Algorithm",
+    description="Upload an image and scribble on it to perform segmentation using the Watershed Algorithm.",
+).launch(share=True)