cornucopia 0.0.0__py3-none-any.whl

cornucopia/fov.py

@@ -0,0 +1,707 @@
+"""
+This module contains transforms that operate on the field of view (FOV)
+of the image.
+"""
+__all__ = [
+    'FlipTransform',
+    'RandomFlipTransform',
+    'PermuteAxesTransform',
+    'RandomPermuteAxesTransform',
+    'PatchTransform',
+    'RandomPatchTransform',
+    'CropTransform',
+    'PadTransform',
+    'PowerTwoTransform',
+    'Rot90Transform',
+    'Rot180Transform',
+    'RandomRot90Transform',
+]
+# stdlib
+import math
+from math import inf
+from numbers import Number
+from random import shuffle
+
+# dependencies
+from torch import Tensor
+import typing_extensions as tx
+
+# internals
+from .base import Transform
+from .base import FinalTransform, NonFinalTransform, PerChannelTransform
+from .utils.py import ensure_list
+from .utils.padding import pad
+from .random import Uniform, RandKFrom, Sampler, RandInt, make_range
+from . import typing as cct
+
+
+class FlipTransform(FinalTransform):
+    """Flip one or more axes."""
+
+    def __init__(
+        self, axis: tx.Optional[cct.ScalarOrSequence[int]] = None, **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        axis : [list of] int
+            Axes to flip. By default, flip all spatial axes.
+
+        Other Parameters
+        ----------------
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(**kwargs)
+        self.axis = axis
+
+    def _xform(self, x: Tensor) -> Tensor:
+        axis = self.axis
+        if axis is None:
+            axis = list(range(1, x.ndim))
+        axis = ensure_list(axis)
+        return x.flip(axis)
+
+    def make_inverse(self) -> 'FlipTransform':
+        return self
+
+
+class RandomFlipTransform(NonFinalTransform):
+    """Randomly flip one or more axes."""
+
+    Final = Next = FlipTransform
+    """The transform type returned by `unroll`, `next` and `final`."""
+
+    def __init__(
+        self,
+        axes: tx.Union[Sampler, cct.ScalarOrSequence[int], None] = None,
+        *,
+        shared: bool = True,
+        **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        axes : Sampler | [list of] int
+            Axes that can be flipped (default: all spatial axes)
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        axes = kwargs.pop('axis', axes)
+        super().__init__(shared=shared, **kwargs)
+        self.axes = axes
+
+    def _unroll(self, x: Tensor, max_depth: int = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        if 'channels' not in self.shared and len(x) > 1:
+            return PerChannelTransform(
+                [self.unroll(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).unroll(x, max_depth-1)
+        axes = self.axes or range(1, x.ndim)
+        if not isinstance(axes, Sampler):
+            rand_axes = RandKFrom(ensure_list(axes))
+        rand_axes = rand_axes()
+        return FlipTransform(rand_axes).unroll(x, max_depth-1)
+
+
+class PermuteAxesTransform(FinalTransform):
+    """Permute axes"""
+
+    def __init__(
+        self, permutation: tx.Optional[tx.Sequence[int]] = None, **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        permutation : [list of] int
+            Axes permutation. By default, reverse axes.
+            Only applies to spatial axes, so axes are numbered [C, 0, 1, 2]
+
+        Other Parameters
+        ----------------
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(**kwargs)
+        self.permutation = permutation
+
+    def _xform(self, x: Tensor) -> Tensor:
+        permutation = self.permutation
+        if permutation is None:
+            permutation = list(reversed(range(x.dim()-1)))
+        permutation = [0] + [p+1 for p in permutation]
+        return x.permute(permutation)
+
+    def make_inverse(self) -> 'PermuteAxesTransform':
+        if self.permutation:
+            i = range(len(self.permutation))
+            iperm = [i[p] for p in self.permutation]
+            return PermuteAxesTransform(iperm, **self.get_prm())
+        else:
+            return self
+
+
+class RandomPermuteAxesTransform(NonFinalTransform):
+    """Randomly permute axes."""
+
+    Final = Next = PermuteAxesTransform
+    """The transform type returned by `unroll`, `next` and `final`."""
+
+    def __init__(
+        self,
+        axes: tx.Optional[tx.Sequence[int]] =  None,
+        *,
+        shared: bool =True,
+        **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        axes : [list of] int
+            Axes that can be permuted (default: all)
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(shared=shared, **kwargs)
+        self.axes = axes
+
+    def _unroll(self, x: Tensor, max_depth: float = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        if 'channels' not in self.shared and len(x) > 1:
+            return PerChannelTransform(
+                [self.unroll(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).unroll(x, max_depth-1)
+        axes = list(self.axes or range(x.ndim-1))
+        shuffle(axes)
+        return PermuteAxesTransform(
+            axes, **self.get_prm()
+        ).unroll(x, max_depth-1)
+
+
+class Rot90Transform(FinalTransform):
+    """Apply a 90 (or 180) rotation along one or several axes."""
+
+    def __init__(
+        self,
+        axis: cct.ScalarOrSequence[int] = 0,
+        negative: cct.ScalarOrSequence[bool] = False,
+        double: cct.ScalarOrSequence[bool] = False,
+        **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        axis : [list of] int
+            Rotation axis (indexing does not account for the channel axis)
+        negative : [list of] bool
+            Rotate by -90 deg instead of 90 deg
+        double : [list of] bool
+            Rotate be 180 instead of 90 (`negative` is then unused)
+
+        Other Parameters
+        ----------------
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(**kwargs)
+        self.axis = ensure_list(axis)
+        self.negative = ensure_list(negative, len(self.axis))
+        self.double = ensure_list(double, len(self.axis))
+
+    def _xform(self, x: Tensor) -> Tensor:
+        # this implementation is suboptimal. We should fuse all transpose
+        # and all flips into a single "transpose + flip" operation so that
+        # a single allocation happens. This will be fine for now.
+
+        ndim = x.ndim - 1
+        axis = [1 + (ndim + a if a < 0 else a) for a in self.axis]
+        for ax, neg, dbl in zip(axis, self.negative, self.double):
+            if dbl:
+                if ndim == 2:
+                    dims = [1, 2]
+                else:
+                    assert ndim == 3
+                    dims = [d for d in (1, 2, 3) if d != ax]
+                x = x.flip(dims)
+            else:
+                if ndim == 2:
+                    dims = [1, 2]
+                else:
+                    assert ndim == 3
+                    dims = [d for d in (1, 2, 3) if d != ax]
+                x = x.transpose(*dims).flip(dims[1] if neg else dims[0])
+        return x
+
+
+class Rot180Transform(Rot90Transform):
+    """Apply a 180 deg rotation along one or several axes"""
+
+    def __init__(self, axis: cct.ScalarOrSequence[int] = 0, **kwargs) -> None:
+        """
+        Parameters
+        ----------
+        axis : [list of] int
+            Rotation axis (indexing does not account for the channel axis)
+
+        Other Parameters
+        ----------------
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(axis, double=True, **kwargs)
+
+
+class RandomRot90Transform(NonFinalTransform):
+    """Random set of 90 transforms."""
+
+    Final = Next = Rot90Transform
+    """The transform type returned by `unroll`, `next` and `final`."""
+
+    def __init__(
+        self,
+        axes: tx.Optional[cct.ScalarOrSequence[int]] = None,
+        max_rot: cct.SamplerOrBound[int] = 2,
+        negative: bool = True,
+        *,
+        shared: cct.SharedT = True,
+        **kwargs
+     ) -> None:
+        """
+        Parameters
+        ----------
+        axes : [list of] int
+            Axes along which rotations can happen.
+            If `None`, all axes.
+        max_rot : Sampler | int
+            Maximum number of consecutive rotations.
+        negative : bool
+            Whether to authorize negative rotations.
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(shared=shared, **kwargs)
+        self.axes = axes
+        self.max_rot = RandInt.make(make_range(1, max_rot))
+        self.negative = negative
+
+    def _unroll(self, x, max_depth=float('inf')):
+        if max_depth == 0:
+            return self
+        if 'channels' not in self.shared and len(x) > 1:
+            return PerChannelTransform(
+                [self.unroll(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).unroll(x, max_depth-1)
+        ndim = x.ndim - 1
+        max_rot = self.max_rot
+        if isinstance(max_rot, Sampler):
+            max_rot = max_rot()
+        axes = self.axes
+        if axes is None:
+            axes = list(range(ndim))
+        if isinstance(axes, (int, list, tuple)):
+            axes = ensure_list(axes, max_rot, crop=False)
+        if not isinstance(axes, Sampler):
+            axes = RandKFrom(axes, max_rot, replacement=True)
+
+        axes = ensure_list(axes(), max_rot)
+        negative = RandKFrom([False, True], max_rot, replacement=True)() \
+            if self.negative else [False] * max_rot
+        return Rot90Transform(
+            axes, negative, **self.get_prm()
+        ).unroll(max_depth-1)
+
+
+class CropPadTransform(FinalTransform):
+    """Crop and/or pad a tensor."""
+
+    def __init__(
+        self,
+        crop: tx.Sequence[slice] = (),
+        pad: tx.Sequence[int] = (),
+        bound: cct.ItemOrSequence[str] = 'zero',
+        value: Number = 0,
+        **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        crop : list[slice]
+            Slicing operator per dimension.
+        pad : list[int]
+            Left and right padding per dimensions
+        bound : [list of] str
+            Boundary condition for padding
+        value : number
+            Padding value in case `bound='constant`
+
+        Other Parameters
+        ----------------
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        super().__init__(**kwargs)
+        self.crop = crop
+        self.pad = pad
+        self.bound = bound
+        self.value = value
+
+    def _xform(self, x: Tensor) -> Tensor:
+        crop = tuple([Ellipsis, *self.crop])
+        x = x[crop]
+        x = pad(x, self.pad, mode=self.bound, value=self.value)
+        return x
+
+    def make_inverse(self) -> 'CropPadTransform':
+        ipad = [slice(left, (-right) or None) for left, right in self.pad]
+        icrop = [[s.start or 0, -s.stop if s.stop else 0] for s in self.crop]
+        return CropPadTransform(
+            ipad, icrop, bound=self.bound, value=self.value, **self.get_prm()
+        )
+
+
+class PatchTransform(NonFinalTransform):
+    """Extract a patch from the volume"""
+
+    Final = Next = CropPadTransform
+    """The transform type returned by `unroll`, `next` and `final`."""
+
+    def __init__(
+        self,
+        shape: cct.ScalarOrSequence[int] = 64,
+        center: cct.ScalarOrSequence[float] = 0,
+        bound: cct.ItemOrSequence[str] = 'zero',
+        *,
+        shared: cct.SharedT = True,
+        **kwargs
+    ) -> None:
+        """
+        Parameters
+        ----------
+        shape : [list of] int
+            Patch shape
+        center : [list of] float
+            Patch center, in relative coordinates -1..1
+        bound : [list of]str
+            Boundary condition in case padding is needed
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+
+            !!! changedin "![v0.5](https://img.shields.io/badge/v0.5-yellow) \
+                Default for `shared` changed from `"channels"` to `True`"
+
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        kwargs.setdefault('shared', shared)
+        super().__init__(**kwargs)
+        self.shape = shape
+        self.center = center
+        self.bound = bound
+
+    def _unroll(self, x: Tensor, max_depth: int = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        ndim = x.dim() - 1
+        shape = ensure_list(self.shape, ndim)
+        center = ensure_list(self.center, ndim)
+        center = [(c + 1) / 2 * (s - 1) for c, s in zip(center, x.shape[1:])]
+        crop = []
+        padding = []
+        for ss, cc, sv in zip(shape, center, x.shape[1:]):
+            first = int(math.floor(cc - ss/2))
+            pad_first = max(0, -first)
+            last = first + ss
+            pad_last = max(0, last - sv)
+            first = max(0, first)
+            last = min(sv, last)
+            last = (last - sv) or None  # ensure negative for CropPad
+            crop.append(slice(first, last))
+            padding.extend([pad_first, pad_last])
+        return CropPadTransform(
+            crop, padding, bound=self.bound, **self.get_prm()
+        ).unroll(x, max_depth-1)
+
+
+class RandomPatchTransform(NonFinalTransform):
+    """Extract a (randomly located) patch from the volume.
+
+    This transform ensures that the patch is fully contained within the
+    original field of view (unless the patch size is larger than the
+    input shape).
+    """
+
+    Next = PatchTransform
+    """The transform type returned by `next`."""
+
+    Final = CropPadTransform
+    """The transform type returned by `final`."""
+
+    def __init__(
+        self,
+        shape: cct.ScalarOrSequence[int],
+        bound: cct.ItemOrSequence[str] = 'zero',
+        *,
+        shared: cct.SharedT = True,
+        **kwargs
+    ) -> None:
+        """
+
+        Parameters
+        ----------
+        shape : [list of] int
+            Patch shape
+        bound : [list of] str
+            Boundary condition in case padding is needed
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+
+            !!! changedin "![v0.5](https://img.shields.io/badge/v0.5-yellow) \
+                        Default for `shared` changed from `"channels"` to `True`"
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        shape = kwargs.pop('patch_size', shape)  # support legacy name
+        kwargs.setdefault('shared', shared)
+        super().__init__(**kwargs)
+        self.shape = shape
+        self.bound = bound
+
+    def _unroll(self, x, max_depth=float('inf')):
+        if max_depth == 0:
+            return self
+        shape = x.shape[1:]
+        patch_size = ensure_list(self.shape, len(shape))
+        min_center = [max(p/s - 1, -1) for p, s in zip(patch_size, shape)]
+        max_center = [min(1 - p/s, 1) for p, s in zip(patch_size, shape)]
+        center = [Uniform(mn, mx)() for mn, mx in zip(min_center, max_center)]
+        return PatchTransform(
+            patch_size, center, self.bound, **self.get_prm()
+        ).unroll(x, max_depth-1)
+
+
+class CropTransform(NonFinalTransform):
+    """Crop a tensor by some amount"""
+
+    def __init__(
+        self,
+        cropping: cct.ScalarOrSequence[tx.Union[int, float]],
+        unit: str = 'vox',
+        side: str = 'both',
+        **kwargs
+    ) -> None:
+        """
+
+        Parameters
+        ----------
+        cropping : [list of] int or float
+            Amount of cropping. If `side` is `None`, pre and post cropping
+            must be provided in turn.
+        unit : {'vox', 'pct'}
+            Padding unit
+        side : {'pre', 'post', 'both', None}
+            Side to crop
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+
+            !!! changedin "![v0.5](https://img.shields.io/badge/v0.5-yellow) \
+                        Default for `shared` changed from `"channels"` to `True`"
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        kwargs.setdefault('shared', True)
+        super().__init__(**kwargs)
+        self.cropping = cropping
+        self.unit = unit
+        self.side = side
+
+    def _unroll(self, x: Tensor, max_depth: int = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        ndim = x.dim() - 1
+        cropping = self.cropping
+        if self.side is not None:
+            cropping = ensure_list(cropping, ndim)
+            if self.unit[0] == 'p':
+                cropping = [int(math.ceil(c * s))
+                            for c, s in zip(cropping, x.shape[1:])]
+            cropping = [slice(c, -c if c else None) for c in cropping]
+        else:
+            cropping = ensure_list(cropping)
+            cropping = [0] * (2*ndim - len(cropping))
+            if self.unit[0] == 'p':
+                shape2 = [s for s in x.shape[1:] for _ in range(2)]
+                cropping = [int(math.ceil(c * s))
+                            for c, s in zip(cropping, shape2)]
+            cropping = [slice(c0, -c1 if c1 else None)
+                        for c0, c1 in zip(cropping[::2], cropping[1::2])]
+        return CropPadTransform(
+            cropping, [0]*(2*ndim), **self.get_prm()
+        ).unroll(x, max_depth-1)
+
+
+class PadTransform(NonFinalTransform):
+    """Pad a tensor by some amount"""
+
+    Final = Next = CropPadTransform
+    """The transform type returned by `unroll`, `next` and `final`."""
+
+    def __init__(
+        self,
+        padding: cct.ScalarOrSequence[tx.Union[int, float]],
+        unit: str = 'vox',
+        side: str = 'both',
+        bound: str = 'zero',
+        value: Number = 0,
+        **kwargs
+    ) -> None:
+        """
+
+        Parameters
+        ----------
+        padding : [list of] int or float
+            Amount of padding. If `side` is `None`, pre and post padding
+            must be provided in turn.
+        unit : {'vox', 'pct'}
+            Padding unit
+        side : {'pre', 'post', 'both', None}
+            Side to pad
+        bound : str
+            Boundary condition
+        value : float
+            Value for case `bound='constant'`
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+
+            !!! changedin "![v0.5](https://img.shields.io/badge/v0.5-yellow) \
+                        Default for `shared` changed from `"channels"` to `True`"
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        kwargs.setdefault('shared', True)
+        super().__init__(**kwargs)
+        self.padding = padding
+        self.unit = unit
+        self.side = side
+        self.bound = bound
+        self.value = value
+
+    def _unroll(self, x: Tensor, max_depth: int = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        ndim = x.dim() - 1
+        padding = self.padding
+        if self.side is not None:
+            padding = ensure_list(padding, ndim)
+            if self.unit[0] == 'p':
+                padding = [int(math.ceil(p * s))
+                           for p, s in zip(padding, x.shape[1:])]
+
+        else:
+            padding = ensure_list(padding)
+            padding = [0] * (2 * ndim - len(padding)) + padding
+            if self.unit[0] == 'p':
+                shape2 = [s for s in x.shape[1:] for _ in range(2)]
+                padding = [int(math.ceil(p * s))
+                           for p, s in zip(padding, shape2)]
+
+        if self.side == 'pre':
+            padding = [p for pz in zip(padding, [0]*ndim) for p in pz]
+        elif self.side == 'post':
+            padding = [p for zp in zip([0]*ndim, padding) for p in zp]
+        elif self.side == 'both':
+            padding = [p for pp in zip(padding, padding) for p in pp]
+
+        return CropPadTransform(
+            [slice(None)]*ndim, padding, bound=self.bound, value=self.value,
+            **self.get_prm()
+        ).unroll(x, max_depth-1)
+
+
+class PowerTwoTransform(NonFinalTransform):
+    """Pad the volume such that the tensor shape can be divided by 2**x"""
+
+    Next = PatchTransform
+    """The transform type returned by `next`."""
+
+    Final = CropPadTransform
+    """The transform type returned by `final`."""
+
+    def __init__(
+        self,
+        exponent: cct.ScalarOrSequence[int] = 1,
+        bound: cct.ItemOrSequence[str] = 'zero',
+        **kwargs
+    ) -> None:
+        """
+
+        Parameters
+        ----------
+        exponent : [list of] int
+            Ensure that the shape can be divided by 2 ** exponent
+        bound : [list of] str
+            Boundary condition for padding
+
+        Other Parameters
+        ----------------
+        shared
+            See [`NonFinalTransform`][cornucopia.base.NonFinalTransform]
+            for details.
+
+            !!! changedin "![v0.5](https://img.shields.io/badge/v0.5-yellow) \
+                        Default for `shared` changed from `"channels"` to `True`"
+        returns, append, prefix, include, exclude, consume
+            See [`Transform`][cornucopia.base.Transform] for details.
+        """
+        kwargs.setdefault('shared', True)
+        super().__init__(**kwargs)
+        self.exponent = exponent
+        self.bound = bound
+
+    def _unroll(self, x: Tensor, max_depth: int = inf) -> Transform:
+        if max_depth == 0:
+            return self
+        shape = x.shape[1:]
+        exponent = ensure_list(self.exponent, len(shape))
+        bigshape = [max(2 ** e, s) for e, s in zip(exponent, shape)]
+        return PatchTransform(
+            bigshape, bound=self.bound, **self.get_prm()
+        ).unroll(x, max_depth-1)