fix axis bug in multichannel case and add a corresponding test

scikit-image · Jan 24, 2016 · 53742be · 53742be
1 parent a30592a
commit 53742be
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/skimage/restoration/_denoise.py b/skimage/restoration/_denoise.py
@@ -276,7 +276,7 @@ def denoise_tv_chambolle(im, weight=0.1, eps=2.e-4, n_iter_max=200,
 
     if multichannel:
         out = np.zeros_like(im)
-        for c in range(im.shape[2]):
+        for c in range(im.shape[-1]):
             out[..., c] = _denoise_tv_chambolle_nd(im[..., c], weight, eps,
                                                    n_iter_max)
     else:

diff --git a/skimage/restoration/tests/test_denoise.py b/skimage/restoration/tests/test_denoise.py
@@ -39,6 +39,15 @@ def test_denoise_tv_chambolle_multichannel():
                                                 multichannel=True)
     assert_equal(denoised[..., 0], denoised0)
 
+    # tile astronaut subset to generate 3D+channels data
+    astro3 = np.tile(astro[:64, :64, np.newaxis, :], [1, 1, 2, 1])
+    # modify along tiled dimension to give non-zero gradient on 3rd axis
+    astro3[:, :, 0, :] = 2*astro3[:, :, 0, :]
+    denoised0 = restoration.denoise_tv_chambolle(astro3[..., 0], weight=0.1)
+    denoised = restoration.denoise_tv_chambolle(astro3, weight=0.1,
+                                                multichannel=True)
+    assert_equal(denoised[..., 0], denoised0)
+
 
 def test_denoise_tv_chambolle_float_result_range():
     # astronaut image