ENH felzenszwalb for color images

doc/examples/plot_felzenszwalb_segmentation.py

@@ -3,11 +3,25 @@ import numpy as np
 
 from skimage.data import lena
 from skimage.segmentation import felzenszwalb_segmentation
+from skimage.util import img_as_float
 
-img = lena()
-segments = felzenszwalb_segmentation(img, k=1000)
-plt.imshow(img)
-plt.figure()
-plt.imshow(segments)
+img = img_as_float(lena())
+segments = felzenszwalb_segmentation(img, scale=1)
+segments = np.unique(segments, return_inverse=True)[1].reshape(img.shape[:2])
+
+plt.subplot(131, title="original")
+plt.imshow(img, interpolation='nearest')
+
+plt.subplot(132, title="superpixels")
+# shuffle the labels for better visualization
+permuted_labels = np.random.permutation(segments.max() + 1)
+plt.imshow(permuted_labels[segments], interpolation='nearest')
+
+plt.subplot(133, title="mean color")
+colors = [np.bincount(segments.ravel(), img[:, :, c].ravel()) for c in
+        xrange(img.shape[2])]
+counts = np.bincount(segments.ravel())
+colors = np.vstack(colors) / counts
+plt.imshow(colors.T[segments], interpolation='nearest')
 plt.show()
 print("num segments: %d" % len(np.unique(segments)))

skimage/segmentation/_felzenszwalb.pyx

@@ -51,11 +51,30 @@ cdef join_trees(np.int_t *forest, np.int_t n, np.int_t m):
         set_root(forest, m, root)
 
 
-def felzenszwalb_segmentation(image, k, sigma=0.8):
-    k = float(k)
-    #image = img_as_float(image)
-    #image = rgb2grey(image)
-    image = image[:, :, 0]
+def felzenszwalb_segmentation_grey(image, scale=200, sigma=0.8):
+    """Computes Felsenszwalb's efficient graph based segmentation for a single channel.
+
+    Parameters
+    ----------
+    image: ndarray, [width, height]
+        Input image
+
+    scale: float
+        Free parameter. Higher means larger clusters.
+        For 0-255 data, hundereds are good.
+
+    sigma: float
+        Width of Gaussian kernel used in preprocessing.
+
+    Returns
+    -------
+    segment_mask: ndarray, [width, height]
+        Integer mask indicating segment labels.
+    """
+    if image.ndim != 2:
+        raise ValueError("This algorithm works only on single-channel 2d images."
+                "Got image of shape %s" % str(image.shape))
+    scale = float(scale)
     image = scipy.ndimage.gaussian_filter(image, sigma=sigma)
 
     # compute edge weights in 8 connectivity:
@@ -96,8 +115,8 @@ def felzenszwalb_segmentation(image, k, sigma=0.8):
         costs_p += 1
         if seg0 == seg1:
             continue
-        inner_cost0 = cint[seg0] + k / segment_size[seg0]
-        inner_cost1 = cint[seg1] + k / segment_size[seg1]
+        inner_cost0 = cint[seg0] + scale / segment_size[seg0]
+        inner_cost1 = cint[seg1] + scale / segment_size[seg1]
         if costs_p[0] < min(inner_cost0, inner_cost1):
             # update size and cost
             join_trees(segments_p, seg0, seg1)

skimage/segmentation/felzenszwalb.py

@@ -0,0 +1,67 @@
+import warnings
+import numpy as np
+
+from ._felzenszwalb import felzenszwalb_segmentation_grey
+
+from IPython.core.debugger import Tracer
+tracer = Tracer()
+
+
+def felzenszwalb_segmentation(image, scale=200, sigma=0.8):
+    """Computes Felsenszwalb's segmentation for multi channel images.
+
+    Calls the algorithm on each channel separately, then combines
+    using "and", i.e. two pixels are in the same segment if they are
+    in the same segment for each channel.
+
+    Parameters
+    ----------
+    image: ndarray, [width, height]
+        Input image
+
+    scale: float
+        Free parameter. Higher means larger clusters.
+        For 0-255 data, hundereds are good.
+
+    sigma: float
+        Width of Gaussian kernel used in preprocessing.
+
+    Returns
+    -------
+    segment_mask: ndarray, [width, height]
+        Integer mask indicating segment labels.
+    """
+
+    #image = img_as_float(image)
+    if image.ndim == 2:
+        # assume single channel image
+        return felzenszwalb_segmentation_grey(image, scale=scale, sigma=sigma)
+
+    elif image.ndim != 3:
+        raise ValueError("Got image with ndim=%d, don't know"
+                " what to do." % image.ndim)
+
+    # assume we got 2d image with multiple channels
+    n_channels = image.shape[2]
+    if n_channels != 3:
+        warnings.warn("Got image with %d channels. Is that really what you"
+                " wanted?" % image.shape[2])
+    segmentations = []
+    # compute quickshift for each channel
+    for c in xrange(n_channels):
+        channel = np.ascontiguousarray(image[:, :, c])
+        seg = felzenszwalb_segmentation_grey(channel, scale=scale, sigma=sigma)
+        segmentations.append(seg)
+
+    # put pixels in same segment only if in the same segment in all images
+    # we do this by combining the channels to one number
+    segmentations = [np.unique(s, return_inverse=True)[1] for s in
+            segmentations]
+    n0 = max(segmentations[0])
+    n1 = max(segmentations[1])
+    hasher = np.array([n1 * n0, n0, 1])
+    segmentations = np.dstack(segmentations).reshape(-1, n_channels)
+    segmentation = np.dot(segmentations, hasher)
+    # make segment labels consecutive numbers starting at 0
+    labels = np.unique(segmentation, return_inverse=True)[1]
+    return labels.reshape(image.shape[:2])

skimage/segmentation/setup.py

@@ -11,8 +11,8 @@ def configuration(parent_package='', top_path=None):
 
     config = Configuration('segmentation', parent_package, top_path)
 
-    cython(['felzenszwalb.pyx'], working_path=base_path)
-    config.add_extension('felzenszwalb', sources=['felzenszwalb.c'],
+    cython(['_felzenszwalb.pyx'], working_path=base_path)
+    config.add_extension('_felzenszwalb', sources=['_felzenszwalb.c'],
                          include_dirs=[get_numpy_include_dirs()])
     cython(['quickshift.pyx'], working_path=base_path)
     config.add_extension('quickshift', sources=['quickshift.c'],