datasets

Implement data handling classes.

Classes:

• Augmentation –

  Base class for data augmentations.

• AugmentationFlip –

  Random flip augmentation.

• AugmentationGaussianNoise –

  Random Gaussian noise augmentation.

• AugmentationPoissonNoise –

  Random Poisson noise augmentation.

• AugmentationRotation –

  Random rotation augmentation.

• DataHandler –

  Provide base interface.

• DatasetImagesStack –

  Handle on-disk datasets made of a stack of images.

• DatasetNumpy –

  Handle in-memory datasets.

• DatasetsList –

  Handle lists of datasets.

Functions:

• data_to_tensor –

  Convert a NumPy array to a PyTorch tensor.

• get_batches –

  Generate batches of indices or a single slice for the entire dataset.

• get_flip_axes –

  Generate all possible combinations of dimensions to flip for a given list of axes.

• get_flip_dims –

  Generate all possible combinations of dimensions to flip for a given number of dimensions.

• random_flips –

  Randomly flip images along specified dimensions.

• random_rotations –

  Randomly rotate images by multiples of 90 degrees.

Augmentation

Augmentation(rng: Generator | None = None)

Bases: ABC

Base class for data augmentations.

Source code in src/autoden/algorithms/datasets.py

def __init__(self, rng: np.random.Generator | None = None) -> None:
    if rng is None:
        rng = np.random.default_rng()
    self.rng = rng
    super().__init__()

AugmentationFlip

AugmentationFlip(
    axes: Sequence[int] | None = None,
    n_dims: int | None = None,
    rng: Generator | None = None,
)

Bases: Augmentation

Random flip augmentation.

The axes or n_dims parameter should be set at the same time.

Parameters:

• axes (Sequence[int] | None, default: None ) –

  The axes of the flips, by default None

• n_dims (int | None, default: None ) –

  The dimensions of the flips, by default None

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self, axes: Sequence[int] | None = None, n_dims: int | None = None, rng: np.random.Generator | None = None
) -> None:
    """Initialize the random flip augmentation class.

    The `axes` or `n_dims` parameter should be set at the same time.

    Parameters
    ----------
    axes : Sequence[int] | None, optional
        The axes of the flips, by default None
    n_dims : int | None, optional
        The dimensions of the flips, by default None
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.
    """
    super().__init__(rng)

    if axes is None and n_dims is None:
        self.flips = None
    elif n_dims is None and axes is not None:
        self.flips = get_flip_axes(axes)
    elif axes is None and n_dims is not None:
        self.flips = get_flip_dims(n_dims)
    else:
        raise ValueError("The parameters `axes` and `n_dims` cannot be used at the same time.")

AugmentationGaussianNoise

AugmentationGaussianNoise(
    sigma: float | Sequence[float] | tuple[float, float],
    n: int = 1,
    rng: Generator | None = None,
)

Bases: Augmentation

Random Gaussian noise augmentation.

Parameters:

• sigma (float | Sequence[float] | tuple[float, float]) –

  The standard deviation(s) of the Gaussian noise. If a single float is provided, it will be used for all elements. If a sequence is provided, it will be rotated and used for the first n elements. If a tuple is provided, it should be a range (min, max), and a random value will be chosen from this range for each element.

• n (int, default: 1 ) –

  The number of elements to add noise to, by default 1

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self, sigma: float | Sequence[float] | tuple[float, float], n: int = 1, rng: np.random.Generator | None = None
) -> None:
    """Initialize the Gaussian noise augmentation class.

    Parameters
    ----------
    sigma : float | Sequence[float] | tuple[float, float]
        The standard deviation(s) of the Gaussian noise.
        If a single float is provided, it will be used for all elements.
        If a sequence is provided, it will be rotated and used for the first `n` elements.
        If a tuple is provided, it should be a range (min, max), and a random value will be chosen from this range for each element.
    n : int, optional
        The number of elements to add noise to, by default 1
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.
    """
    super().__init__(rng)

    self.sigma = sigma
    self.n = n

AugmentationPoissonNoise

AugmentationPoissonNoise(
    n_10_counts: (
        float | Sequence[float] | tuple[float, float]
    ),
    n: int = 1,
    rng: Generator | None = None,
)

Bases: Augmentation

Random Poisson noise augmentation.

Parameters:

• n_10_counts (float | Sequence[float] | tuple[float, float]) –

  The average number of counts (in log10) to multiply and de-multiply to bring the values in the desired intensity range. If a single float is provided, it will be used for all elements. If a sequence is provided, it will be rotated and used for the first n elements. If a tuple is provided, it should be a range (min, max), and a random value will be chosen from this range for each element.

• n (int, default: 1 ) –

  The number of elements to add noise to, by default 1

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self, n_10_counts: float | Sequence[float] | tuple[float, float], n: int = 1, rng: np.random.Generator | None = None
) -> None:
    """Initialize the Poisson noise augmentation class.

    Parameters
    ----------
    n_10_counts : float | Sequence[float] | tuple[float, float]
        The average number of counts (in log10) to multiply and de-multiply to bring the values in the desired intensity range.
        If a single float is provided, it will be used for all elements.
        If a sequence is provided, it will be rotated and used for the first `n` elements.
        If a tuple is provided, it should be a range (min, max), and a random value will be chosen from this range for each element.
    n : int, optional
        The number of elements to add noise to, by default 1
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.
    """
    super().__init__(rng)

    self.n_10_counts = n_10_counts
    self.n = n

AugmentationRotation

AugmentationRotation(
    dims: tuple[int, int] | None = None,
    rng: Generator | None = None,
)

Bases: Augmentation

Random rotation augmentation.

Parameters:

• dims (tuple[int, int], default: None ) –

  The dimensions to rotate, by default (-2, -1)

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Source code in src/autoden/algorithms/datasets.py

def __init__(self, dims: tuple[int, int] | None = None, rng: np.random.Generator | None = None) -> None:
    """Initialize the rotation augmentation class.

    Parameters
    ----------
    dims : tuple[int, int]
        The dimensions to rotate, by default (-2, -1)
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.
    """
    super().__init__(rng)

    self.dims = dims

DataHandler

Bases: Dataset, ABC

Provide base interface.

Attributes:

• shape (tuple) –

  Return shape of the dataset.

shape abstractmethod property

shape: tuple

Return shape of the dataset.

DatasetImagesStack

DatasetImagesStack(
    files_pattern: str | Path,
    device: str,
    n_dims: int = 2,
    channel_axis: int | None = None,
    dtype: DTypeLike = float32,
    verbose: bool = False,
)

Bases: DataHandler

Handle on-disk datasets made of a stack of images.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self,
    files_pattern: str | Path,
    device: str,
    n_dims: int = 2,
    channel_axis: int | None = None,
    dtype: DTypeLike = np.float32,
    verbose: bool = False,
) -> None:
    super().__init__()
    files_pattern = Path(files_pattern).expanduser().absolute()
    self.paths = sorted(Path(files_pattern.parent).glob(files_pattern.name))
    if verbose:
        print(f"{self.__class__.__name__}: Found the following images: {self.paths}")
    self.device = device
    self.n_dims = n_dims
    self.channel_axis = channel_axis
    self.dtype = dtype
    self.verbose = verbose

    if len(self.paths) == 0:
        raise ValueError(f"{self.__class__.__name__}: No images found for path: {files_pattern}")

    self._shape = self[0].shape
    if self.shape[0] == 1:
        self._shape = (len(self), *self._shape[1:])
    else:
        self._shape = (len(self), *self._shape)

DatasetNumpy

DatasetNumpy(
    data: NDArray,
    device: str,
    n_dims: int = 2,
    channel_axis: int | None = None,
    dtype: DTypeLike = float32,
    verbose: bool = False,
    pre_load_device: bool = True,
)

Bases: DataHandler

Handle in-memory datasets.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self,
    data: NDArray,
    device: str,
    n_dims: int = 2,
    channel_axis: int | None = None,
    dtype: DTypeLike = np.float32,
    verbose: bool = False,
    pre_load_device: bool = True,
) -> None:
    super().__init__()

    self.device = device
    self.n_dims = n_dims
    self.channel_axis = channel_axis
    self.dtype = dtype
    self.verbose = verbose
    self.pre_load_device = pre_load_device

    self._shape = data.shape

    device_to_use = self.device if self.pre_load_device else None
    self.data = data_to_tensor(
        data, device=device_to_use, n_dims=self.n_dims, channel_axis=self.channel_axis, dtype=self.dtype
    )

DatasetsList

DatasetsList(
    datasets: Sequence[DataHandler],
    augmentation: (
        str
        | Augmentation
        | Sequence[str | Augmentation]
        | None
    ) = None,
)

Bases: Dataset

Handle lists of datasets.

Source code in src/autoden/algorithms/datasets.py

def __init__(
    self, datasets: Sequence[DataHandler], augmentation: str | Augmentation | Sequence[str | Augmentation] | None = None
) -> None:
    super().__init__()

    def _convert_augmentation(aug: str | Augmentation) -> Augmentation:
        if isinstance(aug, Augmentation):
            return aug
        if aug.lower() == "flip":
            return AugmentationFlip()
        if aug.lower() == "rot":
            return AugmentationRotation()
        raise ValueError(f"Unrecognized augmentation: {aug}")

    self.datasets = list(datasets)

    if augmentation is None:
        augmentation = []
    elif isinstance(augmentation, str | Augmentation):
        augmentation = [augmentation]
    self.augmentation = [_convert_augmentation(aug) for aug in augmentation]

    if not self.datasets:
        raise ValueError("The argument `datasets` cannot be an empty Sequence")

    self._length = len(self.datasets[0])
    if any(self._length != len(d) for d in self.datasets[1:]):
        raise ValueError(
            f"Datasets should all have the same length, but these lengths were found: {[len(d) for d in self.datasets]}"
        )

    self._min_n_dims = min(len(d.shape) for d in self.datasets)

data_to_tensor

data_to_tensor(
    data: NDArray,
    device: str | None,
    n_dims: int = 2,
    channel_axis: int | None = None,
    dtype: DTypeLike | None = float32,
) -> Tensor

Convert a NumPy array to a PyTorch tensor.

Parameters:

• data (NDArray) –

  The input data to be converted to a tensor.

• device (str or None) –

  The device to which the tensor should be moved (e.g., 'cpu', 'cuda').

• n_dims (int, default: 2 ) –

  The number of dimensions to consider for the data shape, by default 2.

• channel_axis (int or None, default: None ) –

  The axis along which the channels are stacked, by default None.

• dtype (DTypeLike or None, default: float32 ) –

  The data type to which the data should be converted, by default np.float32.

Returns:

• Tensor –

  The converted PyTorch tensor.

Notes

If channel_axis is provided, the data is moved to the specified axis. Otherwise, the data is expanded to include an additional dimension. The data is then reshaped and converted to the specified data type before being converted to a PyTorch tensor and moved to the specified device.

Source code in src/autoden/algorithms/datasets.py

def data_to_tensor(
    data: NDArray, device: str | None, n_dims: int = 2, channel_axis: int | None = None, dtype: DTypeLike | None = np.float32
) -> pt.Tensor:
    """
    Convert a NumPy array to a PyTorch tensor.

    Parameters
    ----------
    data : NDArray
        The input data to be converted to a tensor.
    device : str or None
        The device to which the tensor should be moved (e.g., 'cpu', 'cuda').
    n_dims : int, optional
        The number of dimensions to consider for the data shape, by default 2.
    channel_axis : int or None, optional
        The axis along which the channels are stacked, by default None.
    dtype : DTypeLike or None, optional
        The data type to which the data should be converted, by default np.float32.

    Returns
    -------
    pt.Tensor
        The converted PyTorch tensor.

    Notes
    -----
    If `channel_axis` is provided, the data is moved to the specified axis.
    Otherwise, the data is expanded to include an additional dimension.
    The data is then reshaped and converted to the specified data type before
    being converted to a PyTorch tensor and moved to the specified device.
    """
    if channel_axis is not None:
        num_channels = data.shape[channel_axis]
        data = np.moveaxis(data, channel_axis, -n_dims - 1)
    else:
        num_channels = 1
        data = np.expand_dims(data, -n_dims - 1)
    data_shape = data.shape[-n_dims:]
    data = data.reshape([-1, num_channels, *data_shape])
    if dtype is not None:
        # # If complex, we promote the type to the lowest required and available precision
        # # Deactivated for the moment, because we don't handle complex weights. We handle
        # # this in the algorithm prepare functions.
        # if np.iscomplexobj(data):
        #     dtype = np.promote_types(dtype, np.complex64)
        data = data.astype(dtype)
    return pt.tensor(data, device=device)

get_batches

get_batches(
    num_instances: int, batch_size: int | None = None
) -> list[slice]

Generate batches of indices or a single slice for the entire dataset.

Parameters:

• num_instances (int) –

  The total number of instances.

• batch_size (int | None, default: None ) –

  The size of each batch. If None, a single slice covering the entire dataset is returned. Default is None.

Returns:

• list[slice] –

  A list of slice objects representing batch indices or a single slice object covering the entire dataset.

Examples:

>>> get_batches(10, 3)
[range(0, 3), range(3, 6), range(6, 9), range(9, 10)]

>>> get_batches(10, None)
[slice(None)]

Source code in src/autoden/algorithms/datasets.py

def get_batches(num_instances: int, batch_size: int | None = None) -> list[slice]:
    """
    Generate batches of indices or a single slice for the entire dataset.

    Parameters
    ----------
    num_instances : int
        The total number of instances.
    batch_size : int | None, optional
        The size of each batch. If None, a single slice covering the entire dataset is returned. Default is None.

    Returns
    -------
    list[slice]
        A list of slice objects representing batch indices or a single slice object covering the entire dataset.

    Examples
    --------
    >>> get_batches(10, 3)
    [range(0, 3), range(3, 6), range(6, 9), range(9, 10)]

    >>> get_batches(10, None)
    [slice(None)]
    """
    if batch_size is not None:
        return [slice(ii, min(ii + batch_size, num_instances)) for ii in range(0, num_instances, batch_size)]
    else:
        return [slice(None)]

get_flip_axes

get_flip_axes(
    axes: Sequence[int],
) -> Sequence[tuple[int, ...]]

Generate all possible combinations of dimensions to flip for a given list of axes.

Parameters:

• axes (Sequence[int]) –

  The list of axes.

Returns:

• Sequence[tuple[int, ...]] –

  A sequence of tuples, where each tuple represents a combination of dimensions to flip. The dimensions are represented by negative indices, ranging from -n_dims to -1.

Examples:

>>> _get_flip_axes((-2, -1))
[(), (-2,), (-1,), (-2, -1)]

Source code in src/autoden/algorithms/datasets.py

def get_flip_axes(axes: Sequence[int]) -> Sequence[tuple[int, ...]]:
    """
    Generate all possible combinations of dimensions to flip for a given list of axes.

    Parameters
    ----------
    axes : Sequence[int]
        The list of axes.

    Returns
    -------
    Sequence[tuple[int, ...]]
        A sequence of tuples, where each tuple represents a combination of dimensions to flip.
        The dimensions are represented by negative indices, ranging from -n_dims to -1.

    Examples
    --------
    >>> _get_flip_axes((-2, -1))
    [(), (-2,), (-1,), (-2, -1)]
    """
    return sum([[*combinations(axes, d)] for d in range(len(axes) + 1)], [])

get_flip_dims

get_flip_dims(n_dims: int) -> Sequence[tuple[int, ...]]

Generate all possible combinations of dimensions to flip for a given number of dimensions.

Parameters:

• n_dims (int) –

  The number of dimensions.

Returns:

• Sequence[tuple[int, ...]] –

  A sequence of tuples, where each tuple represents a combination of dimensions to flip. The dimensions are represented by negative indices, ranging from -n_dims to -1.

Examples:

>>> _get_flip_dims(2)
[(), (-2,), (-1,), (-2, -1)]

Source code in src/autoden/algorithms/datasets.py

def get_flip_dims(n_dims: int) -> Sequence[tuple[int, ...]]:
    """
    Generate all possible combinations of dimensions to flip for a given number of dimensions.

    Parameters
    ----------
    n_dims : int
        The number of dimensions.

    Returns
    -------
    Sequence[tuple[int, ...]]
        A sequence of tuples, where each tuple represents a combination of dimensions to flip.
        The dimensions are represented by negative indices, ranging from -n_dims to -1.

    Examples
    --------
    >>> _get_flip_dims(2)
    [(), (-2,), (-1,), (-2, -1)]
    """
    return get_flip_axes(range(-n_dims, 0))

random_flips

random_flips(
    *imgs: Tensor,
    flips: Sequence[tuple[int, ...]] | None = None,
    rng: Generator | None = None
) -> Sequence[Tensor]

Randomly flip images along specified dimensions.

Parameters:

• *imgs (Tensor, default: () ) –

  The input images to be flipped.

• flips (Sequence[tuple[int, ...]] | None, default: None ) –

  The possible flip dimensions to choose from. If None, it will call _get_flip_dims on the ndim of the first image. By default None.

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Returns:

• Sequence[Tensor] –

  The flipped images.

Source code in src/autoden/algorithms/datasets.py

def random_flips(
    *imgs: pt.Tensor, flips: Sequence[tuple[int, ...]] | None = None, rng: np.random.Generator | None = None
) -> Sequence[pt.Tensor]:
    """Randomly flip images along specified dimensions.

    Parameters
    ----------
    *imgs : torch.Tensor
        The input images to be flipped.
    flips : Sequence[tuple[int, ...]] | None, optional
        The possible flip dimensions to choose from. If None, it will call _get_flip_dims on the ndim of the first image.
        By default None.
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.

    Returns
    -------
    Sequence[torch.Tensor]
        The flipped images.
    """
    if flips is None:
        flips = get_flip_dims(imgs[0].ndim - 2)
    if rng is None:
        rng = np.random.default_rng()
    rand_val = int(rng.integers(0, len(flips)))

    flip = flips[rand_val]
    return [pt.flip(im, flip) for im in imgs]

random_rotations

random_rotations(
    *imgs: Tensor,
    dims: tuple[int, int] | None = None,
    rng: Generator | None = None
) -> Sequence[Tensor]

Randomly rotate images by multiples of 90 degrees.

Parameters:

• *imgs (Tensor, default: () ) –

  The input images to be rotated.

• dims (tuple[int, int], default: None ) –

  The dimensions to rotate. By default (-2, -1).

• rng (Generator | None, default: None ) –

  The random number generator to use. If None, a default generator will be used. By default None.

Returns:

• Sequence[Tensor] –

  The rotated images.

Source code in src/autoden/algorithms/datasets.py

def random_rotations(
    *imgs: pt.Tensor, dims: tuple[int, int] | None = None, rng: np.random.Generator | None = None
) -> Sequence[pt.Tensor]:
    """Randomly rotate images by multiples of 90 degrees.

    Parameters
    ----------
    *imgs : torch.Tensor
        The input images to be rotated.
    dims : tuple[int, int], optional
        The dimensions to rotate. By default (-2, -1).
    rng : np.random.Generator | None, optional
        The random number generator to use. If None, a default generator will be used.
        By default None.

    Returns
    -------
    Sequence[torch.Tensor]
        The rotated images.
    """
    rand_val = np.random.randint(4)
    if dims is None:
        dims = (-2, -1)
    if rng is None:
        rng = np.random.default_rng()
    rand_val = int(rng.integers(0, 4))

    if rand_val > 0:
        return [pt.rot90(im, k=rand_val, dims=dims) for im in imgs]
    else:
        return imgs