ENH: Add grayscale support and color adjustment to visualize_boundaries

skimage/segmentation/boundaries.py

@@ -1,19 +1,39 @@
 import numpy as np
 from ..morphology import dilation, square
 from ..util import img_as_float
+from ..color import gray2rgb
 
 
 def find_boundaries(label_img):
+    """Return bool array where boundaries between labeled regions are True."""
     boundaries = np.zeros(label_img.shape, dtype=np.bool)
     boundaries[1:, :] += label_img[1:, :] != label_img[:-1, :]
     boundaries[:, 1:] += label_img[:, 1:] != label_img[:, :-1]
     return boundaries
 
 
-def visualize_boundaries(img, label_img):
-    img = img_as_float(img, force_copy=True)
+def visualize_boundaries(image, label_img, color=(1, 1, 0), outline_color=None):
+    """Return image with boundaries between labeled regions highlighted.
+
+    Parameters
+    ----------
+    image : (M, N[, 3]) array
+        Grayscale or RGB image.
+    label_img : (M, N) array
+        Label array where regions are marked by different integer values.
+    color : length-3 sequence
+        RGB color of boundaries in the output image.
+    outline_color : length-3 sequence
+        RGB color surrounding boundaries in the output image. If None, no
+        outline is drawn.
+    """
+    if image.ndim == 2:
+        image = gray2rgb(image)
+    image = img_as_float(image, force_copy=True)
+
     boundaries = find_boundaries(label_img)
-    outer_boundaries = dilation(boundaries.astype(np.uint8), square(2))
-    img[outer_boundaries != 0, :] = np.array([0, 0, 0])  # black
-    img[boundaries, :] = np.array([1, 1, 0])  # yellow
-    return img
+    if outline_color is not None:
+        outer_boundaries = dilation(boundaries.astype(np.uint8), square(2))
+        image[outer_boundaries != 0, :] = np.array(outline_color)
+    image[boundaries, :] = np.array(color)
+    return image