cornucopia 0.2.0__tar.gz → 0.4.0__tar.gz

cornucopia-0.4.0/PKG-INFO

@@ -0,0 +1,93 @@
+Metadata-Version: 2.4
+Name: cornucopia
+Version: 0.4.0
+Summary: An abundance of augmentation layers
+Author: Yael Balbastre
+Author-email: yael.balbastre@gmail.com
+License: MIT
+Project-URL: Source Code, https://github.com/balbasty/cornucopia
+Platform: OS Independent
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.8
+Requires-Dist: numpy
+Requires-Dist: nibabel
+Requires-Dist: torch-interpol>=0.2.4
+Requires-Dist: torch-distmap
+Dynamic: license-file
+
+<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="https://github.com/balbasty/cornucopia/raw/main/docs/icons/cornucopia_orange.svg">
+</picture>
+
+The `cornucopia` package provides a generic framework for preprocessing,
+augmentation, and domain randomization; along with an abundance of specific layers,
+mostly targeted at (medical) imaging. `cornucopia` is written using a PyTorch
+backend, and therefore runs **on the CPU or GPU**.
+
+Cornucopia is *intended* to be used on the GPU for on-line augmentation.
+A quick [benchmark](docs/examples/benchmark.ipynb) of affine and elastic augmentation
+shows that while cornucopia is slower than [TorchIO](https://github.com/fepegar/torchio)
+on the CPU (~ 3s vs 1s), it is greatly accelerated on the GPU (~ 50ms).
+
+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
+
+### Dependencies
+
+- `pytorch >= 1.8`
+- `numpy`
+- `nibabel`
+- `torch-interpol`
+- `torch-distmap`
+
+### Conda
+
+```sh
+conda install cornucopia -c balbasty -c pytorch -c conda-forge
+```
+
+### 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:
+- [installation](https://cornucopia.readthedocs.io/en/latest/install/)
+- [get started](https://cornucopia.readthedocs.io/en/latest/start/)
+- [examples](https://cornucopia.readthedocs.io/en/latest/examples/overview/)
+- [API](https://cornucopia.readthedocs.io/en/latest/api/overview/)
+
+## Other augmentation packages
+
+There are other great, and much more mature, augmentation packages 
+out-there (although few run on the GPU). Here's a non-exhaustive list:
+- [MONAI](https://github.com/Project-MONAI/MONAI)
+- [TorchIO](https://github.com/fepegar/torchio)
+- [Albumentations](https://github.com/albumentations-team/albumentations) (2D only)
+- [Volumentations](https://github.com/ZFTurbo/volumentations) (3D extension of Albumentations)
+
+## Contributions
+
+If you find this project useful and wish to contribute, please reach out!

{cornucopia-0.2.0 → cornucopia-0.4.0}/README.md

@@ -10,7 +10,7 @@ mostly targeted at (medical) imaging. `cornucopia` is written using a PyTorch
 backend, and therefore runs **on the CPU or GPU**.
 
 Cornucopia is *intended* to be used on the GPU for on-line augmentation.
-A quick [benchmark](examples/benchmark.ipynb) of affine and elastic augmentation
+A quick [benchmark](docs/examples/benchmark.ipynb) of affine and elastic augmentation
 shows that while cornucopia is slower than [TorchIO](https://github.com/fepegar/torchio)
 on the CPU (~ 3s vs 1s), it is greatly accelerated on the GPU (~ 50ms).
 
@@ -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
@@ -37,12 +34,18 @@ independent of the downstream learning framework (pytorch, tensorflow, jax, ...)
 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.2.0 → cornucopia-0.4.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.2.0 → cornucopia-0.4.0}/cornucopia/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-11-21T15:24:48-0500",
+ "date": "2025-04-16T18:13:05+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "6d09025573589db13d99c6df247f138f22f5ab61",
- "version": "0.2.0"
+ "full-revisionid": "6f8ab58dfcfe8978c9aa9e8b05898dcf7d75bb5b",
+ "version": "0.4.0"
 }
 '''  # END VERSION_JSON
 

{cornucopia-0.2.0 → cornucopia-0.4.0}/cornucopia/base.py

@@ -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,6 +447,8 @@ class SequentialTransform(SpecialMixin, SharedMixin, Transform):
     def make_final(self, x, max_depth=float('inf')):
         if max_depth == 0:
             return self
+        if self.is_final:
+            return self
         # x = VirtualTensor.from_any(x, compute_stats=True)
         trf = []
         for t in self:
@@ -477,12 +479,23 @@ class SequentialTransform(SpecialMixin, SharedMixin, Transform):
             x = args[0]
         else:
             return None
-        for trf in self.transforms:
+        for trf in self:
             with IncludeKeysTransform(trf, self.include), \
                  ExcludeKeysTransform(trf, self.exclude):
                 x = trf(x)
         return x
 
+    def xform(self, x):
+        # This should only be called when a Layer's `make_final` returns
+        # a `SequentialTransform`` (i.e., it is created implictly under
+        # the hood, not explicitly by the user).
+        # In such cases, `shared=False` and hopefully we can just fallback
+        # to `forward()`.
+        #
+        # FIXME
+        #   what happens if there's weird stuff in returns/include/exclude?
+        return self(x)
+
     def __len__(self):
         return len(self.transforms)
 
@@ -556,9 +569,11 @@ class MaybeTransform(SpecialMixin, SharedMixin, Transform):
         img = (0.2 * gauss)(img)
         ```
         ```
-    """
 
-    def __init__(self, transform, prob=0.5, *, shared=False, **kwargs):
+    !!! changedin "![v0.4](https://img.shields.io/badge/v0.4-yellow) \
+                   Default for `shared` changed from `False` to `True`"
+    """
+    def __init__(self, transform, prob=0.5, *, shared=True, **kwargs):
         """
 
         Parameters
@@ -617,9 +632,12 @@ class SwitchTransform(SpecialMixin, SharedMixin, Transform):
         ```python
         img = cc.switch({gauss: 0.5, chi: 0.5})(img)
         ```
+
+    !!! changedin "![v0.4](https://img.shields.io/badge/v0.4-yellow) \
+                   Default for `shared` changed from `False` to `True`"
     """
 
-    def __init__(self, transforms, prob=0, *, shared=False, **kwargs):
+    def __init__(self, transforms, prob=0, *, shared=True, **kwargs):
         """
 
         Parameters
@@ -1068,33 +1086,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 +1182,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.2.0 → cornucopia-0.4.0}/cornucopia/baseutils.py

@@ -117,6 +117,29 @@ def returns_find(flag, returned, returns):
             return None
 
 
+def returns_update(value, flag, returned, returns):
+    """Find tensor corresponding to flag in returned structure"""
+    if returns is None:
+        if flag == 'output':
+            return value
+        else:
+            return None
+    if isinstance(returns, dict):
+        if flag in returns:
+            returned[flag] = value
+        return returned
+    elif isinstance(returns, (list, tuple)):
+        if flag in returns:
+            returned[returns.index(flag)] = value
+        return returned
+    else:
+        assert isinstance(returns, str)
+        if returns == flag:
+            return value
+        else:
+            return None
+
+
 def flatstruct(x):
     """Flatten a nested structure of tensors"""
 

{cornucopia-0.2.0 → cornucopia-0.4.0}/cornucopia/contrast.py

@@ -184,7 +184,7 @@ class ContrastLookupTransform(NonFinalTransform):
 
         vmin, vmax = x.min(), x.max()
         edges = torch.linspace(vmin, vmax, self.nk+1)
-        new_mu = torch.rand(self.nk) * (vmax - vmin) + vmin
+        new_mu = torch.rand(self.nk).to(x) * (vmax - vmin) + vmin
         return self.LookupFinalTransform(
             edges, new_mu, **self.get_prm()
         ).make_final(x, max_depth-1)

{cornucopia-0.2.0 → cornucopia-0.4.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
@@ -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
@@ -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"""