Improve SLIC

skimage/segmentation/_slic.pyx

@@ -2,31 +2,32 @@
 #cython: boundscheck=False
 #cython: nonecheck=False
 #cython: wraparound=False
-import numpy as np
-from scipy import ndimage
+from libc.float cimport DBL_MAX
 
-from skimage.util import img_as_float, regular_grid
-from skimage.color import rgb2lab, gray2rgb
+import numpy as np
+cimport numpy as cnp
+
+from skimage.util import regular_grid
 
 
 def _slic_cython(double[:, :, :, ::1] image_zyx,
                  Py_ssize_t[:, :, ::1] nearest_mean,
                  double[:, :, ::1] distance,
-                 double[:, ::1] means,
+                 double[:, ::1] clusters,
                  Py_ssize_t max_iter, Py_ssize_t n_segments):
     """Helper function for SLIC segmentation.
 
     Parameters
     ----------
-    image_zyx : 4D array of double, shape (Z, Y, X, 6)
+    image_zyx : 4D array of double, shape (Z, Y, X, C)
         The image with embedded coordinates, that is, `image_zyx[i, j, k]` is
-        `array([i, j, k, r, g, b])` or `array([i, j, k, L, a, b])`, depending
+        `array([i, j, k, c])`, depending
         on the colorspace.
     nearest_mean : 3D array of int, shape (Z, Y, X)
         The (initially empty) label field.
     distance : 3D array of double, shape (Z, Y, X)
         The (initially infinity) array of distances to the nearest centroid.
-    means : 2D array of double, shape (n_segments, 6)
+    clusters : 2D array of double, shape (n_segments, 6)
         The centroids obtained by SLIC.
     max_iter : int
         The maximum number of k-means iterations.
@@ -43,36 +44,43 @@ def _slic_cython(double[:, :, :, ::1] image_zyx,
     cdef Py_ssize_t depth, height, width
     depth, height, width = (image_zyx.shape[0], image_zyx.shape[1],
                             image_zyx.shape[2])
+
+    cdef Py_ssize_t n_features = clusters.shape[1]
+    cdef Py_ssize_t n_clusters = clusters.shape[0]
+
     # approximate grid size for desired n_segments
     cdef Py_ssize_t step_z, step_y, step_x
     slices = regular_grid((depth, height, width), n_segments)
     step_z, step_y, step_x = [int(s.step) for s in slices]
 
-    n_means = means.shape[0]
     cdef Py_ssize_t i, k, x, y, z, x_min, x_max, y_min, y_max, z_min, z_max, \
             changes
     cdef double dist_mean
     cdef double tmp
+
+    cdef Py_ssize_t[:] n_cluster_elems = np.zeros(n_clusters, dtype=np.intp)
+
     for i in range(max_iter):
         changes = 0
-        distance[:, :, :] = np.inf
-        # assign pixels to means
-        for k in range(n_means):
+        distance[:, :, :] = DBL_MAX
+
+        # assign pixels to clusters
+        for k in range(n_clusters):
             # compute windows:
-            z_min = int(max(means[k, 0] - 2 * step_z, 0))
-            z_max = int(min(means[k, 0] + 2 * step_z, depth))
-            y_min = int(max(means[k, 1] - 2 * step_y, 0))
-            y_max = int(min(means[k, 1] + 2 * step_y, height))
-            x_min = int(max(means[k, 2] - 2 * step_x, 0))
-            x_max = int(min(means[k, 2] + 2 * step_x, width))
+            z_min = int(max(clusters[k, 0] - 2 * step_z, 0))
+            z_max = int(min(clusters[k, 0] + 2 * step_z, depth))
+            y_min = int(max(clusters[k, 1] - 2 * step_y, 0))
+            y_max = int(min(clusters[k, 1] + 2 * step_y, height))
+            x_min = int(max(clusters[k, 2] - 2 * step_x, 0))
+            x_max = int(min(clusters[k, 2] + 2 * step_x, width))
             for z in range(z_min, z_max):
                 for y in range(y_min, y_max):
                     for x in range(x_min, x_max):
                         dist_mean = 0
-                        for c in range(6):
+                        for c in range(n_features):
                             # you would think the compiler can optimize the
                             # squaring itself. mine can't (with O2)
-                            tmp = image_zyx[z, y, x, c] - means[k, c]
+                            tmp = image_zyx[z, y, x, c] - clusters[k, c]
                             dist_mean += tmp * tmp
                         if distance[z, y, x] > dist_mean:
                             nearest_mean[z, y, x] = k
@@ -80,15 +88,23 @@ def _slic_cython(double[:, :, :, ::1] image_zyx,
                             changes = 1
         if changes == 0:
             break
-        # recompute means:
-        nearest_mean_ravel = np.asarray(nearest_mean).ravel()
-        means_list = []
-        for j in range(6):
-            image_zyx_ravel = (
-                        np.ascontiguousarray(image_zyx[:, :, :, j]).ravel())
-            means_list.append(np.bincount(nearest_mean_ravel,
-                                          image_zyx_ravel))
-        in_mean = np.bincount(nearest_mean_ravel)
-        in_mean[in_mean == 0] = 1
-        means = (np.vstack(means_list) / in_mean).T.copy("C")
-    return np.ascontiguousarray(nearest_mean)
+
+        # recompute clusters
+
+        # sum features for all clusters
+        n_cluster_elems[:] = 0
+        clusters[:, :] = 0
+        for z in range(depth):
+            for y in range(height):
+                for x in range(width):
+                    k = nearest_mean[z, y, x]
+                    n_cluster_elems[k] += 1
+                    for c in range(n_features):
+                        clusters[k, c] += image_zyx[z, y, x, c]
+
+        # divide by number of elements per cluster to obtain mean
+        for k in range(n_clusters):
+            for c in range(n_features):
+                clusters[k, c] /= n_cluster_elems[k]
+
+    return np.asarray(nearest_mean)

skimage/segmentation/slic_superpixels.py

@@ -21,7 +21,7 @@ def slic(image, n_segments=100, compactness=10., max_iter=10, sigma=1,
         (see `multichannel` parameter).
     n_segments : int, optional (default: 100)
         The (approximate) number of labels in the segmented output image.
-    compactness: float, optional (default: 10)
+    compactness : float, optional (default: 10)
         Balances color-space proximity and image-space proximity. Higher
         values give more weight to image-space. As `compactness` tends to
         infinity, superpixel shapes become square/cubic.
@@ -37,14 +37,14 @@ def slic(image, n_segments=100, compactness=10., max_iter=10, sigma=1,
         array has shape (M, N, 3).
     convert2lab : bool, optional (default: True)
         Whether the input should be converted to Lab colorspace prior to
-        segmentation.  For this purpose, the input is assumed to be RGB. Highly
+        segmentation. For this purpose, the input is assumed to be RGB. Highly
         recommended.
     ratio : float, optional
         Synonym for `compactness`. This keyword is deprecated.
 
     Returns
     -------
-    segment_mask : (width, height) ndarray
+    segment_mask : (width, height, depth) array
         Integer mask indicating segment labels.
 
     Raises
@@ -99,34 +99,41 @@ def slic(image, n_segments=100, compactness=10., max_iter=10, sigma=1,
             (multichannel and image.ndim not in [3, 4]) or
             (multichannel and image.shape[-1] != 3)):
         ValueError("Only 1- or 3-channel 2- or 3-D images are supported.")
+
     image = img_as_float(image)
-    if not multichannel:
-        image = gray2rgb(image)
+    image = np.atleast_3d(image)
+
     if image.ndim == 3:
         # See 2D RGB image as 3D RGB image with Z = 1
         image = image[np.newaxis, ...]
+
     if not isinstance(sigma, coll.Iterable):
         sigma = np.array([sigma, sigma, sigma, 0])
     if (sigma > 0).any():
         image = ndimage.gaussian_filter(image, sigma)
-    if convert2lab:
+
+    if image.shape[3] == 3 and convert2lab:
         image = rgb2lab(image)
 
-    # initialize on grid:
+    # initialize on grid
     depth, height, width = image.shape[:3]
+
     # approximate grid size for desired n_segments
     grid_z, grid_y, grid_x = np.mgrid[:depth, :height, :width]
     slices = regular_grid(image.shape[:3], n_segments)
     step_z, step_y, step_x = [int(s.step) for s in slices]
-    means_z = grid_z[slices]
-    means_y = grid_y[slices]
-    means_x = grid_x[slices]
+    clusters_z = grid_z[slices]
+    clusters_y = grid_y[slices]
+    clusters_x = grid_x[slices]
+
+    clusters_color = np.zeros(clusters_z.shape + (image.shape[3],))
+    clusters = np.concatenate([clusters_z[..., np.newaxis],
+                               clusters_y[..., np.newaxis],
+                               clusters_x[..., np.newaxis],
+                               clusters_color
+                              ], axis=-1).reshape(-1, 3 + image.shape[3])
+    clusters = np.ascontiguousarray(clusters)
 
-    means_color = np.zeros(means_z.shape + (3,))
-    means = np.concatenate([means_z[..., np.newaxis], means_y[..., np.newaxis],
-                            means_x[..., np.newaxis], means_color
-                           ], axis=-1).reshape(-1, 6)
-    means = np.ascontiguousarray(means)
     # we do the scaling of ratio in the same way as in the SLIC paper
     # so the values have the same meaning
     ratio = float(max((step_z, step_y, step_x))) / compactness
@@ -134,9 +141,9 @@ def slic(image, n_segments=100, compactness=10., max_iter=10, sigma=1,
                                 grid_y[..., np.newaxis],
                                 grid_x[..., np.newaxis],
                                 image * ratio], axis=-1).copy("C")
-    nearest_mean = np.zeros((depth, height, width), dtype=np.intp)
+    nearest_cluster = np.empty((depth, height, width), dtype=np.intp)
     distance = np.empty((depth, height, width), dtype=np.float)
-    segment_map = _slic_cython(image_zyx, nearest_mean, distance, means,
+    segment_map = _slic_cython(image_zyx, nearest_cluster, distance, clusters,
                                max_iter, n_segments)
     if segment_map.shape[0] == 1:
         segment_map = segment_map[0]

skimage/segmentation/tests/test_slic.py

@@ -28,7 +28,7 @@ def test_color_2d():
 
 def test_gray_2d():
     rnd = np.random.RandomState(0)
-    img = np.zeros((20, 21))
+    img = np.zeros((20, 20))
     img[:10, :10] = 0.33
     img[10:, :10] = 0.67
     img[10:, 10:] = 1.00