Spaces:
Sleeping
Sleeping
Commit
•
9754079
1
Parent(s):
c39487a
add watershed
Browse files
app.py
CHANGED
@@ -105,6 +105,20 @@ def segmentation_inference(img, algorithm, scale):
|
|
105 |
compactness=0.00001,
|
106 |
sigma=1,
|
107 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
108 |
else:
|
109 |
raise ValueError("algorithm should be kmeans, felzenszwalb or slic")
|
110 |
print("Done")
|
@@ -172,7 +186,7 @@ with gr.Blocks() as demo:
|
|
172 |
with gr.Row():
|
173 |
with gr.Column():
|
174 |
image_input = gr.Image(label="Input Image")
|
175 |
-
algorithm = gr.Radio(['KMeans', 'Felzenszwalb', 'SLIC'], label="Algorithm", value="SLIC")
|
176 |
scale = gr.Slider(0.1, 1.0, 0.5, label="Scale")
|
177 |
run_button = gr.Button("Run")
|
178 |
with gr.Column():
|
|
|
105 |
compactness=0.00001,
|
106 |
sigma=1,
|
107 |
)
|
108 |
+
elif algorithm == 'watershed':
|
109 |
+
from skimage.segmentation import watershed
|
110 |
+
from skimage.feature import peak_local_max
|
111 |
+
from scipy import ndimage as ndi
|
112 |
+
sobel = busam.sobel_from_pred(pred, size)
|
113 |
+
sobel = sobel.cpu().numpy() if isinstance(sobel, torch.Tensor) else sobel
|
114 |
+
# contrast stretch sobel with 5% largest
|
115 |
+
sobel = np.clip(sobel / np.percentile(sobel, 95), 0, 1)
|
116 |
+
distance = ndi.distance_transform_edt(sobel < 1) # distance to the borders
|
117 |
+
coords = peak_local_max(distance, min_distance=int(1+100*scale), labels=sobel<1)
|
118 |
+
mask = np.zeros(sobel.shape, dtype=bool)
|
119 |
+
mask[tuple(coords.T)] = True
|
120 |
+
markers, _ = ndi.label(mask)
|
121 |
+
labels = watershed(sobel, markers)
|
122 |
else:
|
123 |
raise ValueError("algorithm should be kmeans, felzenszwalb or slic")
|
124 |
print("Done")
|
|
|
186 |
with gr.Row():
|
187 |
with gr.Column():
|
188 |
image_input = gr.Image(label="Input Image")
|
189 |
+
algorithm = gr.Radio(['KMeans', 'Felzenszwalb', 'SLIC', 'Watershed'], label="Algorithm", value="SLIC")
|
190 |
scale = gr.Slider(0.1, 1.0, 0.5, label="Scale")
|
191 |
run_button = gr.Button("Run")
|
192 |
with gr.Column():
|