Make rotate available

dask_image/ndinterp/__init__.py

@@ -3,6 +3,7 @@
 import functools
 import math
 from itertools import product
+from numbers import Number
 import warnings
 
 import dask.array as da
@@ -11,6 +12,7 @@
 from dask.highlevelgraph import HighLevelGraph
 import scipy
 from scipy.ndimage import affine_transform as ndimage_affine_transform
+from scipy import special
 
 from ..dispatch._dispatch_ndinterp import (
     dispatch_affine_transform,
@@ -25,6 +27,7 @@
 
 __all__ = [
     "affine_transform",
+    "rotate"
 ]
 
 
@@ -247,6 +250,207 @@ def affine_transform(
     return transformed
 
 
+def rotate(input_arr, angle, axes=(1, 0), reshape=True, output=None, order=1,
+           mode='constant', cval=0.0, prefilter=False, output_chunks=None, output_shape=None):
+    """Rotate an array using Dask.
+
+    The array is rotated in the plane defined by the two axes given by the
+    `axes` parameter using spline interpolation of the requested order.
+    Chunkwise processing is performed using dask_image.ndinterp.affine_transform.
+
+    Parameters
+    ----------
+    input_arr : array_like (Numpy Array, Cupy Array, Dask Array...)
+        The image array.
+    angle : float
+        The rotation angle in degrees.
+    axes : tuple of 2 ints, optional
+        The two axes that define the plane of rotation. Default is the first
+        two axes.
+    reshape : bool, optional
+        If `reshape` is true, the output shape is adapted so that the input
+        array is contained completely in the output. Default is True.
+    output : dtype, optional
+        The dtype of the returned array.
+        By default, an array of the same dtype as input will be created.
+
+    order : int, optional
+        The order of the spline interpolation, default is 1.
+        The order has to be in the range 0-5. Note that for order>1
+        scipy's affine_transform applies prefiltering, which is not
+        yet supported and skipped in this implementation.
+    mode : {‘constant’, ‘grid-constant’, ‘nearest’}, optional
+        The mode parameter determines how the input array is extended beyond its boundaries.
+        Default is ‘constant’. Behavior for each valid value is as follows (see additional plots and details on boundary modes):
+
+        ‘constant’ (k k k k | a b c d | k k k k)
+            The input is extended by filling all values beyond the edge with the same constant value, defined by the cval parameter. No interpolation is performed beyond the edges of the input.
+        ‘grid-constant’ (k k k k | a b c d | k k k k)
+            The input is extended by filling all values beyond the edge with the same constant value, defined by the cval parameter. Interpolation occurs for samples outside the input’s extent as well.
+        ‘nearest’ (a a a a | a b c d | d d d d)
+            The input is extended by replicating the last pixel.
+    cval : scalar, optional
+        Value to fill past edges of input if mode is ‘constant’. Default is 0.0.
+    prefilter : bool, optional
+        currently not supported
+    output_shape : tuple of ints, optional
+        The shape of the array to be returned.
+    output_chunks : tuple of ints, optional
+        The shape of the chunks of the output Dask Array.
+
+    Returns
+    -------
+    rotate : Dask Array
+        A dask array representing the rotated input.
+
+    Notes
+    -----
+        Differences to `ndimage.affine_transformation`:
+        - currently, prefiltering is not supported
+          (affecting the output in case of interpolation `order > 1`)
+        - default order is 1
+        - modes 'reflect', 'mirror' and 'wrap' are not supported
+        - passing array to `output` is not currently supported.
+
+        Arguments equal to `ndimage.affine_rotate`,
+        except for `output_chunks`, `output_shape`.
+
+    Examples
+    --------
+    >>> from scipy import ndimage, misc
+    >>> import matplotlib.pyplot as plt
+    >>> import dask.array as da
+    >>> fig = plt.figure(figsize=(10, 3))
+    >>> ax1, ax2, ax3 = fig.subplots(1, 3)
+    >>> img = da.from_array(misc.ascent(),chunks=(64,64))
+    >>> img_45 = dask_image.ndinterp.rotate(img, 45, reshape=False)
+    >>> full_img_45 = dask_image.ndinterp.rotate(img, 45, reshape=True)
+    >>> ax1.imshow(img, cmap='gray')
+    >>> ax1.set_axis_off()
+    >>> ax2.imshow(img_45, cmap='gray')
+    >>> ax2.set_axis_off()
+    >>> ax3.imshow(full_img_45, cmap='gray')
+    >>> ax3.set_axis_off()
+    >>> fig.set_tight_layout(True)
+    >>> plt.show()
+    >>> print(img.shape)
+    (512, 512)
+    >>> print(img_45.shape)
+    (512, 512)
+    >>> print(full_img_45.shape)
+    (724, 724)
+
+    """
+    if not type(input_arr) == da.core.Array:
+        input_arr = da.from_array(input_arr)
+
+    if output_chunks is None:
+        output_chunks = input_arr.chunksize
+
+    if reshape & (output_shape is not None):
+        warnings.warn('Both reshaping desired and output_shape provided.'
+                      'Will use the explicit output_shape.', UserWarning)
+
+    if prefilter:
+        warnings.warn('Prefilter currently unsupported.', UserWarning)
+
+    if output_shape is None:
+        output_shape = np.asarray(input_arr.shape)
+
+    ndim = input_arr.ndim
+
+    if ndim < 2:
+        raise ValueError('input array should be at least 2D')
+
+    axes = list(axes)
+
+    if len(axes) != 2:
+        raise ValueError('axes should contain exactly two values')
+
+    if not all([isinstance(ax, Number) for ax in axes]):
+        raise TypeError('axes should contain only integer values')
+    elif not all([float(ax).is_integer() for ax in axes]):
+        raise TypeError('axes should contain only integer values')
+
+    if axes[0] < 0:
+        axes[0] += ndim
+    if axes[1] < 0:
+        axes[1] += ndim
+    if axes[0] < 0 or axes[1] < 0 or axes[0] >= ndim or axes[1] >= ndim:
+        raise ValueError('invalid rotation plane specified')
+
+    if output is not None:
+        try:
+            dtype = np.dtype(output)
+        except TypeError:     # pragma: no cover
+            raise TypeError(  # pragma: no cover
+                "Could not coerce the provided output to a dtype. "
+                "Passing array to output is not currently supported."
+            )
+    else:
+        dtype = input_arr.dtype
+
+    axes.sort()
+
+    c, s = special.cosdg(angle), special.sindg(angle)
+
+    rot_matrix = np.array([[c, s],
+                           [-s, c]])
+
+    img_shape = np.asarray(input_arr.shape)
+    in_plane_shape = img_shape[axes]
+
+    if reshape:
+        # Compute transformed input bounds
+        iy, ix = in_plane_shape
+        out_bounds = rot_matrix @ [[0, 0, iy, iy],
+                                   [0, ix, 0, ix]]
+        # Compute the shape of the transformed input plane
+        out_plane_shape = (out_bounds.ptp(axis=1) + 0.5).astype(int)
+    else:
+        out_plane_shape = img_shape[axes]
+
+    if (output_shape == img_shape).all():
+        output_shape[axes] = out_plane_shape
+    else:
+        out_plane_shape = np.asarray(output_shape)[axes]
+
+    out_center = rot_matrix @ ((out_plane_shape - 1) / 2)
+    in_center = (in_plane_shape - 1) / 2
+    offset = in_center - out_center
+    output_shape = tuple(output_shape)
+
+
+    if ndim <= 2:
+
+        output = affine_transform(input_arr, rot_matrix,
+                                  offset=offset, output_shape=output_shape,
+                                  order=order, mode=mode, cval=cval,
+                                  prefilter=prefilter,output_chunks=output_chunks)
+
+
+    elif ndim >= 3:
+        rotmat_nd = np.eye(ndim)
+        offset_nd = np.zeros(ndim)
+
+        for o_x,idx in enumerate(axes):
+
+            rotmat_nd[idx,axes[0]] = rot_matrix[o_x,0]
+            rotmat_nd[idx,axes[1]] = rot_matrix[o_x,1]
+
+            offset_nd[idx] = offset[o_x]
+
+
+        output = affine_transform(input_arr, rotmat_nd,
+                                  offset=offset_nd, output_shape=output_shape,
+                                  order=order, mode=mode, cval=cval,
+                                  prefilter=prefilter,output_chunks=output_chunks)
+
+    result = output.astype(dtype)
+
+    return result
+
+
 # magnitude of the maximum filter pole for each order
 # (obtained from scipy/ndimage/src/ni_splines.c)
 _maximum_pole = {

docs/coverage.rst

@@ -257,7 +257,7 @@ This table shows which SciPy ndimage functions are supported by dask-image.
      - ✓
    * - ``rotate``
      - ✓
-     -
+     - ✓
      -
    * - ``shift``
      - ✓
@@ -311,4 +311,3 @@ This table shows which SciPy ndimage functions are supported by dask-image.
      - ✓
      -
      -
-