Repository: scikit-image/scikit-image

Issue:
IndexError with util.apply_parallel (RGB images)
I don't really understand the origin of this error. I could only reproduce with RGB images, not with 2D arrays. Does `apply_parallel` only work for 2D images?

## Way to reproduce
```python
from skimage import util, segmentation, data
img = data.astronaut()
_ = util.apply_parallel(segmentation.slic, img, chunks=None)
```

Output
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-3-63f67e7b9ccd> in <module>
----> 1 _ = util.apply_parallel(segmentation.slic, img, chunks=None)

~/code/scikit-image/skimage/util/apply_parallel.py in apply_parallel(function, array, chunks, depth, mode, extra_arguments, extra_keywords, compute)
    143     res = darr.map_overlap(wrapped_func, depth, boundary=mode)
    144     if compute:
--> 145         res = res.compute()
    146 
    147     return res

~/.local/lib/python3.7/site-packages/dask/base.py in compute(self, **kwargs)
    164         dask.base.compute
    165         """
--> 166         (result,) = compute(self, traverse=False, **kwargs)
    167         return result
    168 

~/.local/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)
    443 
    444     results = schedule(dsk, keys, **kwargs)
--> 445     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    446 
    447 

~/.local/lib/python3.7/site-packages/dask/base.py in <listcomp>(.0)
    443 
    444     results = schedule(dsk, keys, **kwargs)
--> 445     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    446 
    447 

~/.local/lib/python3.7/site-packages/dask/array/core.py in finalize(results)
    984     while isinstance(results2, (tuple, list)):
    985         if len(results2) > 1:
--> 986             return concatenate3(results)
    987         else:
    988             results2 = results2[0]

~/.local/lib/python3.7/site-packages/dask/array/core.py in concatenate3(arrays)
   4383     if not ndim:
   4384         return arrays
-> 4385     chunks = chunks_from_arrays(arrays)
   4386     shape = tuple(map(sum, chunks))
   4387 

~/.local/lib/python3.7/site-packages/dask/array/core.py in chunks_from_arrays(arrays)
   4152 
   4153     while isinstance(arrays, (list, tuple)):
-> 4154         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   4155         arrays = arrays[0]
   4156         dim += 1

~/.local/lib/python3.7/site-packages/dask/array/core.py in <listcomp>(.0)
   4152 
   4153     while isinstance(arrays, (list, tuple)):
-> 4154         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   4155         arrays = arrays[0]
   4156         dim += 1

IndexError: tuple index out of range
```

## Version information
```python
# Paste the output of the following python commands
from __future__ import print_function
import sys; print(sys.version)
import platform; print(platform.platform())
import skimage; print("scikit-image version: {}".format(skimage.__version__))
import numpy; print("numpy version: {}".format(numpy.__version__))
```

```python
# your output here
3.7.3 (default, Oct  7 2019, 12:56:13) 
[GCC 8.3.0]
Linux-5.0.0-38-generic-x86_64-with-Ubuntu-19.04-disco
scikit-image version: 0.18.dev0
numpy version: 1.16.4
```



Discussion:
First observation, `_get_chunks` slices the image in all dimensions, which is not expected if you pass a RGB image. It must work on 3D grayscale image though.
Thanks @sciunto. Does it mean that `apply_parallel` cannot be used on RGB images at all? If yes this is a serious flaw which should be either corrected or documented. 
As far as I understand, yes... I remember I fell on the same issue as you, but for some reason, I forgot to open a ticket :) 
I will take a look at this. You can avoid the call to `_get_chunks` by manually specifying the `chunks` argument, but we should fix _get_chunks to work properly for the multichannel case.

Here is an example that works if setting chunks manually. Note that depth=0 also needs to be specified along the channels dimension.
```Python
from functools import partial
from skimage import data, restoration
from skimage.util import apply_parallel
camera = data.camera() / 255.
cat = data.chelsea() / 255.
sigma = 0.05
cat_noisy = cat + sigma * np.random.standard_normal(cat.shape)
camera_noisy = camera + sigma * np.random.standard_normal(camera.shape)

func_rgb = partial(restoration.denoise_wavelet, sigma=sigma, multichannel=True, convert2ycbcr=False)
out_rgb = apply_parallel(func_rgb, cat_noisy, chunks=[128, 128, 3], depth=(8, 8, 0))
```

However, I also found another issue that can also occur when trying to use RGB images. In this case if we change one argument in the above example to `convert2ycbcr=True`, it will fail with:
```
ValueError: `dtype` inference failed in `map_blocks`.

Please specify the dtype explicitly using the `dtype` kwarg.
```

I was only able to fix this error by adding a `dtype=float` argument to the internal call to map overlap. The underlying issue appears to be that when the dtype is not specified, Dask does not know the expected dtype of the output array and so tries to run a minimal computation to determine it. Specifically, it calls the function but with an array of size `(1,) * ndim`. This is problematic when `convert2ycbcr=True` because there is an internal call to `color.rgb2ycbcr` which expects to receive an array of size 3 on the last axis. I would have assumed that a single block of the specified chunk size would be used to infer the dtype, but apparently that is not the case (see dask/dask#1949) and we will have to specify the dtype explicitly instead.






Resolving patch:
diff --git a/skimage/util/apply_parallel.py b/skimage/util/apply_parallel.py
--- a/skimage/util/apply_parallel.py
+++ b/skimage/util/apply_parallel.py
@@ -1,3 +1,6 @@
+import numpy
+
+
 __all__ = ['apply_parallel']
 
 
@@ -53,7 +56,8 @@ def _ensure_dask_array(array, chunks=None):
 
 
 def apply_parallel(function, array, chunks=None, depth=0, mode=None,
-                   extra_arguments=(), extra_keywords={}, *, compute=None):
+                   extra_arguments=(), extra_keywords={}, *, dtype=None,
+                   multichannel=False, compute=None):
     """Map a function in parallel across an array.
 
     Split an array into possibly overlapping chunks of a given depth and
@@ -84,6 +88,25 @@ def apply_parallel(function, array, chunks=None, depth=0, mode=None,
         Tuple of arguments to be passed to the function.
     extra_keywords : dictionary, optional
         Dictionary of keyword arguments to be passed to the function.
+    dtype : data-type or None, optional
+        The data-type of the `function` output. If None, Dask will attempt to
+        infer this by calling the function on data of shape ``(1,) * ndim``.
+        For functions expecting RGB or multichannel data this may be
+        problematic. In such cases, the user should manually specify this dtype
+        argument instead.
+
+        .. versionadded:: 0.18
+           ``dtype`` was added in 0.18.
+    multichannel : bool, optional
+        If `chunks` is None and `multichannel` is True, this function will keep
+        only a single chunk along the channels axis. When `depth` is specified
+        as a scalar value, that depth will be applied only to the non-channels
+        axes (a depth of 0 will be used along the channels axis). If the user
+        manually specified both `chunks` and a `depth` tuple, then this
+        argument will have no effect.
+
+        .. versionadded:: 0.18
+           ``multichannel`` was added in 0.18.
     compute : bool, optional
         If ``True``, compute eagerly returning a NumPy Array.
         If ``False``, compute lazily returning a Dask Array.
@@ -126,7 +149,10 @@ def apply_parallel(function, array, chunks=None, depth=0, mode=None,
             ncpu = cpu_count()
         except NotImplementedError:
             ncpu = 4
-        chunks = _get_chunks(shape, ncpu)
+        if multichannel:
+            chunks = _get_chunks(shape[:-1], ncpu) + (shape[-1],)
+        else:
+            chunks = _get_chunks(shape, ncpu)
 
     if mode == 'wrap':
         mode = 'periodic'
@@ -135,12 +161,16 @@ def apply_parallel(function, array, chunks=None, depth=0, mode=None,
     elif mode == 'edge':
         mode = 'nearest'
 
+    if multichannel and numpy.isscalar(depth):
+        # depth is only used along the non-channel axes
+        depth = (depth,) * (len(array.shape) - 1) + (0,)
+
     def wrapped_func(arr):
         return function(arr, *extra_arguments, **extra_keywords)
 
     darr = _ensure_dask_array(array, chunks=chunks)
 
-    res = darr.map_overlap(wrapped_func, depth, boundary=mode)
+    res = darr.map_overlap(wrapped_func, depth, boundary=mode, dtype=dtype)
     if compute:
         res = res.compute()
 


Test patch:
diff --git a/skimage/util/tests/test_apply_parallel.py b/skimage/util/tests/test_apply_parallel.py
--- a/skimage/util/tests/test_apply_parallel.py
+++ b/skimage/util/tests/test_apply_parallel.py
@@ -1,6 +1,7 @@
 import numpy as np
 
-from skimage._shared.testing import assert_array_almost_equal
+from skimage._shared.testing import assert_array_almost_equal, assert_equal
+from skimage import color, data
 from skimage.filters import threshold_local, gaussian
 from skimage.util.apply_parallel import apply_parallel
 
@@ -94,3 +95,19 @@ def wrapped(arr):
                             mode='nearest')
 
     assert_array_almost_equal(result, expected)
+
+
+@pytest.mark.parametrize('dtype', (np.float32, np.float64))
+@pytest.mark.parametrize('chunks', (None, (128, 128, 3)))
+@pytest.mark.parametrize('depth', (0, 8, (8, 8, 0)))
+def test_apply_parallel_rgb(depth, chunks, dtype):
+    cat = data.chelsea().astype(dtype) / 255.
+
+    func = color.rgb2ycbcr
+    cat_ycbcr_expected = func(cat)
+    cat_ycbcr = apply_parallel(func, cat, chunks=chunks, depth=depth,
+                               dtype=dtype, multichannel=True)
+
+    assert_equal(cat_ycbcr.dtype, cat.dtype)
+
+    assert_array_almost_equal(cat_ycbcr_expected, cat_ycbcr)