add watershed

app.py

@@ -105,6 +105,20 @@ def segmentation_inference(img, algorithm, scale):
             compactness=0.00001,
             sigma=1,
         )
+    elif algorithm == 'watershed':
+        from skimage.segmentation import watershed
+        from skimage.feature import peak_local_max
+        from scipy import ndimage as ndi
+        sobel = busam.sobel_from_pred(pred, size)
+        sobel = sobel.cpu().numpy() if isinstance(sobel, torch.Tensor) else sobel
+        # contrast stretch sobel with 5% largest
+        sobel = np.clip(sobel / np.percentile(sobel, 95), 0, 1)
+        distance = ndi.distance_transform_edt(sobel < 1)  # distance to the borders
+        coords = peak_local_max(distance, min_distance=int(1+100*scale), labels=sobel<1)
+        mask = np.zeros(sobel.shape, dtype=bool)
+        mask[tuple(coords.T)] = True
+        markers, _ = ndi.label(mask)
+        labels = watershed(sobel, markers)
     else:
         raise ValueError("algorithm should be kmeans, felzenszwalb or slic")
     print("Done")
@@ -172,7 +186,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 image_input = gr.Image(label="Input Image")
-                algorithm = gr.Radio(['KMeans', 'Felzenszwalb', 'SLIC'], label="Algorithm", value="SLIC")
+                algorithm = gr.Radio(['KMeans', 'Felzenszwalb', 'SLIC', 'Watershed'], label="Algorithm", value="SLIC")
                 scale = gr.Slider(0.1, 1.0, 0.5, label="Scale")
                 run_button = gr.Button("Run")
             with gr.Column():