Add 3D slic tests (gray not working yet)

skimage/segmentation/tests/test_slic.py

@@ -1,9 +1,10 @@
+import itertools as it
 import numpy as np
 from numpy.testing import assert_equal, assert_array_equal
 from skimage.segmentation import slic
 
 
-def test_color():
+def test_color_2d():
     rnd = np.random.RandomState(0)
     img = np.zeros((20, 21, 3))
     img[:10, :10, 0] = 1
@@ -21,7 +22,8 @@ def test_color():
     assert_array_equal(seg[:10, 10:], 1)
     assert_array_equal(seg[10:, 10:], 3)
 
-def test_gray():
+
+def test_gray_2d():
     rnd = np.random.RandomState(0)
     img = np.zeros((20, 21))
     img[:10, :10] = 0.33
@@ -38,6 +40,49 @@ def test_gray():
     assert_array_equal(seg[:10, 10:], 1)
     assert_array_equal(seg[10:, 10:], 3)
 
+
+def test_color_3d():
+    rnd = np.random.RandomState(0)
+    img = np.zeros((20, 21, 22, 3))
+    slices = []
+    for dim_size in img.shape[:-1]:
+        midpoint = dim_size // 2
+        slices.append((slice(None, midpoint), slice(midpoint, None)))
+    slices = list(it.product(*slices))
+    colors = list(it.product(*(([0, 1],) * 3)))
+    for s, c in zip(slices, colors):
+        img[s] = c
+    img += 0.01 * rnd.normal(size=img.shape)
+    img[img > 1] = 1
+    img[img < 0] = 0
+    seg = slic(img, sigma=0, n_segments=8)
+
+    assert_equal(len(np.unique(seg)), 8)
+    for s, c in zip(slices, range(8)):
+        assert_array_equal(seg[s], c)
+
+
+def test_gray_3d():
+    rnd = np.random.RandomState(0)
+    img = np.zeros((20, 21, 22))
+    slices = []
+    for dim_size in img.shape[:-1]:
+        midpoint = dim_size // 2
+        slices.append((slice(None, midpoint), slice(midpoint, None)))
+    slices = list(it.product(*slices))
+    shades = np.arange(0, 1.000001, 1.0/7)
+    for s, sh in zip(slices, shades):
+        img[s] = sh
+    img += 0.001 * rnd.normal(size=img.shape)
+    img[img > 1] = 1
+    img[img < 0] = 0
+    seg = slic(img, sigma=0, n_segments=8, ratio=40.0, multichannel=False)
+
+    assert_equal(len(np.unique(seg)), 8)
+    for s, c in zip(slices, range(8)):
+        assert_array_equal(seg[s], c)
+
+
 if __name__ == '__main__':
     from numpy import testing
     testing.run_module_suite()