cornucopia 0.1.0__tar.gz → 0.3.0__tar.gz

{cornucopia-0.1.0 → cornucopia-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cornucopia
-Version: 0.1.0
+Version: 0.3.0
 Summary: An abundance of augmentation layers
 Home-page: UNKNOWN
 Author: Yael Balbastre
@@ -10,7 +10,7 @@ Project-URL: Source Code, https://github.com/balbasty/cornucopia
 Description: <picture align="center">
           <source media="(prefers-color-scheme: dark)" srcset="docs/icons/cornucopia_lightorange.svg">
           <source media="(prefers-color-scheme: light)"  srcset="docs/icons/cornucopia_orange.svg">
-          <img alt="Cornucopia logo" src="docs/icons/cornucopia_orange.svg">
+          <img alt="Cornucopia logo" src="https://github.com/balbasty/cornucopia/raw/main/docs/icons/cornucopia_orange.svg">
         </picture>
         
         The `cornucopia` package provides a generic framework for preprocessing,
@@ -27,9 +27,6 @@ Description: <picture align="center">
         theoretically be used within any dataloader pipeline,
         independent of the downstream learning framework (pytorch, tensorflow, jax, ...).
         
-        ## Installation
-        
-        
         ## Installation
         
         ### Dependencies
@@ -43,15 +40,21 @@ Description: <picture align="center">
         ### Conda
         
         ```sh
-        conda install cornucopia -c balbasty -c pytorch 
+        conda install cornucopia -c balbasty -c pytorch -c conda-forge
         ```
         
-        ### Pip
+        ### Pip (release)
         
         ```sh
         pip install cornucopia
         ```
         
+        ### Pip (dev)
+        
+        ```sh
+        pip install cornucopia@git+https://github.com/balbasty/cornucopia
+        ```
+        
         ## Documentation
         
         Read the [documentation](https://cornucopia.readthedocs.io) and in particular:

{cornucopia-0.1.0 → cornucopia-0.3.0}/README.md

@@ -1,7 +1,7 @@
 <picture align="center">
   <source media="(prefers-color-scheme: dark)" srcset="docs/icons/cornucopia_lightorange.svg">
   <source media="(prefers-color-scheme: light)"  srcset="docs/icons/cornucopia_orange.svg">
-  <img alt="Cornucopia logo" src="docs/icons/cornucopia_orange.svg">
+  <img alt="Cornucopia logo" src="https://github.com/balbasty/cornucopia/raw/main/docs/icons/cornucopia_orange.svg">
 </picture>
 
 The `cornucopia` package provides a generic framework for preprocessing,
@@ -18,9 +18,6 @@ Since gradients are not expected to backpropagate through its layers, it can
 theoretically be used within any dataloader pipeline,
 independent of the downstream learning framework (pytorch, tensorflow, jax, ...).
 
-## Installation
-
-
 ## Installation
 
 ### Dependencies
@@ -34,15 +31,21 @@ independent of the downstream learning framework (pytorch, tensorflow, jax, ...)
 ### Conda
 
 ```sh
-conda install cornucopia -c balbasty -c pytorch 
+conda install cornucopia -c balbasty -c pytorch -c conda-forge
 ```
 
-### Pip
+### Pip (release)
 
 ```sh
 pip install cornucopia
 ```
 
+### Pip (dev)
+
+```sh
+pip install cornucopia@git+https://github.com/balbasty/cornucopia
+```
+
 ## Documentation
 
 Read the [documentation](https://cornucopia.readthedocs.io) and in particular:

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/__init__.py

@@ -1,4 +1,5 @@
-"""Flexible transforms for pre-processing and augmentation
+"""
+Flexible transforms for pre-processing and augmentation
 
 Example on how to use this machinery to generate within-subject
 image pairs with a random affine deformation between them::
@@ -30,19 +31,25 @@ image pairs with a random affine deformation between them::
 
 """
 
-# TODO:
-#   [x] Make it a standalone package?
-#   [x] Move samplers in their own file
-#   [x] Add IO transforms (that transform filenames in tensors)
-#   [ ] Better deal with separable/shared transforms
-#       [ ] Add a SharedTransform class (like Randomized) that does the heavy
-#           lifting
-#       [ ] By default (non shared), let Transforms handle multi-channel
-#           data (currently we loop across channels in the base class)
-
 from . import random                        # noqa: F401
 from . import ctx                           # noqa: F401
-from .ctx import batch                      # noqa: F401
+from . import base                          # noqa: F401
+from . import special                       # noqa: F401
+from . import contrast                      # noqa: F401
+from . import geometric                     # noqa: F401
+from . import intensity                     # noqa: F401
+from . import io                            # noqa: F401
+from . import fov                           # noqa: F401
+from . import kspace                        # noqa: F401
+from . import labels                        # noqa: F401
+from . import noise                         # noqa: F401
+from . import psf                           # noqa: F401
+from . import qmri                          # noqa: F401
+from . import synth                         # noqa: F401
+from . import utils                         # noqa: F401
+
+from .random import *                       # noqa: F401,F403
+from .ctx import *                          # noqa: F401,F403
 from .base import *                         # noqa: F401,F403
 from .special import *                      # noqa: F401,F403
 from .contrast import *                     # noqa: F401,F403

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-08-28T17:04:06-0400",
+ "date": "2024-04-19T14:23:50+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "57fb52ad37a27a28c5894754411fb6bb73e60b5a",
- "version": "0.1.0"
+ "full-revisionid": "37de94f181b9a97eebd21460f4df63ae4a0750f8",
+ "version": "0.3.0"
 }
 '''  # END VERSION_JSON
 

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/base.py

@@ -148,7 +148,7 @@ class Transform(nn.Module):
                 return x
 
         # now we're working with a single tensor (or str)
-        y = self.apply(x)
+        y = self.xform(x)
         if not isinstance(y, Returned):
             if not isinstance(y, type(self.returns)):
                 y = dict(input=x, output=y)
@@ -240,7 +240,7 @@ class FinalTransform(Transform):
     def is_final(self):
         return True
 
-    def apply(self, x):
+    def xform(self, x):
         """Apply the transform to a tensor
 
         Parameters
@@ -283,7 +283,7 @@ class IdentityTransform(FinalTransform):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
-    def apply(self, x):
+    def xform(self, x):
         return x
 
     def make_inverse(self):
@@ -305,12 +305,12 @@ class SharedMixin:
             shared = ''
         return shared
 
-    def apply(self, x):
+    def xform(self, x):
         if 'channels' in self.shared:
             xform = self.make_final(x[:1], max_depth=1)
         else:
             xform = self.make_final(x, max_depth=1)
-        return xform.apply(x)
+        return xform.xform(x)
 
     def forward(self, *a, **k):
         return self._shared_forward(*a, **k)
@@ -359,7 +359,7 @@ class NonFinalTransform(SharedMixin, Transform):
         super().__init__(**kwargs)
         self.shared = self._prepare_shared(shared)
 
-    def make_final(self, x, max_depth=float('inf'), *args, **kwargs):
+    def make_final(self, x, max_depth=float('inf')):
         if self.is_final or max_depth == 0:
             return self
         return NotImplemented
@@ -447,7 +447,7 @@ class SequentialTransform(SpecialMixin, SharedMixin, Transform):
     def make_final(self, x, max_depth=float('inf')):
         if max_depth == 0:
             return self
-        x = VirtualTensor.from_any(x, compute_stats=True)
+        # x = VirtualTensor.from_any(x, compute_stats=True)
         trf = []
         for t in self:
             t = t.make_final(x, max_depth=max_depth-1)
@@ -525,7 +525,7 @@ class PerChannelTransform(SpecialMixin, Transform):
         trf = PerChannelTransform(trf, **prm)
         return trf
 
-    def apply(self, x):
+    def xform(self, x):
         results = []
         for i, t in enumerate(self.transforms):
             with ReturningTransform(t, self.returns), \
@@ -1068,33 +1068,67 @@ class RandomizedTransform(NonFinalTransform):
     """
     Transform generated by randomizing some parameters of another transform.
 
+    !!! note "`ctx.randomize` is an alias for `RandomizedTransform`"
+
     !!! example "Gaussian noise with randomized variance"
         Object call
         ```python
         import cornucopia as cc
-        hypernoise = RandomizedTransform(cc.GaussianNoise, [cc.Uniform(0, 10)])
+        hypernoise = cc.RandomizedTransform(cc.GaussianNoise, [cc.Uniform()])
         img = hypernoise(img)
         ```
-        Functional call
+
+        Delayed call
         ```python
         import cornucopia as cc
-        hypernoise = cc.randomize(cc.GaussianNoise)(cc.Uniform(0, 10))
+        MyRandomNoise = cc.randomize(cc.GaussianNoise)
+        hypernoise = MyRandomNoise(cc.Uniform())
         img = hypernoise(img)
         ```
 
     """
 
-    def __init__(self, transform, sample, ksample=None,
+    class Delayed:
+        # Temproary parameter holder for delayed calls
+        def __init__(self, transform, **kwargs):
+            self.transform = transform
+            self.kwargs = kwargs
+
+        def __call__(self, *args, **kwargs):
+            return RandomizedTransform(
+                self.transform, args, kwargs, **self.kwargs)
+
+    def __new__(cls, *args, **kwargs):
+        if cls is RandomizedTransform:
+            return cls._base_new(*args, **kwargs)
+        return super().__new__(cls)
+
+    @classmethod
+    def _base_new(cls, transform, sample=tuple(), ksample=dict(),
+                  *, shared=False, **kwargs):
+        assert cls is RandomizedTransform
+        if not sample and not ksample:
+            # If no arguments are passed, it means that the user calls
+            # this in "delayed/functional" mode. In that case, we return
+            # a callable object that returns the constructed instance
+            # using the call-time arguments.
+            return cls.Delayed(transform, shared=shared, **kwargs)
+        # Otherwise, we're in object mode and we instantiate the
+        # randomized object.
+        return super().__new__(cls)
+
+    def __init__(self, transform, sample=tuple(), ksample=dict(),
                  *, shared=False, **kwargs):
         """
-
         Parameters
         ----------
         transform : callable(...) -> Transform
             A Transform subclass or a function that constructs a Transform.
         sample : [list or dict of] callable
             A collection of functions that generate parameter values provided
-            to `transform`.
+            to `transform`. Can be args-like or kwargs-like arguments.
+        ksample : dict[callable]
+            Must be kwargs-like arguments.
 
         Other Parameters
         ----------------
@@ -1130,9 +1164,7 @@ class RandomizedTransform(NonFinalTransform):
 
     def __repr__(self):
         if type(self) is RandomizedTransform:
-            try:
-                if issubclass(self.subtransform, Transform):
-                    return f'Randomized{self.subtransform.__name__}()'
-            except TypeError:
-                pass
+            xform = self.subtransform
+            if isinstance(xform, type) and issubclass(xform, Transform):
+                return f'Randomized{xform.__name__}()'
         return super().__repr__()

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/baseutils.py

@@ -24,7 +24,7 @@ def get_first_element(x, include=None, exclude=None, types=None):
                 if ok:
                     return v, True
             return None, False
-        if torch.is_tensor(x) or isinstance(x, types):
+        if torch.is_tensor(x) or (types and isinstance(x, types)):
             return x, True
         return x, False
 
@@ -213,9 +213,9 @@ class VirtualTensor:
     @classmethod
     def from_tensor(cls, x, compute_stats=False):
         if compute_stats:
-            vmin = x.min(dim=list(range(1, x.ndim)))
-            vmax = x.max(dim=list(range(1, x.ndim)))
-            vmean = x.mean(dim=list(range(1, x.ndim)))
+            vmin = x.reshape([len(x), -1]).min(dim=-1).values
+            vmax = x.reshape([len(x), -1]).max(dim=-1).values
+            vmean = x.float().mean(dim=list(range(1, x.ndim)))
         else:
             vmin = vmax = vmean = None
         return VirtualTensor(x.shape, dtype=x.dtype, device=x.device,

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/contrast.py

@@ -45,7 +45,7 @@ class ContrastMixtureTransform(NonFinalTransform):
             self.mu = mu
             self.sigma = sigma
 
-        def apply(self, x):
+        def xform(self, x):
             z = self.z.to(x)
             mu0 = self.mu0.to(x)
             sigma0 = self.sigma0.to(x)
@@ -114,7 +114,8 @@ class ContrastMixtureTransform(NonFinalTransform):
         mu = torch.rand_like(
             old_mu).mul_(old_mu_max - old_mu_min).add_(old_mu_min)
         sigma = torch.rand_like(
-            old_sigma_diag).mul_(old_sigma_max - old_sigma_min).add_(old_sigma_min)
+            old_sigma_diag
+        ).mul_(old_sigma_max - old_sigma_min).add_(old_sigma_min)
         corr = torch.rand([len(old_mu), nc*(nc-1)//2], **backend).mul_(0.5)
 
         fullsigma = torch.eye(nc, **backend).expand([nk, nc, nc]).clone()
@@ -145,7 +146,7 @@ class ContrastLookupTransform(NonFinalTransform):
             self.edges = edges
             self.mu = mu
 
-        def apply(self, x):
+        def xform(self, x):
             edges, mu = self.edges.to(x), self.mu.to(x)
             mu0 = (edges[:-1] + edges[1:]) / 2
             nk = len(mu)

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/fov.py

@@ -8,14 +8,16 @@ __all__ = [
     'CropTransform',
     'PadTransform',
     'PowerTwoTransform',
+    'Rot90Transform',
+    'Rot180Transform',
+    'RandomRot90Transform',
 ]
-
 import math
 from random import shuffle
-from .base import FinalTransform, NonFinalTransform
+from .base import FinalTransform, NonFinalTransform, PerChannelTransform
 from .utils.py import ensure_list
 from .utils.padding import pad
-from .random import Uniform, RandKFrom, Sampler
+from .random import Uniform, RandKFrom, Sampler, RandInt, make_range
 
 
 class FlipTransform(FinalTransform):
@@ -23,7 +25,6 @@ class FlipTransform(FinalTransform):
 
     def __init__(self, axis=None, **kwargs):
         """
-
         Parameters
         ----------
         axis : [list of] int
@@ -32,7 +33,7 @@ class FlipTransform(FinalTransform):
         super().__init__(**kwargs)
         self.axis = axis
 
-    def apply(self, x):
+    def xform(self, x):
         axis = self.axis
         if axis is None:
             axis = list(range(1, x.ndim))
@@ -46,24 +47,30 @@ class FlipTransform(FinalTransform):
 class RandomFlipTransform(NonFinalTransform):
     """Randomly flip one or more axes"""
 
-    def __init__(self, axes=None, **kwargs):
+    def __init__(self, axes=None, *, shared=True, **kwargs):
         """
-
         Parameters
         ----------
         axes : Sampler or [list of] int
             Axes that can be flipped (default: all)
+
+        Other Parameters
+        ----------------
         shared : {'channels', 'tensors', 'channels+tensors', ''}
             Apply the same flip to all channels and/or tensors
         """
         axes = kwargs.pop('axis', axes)
-        kwargs.setdefault('shared', True)
-        super().__init__(**kwargs)
+        super().__init__(shared=shared, **kwargs)
         self.axes = axes
 
     def make_final(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.make_final(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).make_final(x, max_depth-1)
         axes = self.axes or range(1, x.ndim)
         if not isinstance(axes, Sampler):
             rand_axes = RandKFrom(ensure_list(axes))
@@ -76,7 +83,6 @@ class PermuteAxesTransform(FinalTransform):
 
     def __init__(self, permutation=None, **kwargs):
         """
-
         Parameters
         ----------
         permutation : [list of] int
@@ -86,7 +92,7 @@ class PermuteAxesTransform(FinalTransform):
         super().__init__(**kwargs)
         self.permutation = permutation
 
-    def apply(self, x):
+    def xform(self, x):
         permutation = self.permutation
         if permutation is None:
             permutation = list(reversed(range(x.dim()-1)))
@@ -105,23 +111,29 @@ class PermuteAxesTransform(FinalTransform):
 class RandomPermuteAxesTransform(NonFinalTransform):
     """Randomly permute axes"""
 
-    def __init__(self, axes=None, **kwargs):
+    def __init__(self, axes=None, *, shared=True, **kwargs):
         """
-
         Parameters
         ----------
         axes : [list of] int
             Axes that can be permuted (default: all)
+
+        Other Parameters
+        ----------------
         shared : {'channels', 'tensors', 'channels+tensors', ''}
             Apply the same permutation to all channels and/or tensors
         """
-        kwargs.setdefault('shared', True)
-        super().__init__(**kwargs)
+        super().__init__(shared=shared, **kwargs)
         self.axes = axes
 
     def make_final(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.make_final(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).make_final(x, max_depth-1)
         axes = list(self.axes or range(x.ndim-1))
         shuffle(axes)
         return PermuteAxesTransform(
@@ -129,6 +141,119 @@ class RandomPermuteAxesTransform(NonFinalTransform):
         ).make_final(x, max_depth-1)
 
 
+class Rot90Transform(FinalTransform):
+    """
+    Apply a 90 (or 180) rotation along one or several axes
+    """
+
+    def __init__(self, axis=0, negative=False, double=False, **kwargs):
+        """
+        Parameters
+        ----------
+        axis : int or list[int]
+            Rotation axis (indexing does not account for the channel axis)
+        negative : bool or list[bool]
+            Rotate by -90 deg instead of 90 deg
+        double : bool or list[bool]
+            Rotate be 180 instead of 90 (`negative` is then unused)
+        """
+        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):
+        # 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=0, **kwargs):
+        """
+        Parameters
+        ----------
+        axis : int or list[int]
+            Rotation axis (indexing does not account for the channel axis)
+        """
+        super().__init__(axis, double=True, **kwargs)
+
+
+class RandomRot90Transform(NonFinalTransform):
+    """Random set of 90 transforms"""
+
+    def __init__(self, axes=None, max_rot=2, negative=True,
+                 *, shared=True, **kwargs):
+        """
+        Parameters
+        ----------
+        axes : int or list[int]
+            Axes along which rotations can happen.
+            If `None`, all axes.
+        max_rot : int or Sampler
+            Maximum number of consecutive rotations.
+        negative : bool
+            Whether to authorize negative rotations.
+
+        Other Parameters
+        ----------------
+        shared : {'channels', 'tensors', 'channels+tensors', ''}
+            Apply the same permutation to all channels and/or tensors
+        """
+        super().__init__(shared=shared, **kwargs)
+        self.axes = axes
+        self.max_rot = RandInt.make(make_range(1, max_rot))
+        self.negative = negative
+
+    def make_final(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.make_final(x[i:i+1], max_depth) for i in range(len(x))],
+                **self.get_prm()
+            ).make_final(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()
+        ).make_final(max_depth-1)
+
+
 class CropPadTransform(FinalTransform):
     """Crop and/or pad a tensor"""
 
@@ -151,7 +276,7 @@ class CropPadTransform(FinalTransform):
         self.bound = bound
         self.value = value
 
-    def apply(self, x):
+    def xform(self, x):
         crop = tuple([Ellipsis, *self.crop])
         x = x[crop]
         x = pad(x, self.pad, mode=self.bound, value=self.value)

{cornucopia-0.1.0 → cornucopia-0.3.0}/cornucopia/geometric.py

@@ -218,7 +218,7 @@ class ElasticTransform(NonFinalTransform):
                 ).movedim(-1, 1)
             return flow
 
-        def apply(self, x):
+        def xform(self, x):
             """Deform the input tensor
 
             Parameters
@@ -503,7 +503,7 @@ class AffineTransform(NonFinalTransform):
         def make_flow(self, matrix, shape):
             return warps.affine_flow(matrix, shape).movedim(-1, 0)
 
-        def apply(self, x):
+        def xform(self, x):
             flow = cast_like(self.flow, x)
             matrix = cast_like(self.matrix, x)
             required = return_requires(self.returns)
@@ -766,7 +766,7 @@ class AffineElasticTransform(NonFinalTransform):
             self.affine = affine
             self.bound = bound
 
-        def apply(self, x):
+        def xform(self, x):
             flow = cast_like(self.flow, x)
             controls = cast_like(self.controls, x)
             affine = cast_like(self.affine, x)
@@ -954,7 +954,7 @@ class MakeAffinePair(NonFinalTransform):
             self.left = left
             self.right = right
 
-        def apply(self, x):
+        def xform(self, x):
             x1 = self.left(x)
             x2 = self.right(x)
             mat1, mat2 = self.left.matrix, self.right.matrix
@@ -1119,7 +1119,6 @@ class SlicewiseAffineTransform(NonFinalTransform):
         F = torch.eye(ndim+1, **backend)
         F[:ndim, -1] = -offsets
         Z = E.clone()
-        print(zooms.shape, Z.shape)
         Z.diagonal(0, -1, -2)[:, :-1].copy_(1 + zooms)
         T = E.clone()
         T[:, :ndim, -1] = translations
@@ -1204,7 +1203,7 @@ class SlicewiseAffineTransform(NonFinalTransform):
             self.subsample = subsample
             self.bound = bound
 
-        def apply(self, x):
+        def xform(self, x):
             flow = cast_like(self.flow, x)
             matrix = cast_like(self.matrix, x)
 
@@ -1362,7 +1361,7 @@ class RandomSlicewiseAffineTransform(NonFinalTransform):
         # get slice direction
         slice = self.slice
         if slice is None:
-            slice = RandInt(0, ndim)
+            slice = RandInt(0, ndim - 1)
         if isinstance(slice, Sampler):
             slice = slice()