Color example :)

doc/examples/plot_quickshift.py

@@ -1,47 +1,31 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from scipy import ndimage
-#from skimage.data import lena
-#from skimage.util import img_as_float
+from skimage.data import lena
 from skimage.segmentation import quickshift
+from skimage.util import img_as_float
 
 from IPython.core.debugger import Tracer
 tracer = Tracer()
 
 
-def microstructure(l=256):
-    """
-    Synthetic binary data: binary microstructure with blobs.
-
-    Parameters
-    ----------
-
-    l: int, optional
-        linear size of the returned image
-    """
-    n = 5
-    x, y = np.ogrid[0:l, 0:l]
-    mask = np.zeros((l, l))
-    generator = np.random.RandomState(1)
-    points = l * generator.rand(2, n ** 2)
-    mask[(points[0]).astype(np.int), (points[1]).astype(np.int)] = 1
-    mask = ndimage.gaussian_filter(mask, sigma=l / (4. * n))
-    return (mask > mask.mean()).astype(np.float)
-
-
-#img = img_as_float(lena()[250:300, 250:300])
-img = microstructure(l=50)
-segments = quickshift(img.reshape(50, 50, 1))
-segments = np.unique(segments, return_inverse=True)[1].reshape(50, 50)
-intensities = np.bincount(segments.ravel(), img.ravel())
-counts = np.bincount(segments.ravel())
-intensities /= counts
+img = img_as_float(lena())[::3, ::3, :].copy("C")
+segments = quickshift(img, sigma=2)
+segments = np.unique(segments, return_inverse=True)[1].reshape(img.shape[:2])
 
+plt.subplot(131, title="original")
 plt.imshow(img, interpolation='nearest')
-plt.figure()
-plt.imshow(segments, interpolation='nearest')
-plt.figure()
-plt.imshow(intensities[segments], 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/quickshift.pyx

@@ -3,11 +3,13 @@ cimport numpy as np
 
 from itertools import product
 
+from time import time
+
 cdef extern from "math.h":
     double exp(double)
 
 
-def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, tau=10):
+def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, tau=10, return_tree=False):
     """Computes quickshift clustering in RGB-(x,y) space.
 
     Parameters
@@ -20,6 +22,8 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
     tau: float
         Cut-off point for data distances.
         Higher means less clusters.
+    return_tree: bool
+        Whether to return the full segmentation hierarchy tree
 
     Returns
     -------
@@ -45,7 +49,7 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
     cdef int x, y, xx, yy, x_, y_
 
     cdef np.ndarray[dtype=np.float_t, ndim=2] densities = np.zeros((width, height))
-
+    start = time()
     # compute densities
     for x, y in product(xrange(width), xrange(height)):
         current_pixel = image[x, y, :]
@@ -57,6 +61,7 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
                     dist += (current_pixel[c] - image[x_, y_, c])**2
                 dist += (x - x_)**2 + (y - y_)**2
                 densities[x, y] += float(exp(-dist / sigma))
+    print("densities: %f" % (time() - start))
 
     # this will break ties that otherwise would give us headache
 
@@ -64,6 +69,7 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
     # default parent to self:
     cdef np.ndarray[dtype=np.int_t, ndim=2] parent = np.arange(width * height).reshape(width, height)
     cdef np.ndarray[dtype=np.float_t, ndim=2] dist_parent = np.zeros((width, height))
+    start = time()
     # find nearest node with higher density
     for x, y in product(xrange(width), xrange(height)):
         current_density = densities[x, y]
@@ -81,7 +87,9 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
                         closest = dist
                         parent[x, y] = x_ * width + y_
         dist_parent[x, y] = closest
+    print("parents: %f" % (time() - start))
 
+    start = time()
     dist_parent_flat = dist_parent.ravel()
     flat = parent.ravel()
     flat[dist_parent_flat > tau] = np.arange(width * height)[dist_parent_flat > tau]
@@ -89,4 +97,8 @@ def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, ta
     while (old != flat).any():
         old = flat
         flat = flat[flat]
-    return flat.reshape(width, height)
+    print("rest: %f" % (time() - start))
+    flat = flat.reshape(width, height)
+    if return_tree:
+        return flat, parent
+    return flat