opentau.datasets.transforms

Image transformation utilities for data augmentation.

This module provides configurable image transformation pipelines for data augmentation during training. It extends torchvision.transforms.v2 with custom transforms and a flexible configuration system that supports weighted random sampling of transform subsets.

The module implements a probabilistic augmentation strategy where a random subset of available transforms is applied to each image, with configurable weights controlling the sampling probability. This approach provides more diverse augmentations compared to applying all transforms deterministically.

Key Features:

• Random subset sampling: Applies a random subset of N transforms from a larger pool, controlled by configurable weights.

• Custom transforms: Includes SharpnessJitter for more diverse sharpness augmentation compared to standard torchvision transforms.

• Configurable pipeline: Dataclass-based configuration system for easy customization of transform parameters and weights.

• Flexible ordering: Option to apply transforms in random order or deterministic order.

• Torchvision v2 integration: Built on top of torchvision.transforms.v2 for modern transform API support.

Classes:

RandomSubsetApply: Transform container that applies a random subset of

transforms with weighted sampling and optional random ordering.

SharpnessJitter: Custom transform for randomly jittering image sharpness

with uniform distribution sampling.

ImageTransformConfig: Dataclass for configuring individual transform

parameters (weight, type, kwargs).

ImageTransformsConfig: Dataclass for configuring the overall transform

pipeline (enable flag, max transforms, random order, transform list).

ImageTransforms: Main transform class that composes transforms based on

configuration.

Functions:

make_transform_from_config: Factory function to create transform instances

from ImageTransformConfig.

Example

Create and use image transforms:

>>> config = ImageTransformsConfig(
...     enable=True,
...     max_num_transforms=3,
...     random_order=True
... )
>>> transforms = ImageTransforms(config)
>>> augmented_image = transforms(image_tensor)

Functions

make_transform_from_config(cfg)

Create a transform instance from an ImageTransformConfig.

Classes

ImageTransformConfig(weight, type, kwargs, ...)	For each transform, the following parameters are available:
ImageTransforms(cfg)	A class to compose image transforms based on configuration.
ImageTransformsConfig(enable, ...)	These transforms are all using standard torchvision.transforms.v2 You can find out how these transformations affect images here: https://pytorch.org/vision/0.18/auto_examples/transforms/plot_transforms_illustrations.html We use a custom RandomSubsetApply container to sample them.
RandomSubsetApply(transforms[, p, n_subset, ...])	Apply a random subset of N transformations from a list of transformations.
SharpnessJitter(sharpness)	Randomly change the sharpness of an image or video.

class opentau.datasets.transforms.ImageTransformConfig(weight: float = 1.0, type: str = 'Identity', kwargs: dict[str, ~typing.Any] = <factory>)

Bases: object

For each transform, the following parameters are available:

weight: This represents the multinomial probability (with no replacement)

used for sampling the transform. If the sum of the weights is not 1, they will be normalized.

type: The name of the class used. This is either a class available under torchvision.transforms.v2 or a

custom transform defined here.

kwargs: Lower & upper bound respectively used for sampling the transform’s parameter

(following uniform distribution) when it’s applied.

__init__(weight: float = 1.0, type: str = 'Identity', kwargs: dict[str, ~typing.Any] = <factory>) → None

kwargs: dict[str, Any]

type: str = 'Identity'

weight: float = 1.0

class opentau.datasets.transforms.ImageTransforms(cfg: ImageTransformsConfig)

Bases: Transform

A class to compose image transforms based on configuration.

__init__(cfg: ImageTransformsConfig) → None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(*inputs: Any) → Any

Do not override this! Use transform() instead.

class opentau.datasets.transforms.ImageTransformsConfig(enable: bool = False, max_num_transforms: int = 3, random_order: bool = False, tfs: dict[str, ~opentau.datasets.transforms.ImageTransformConfig] = <factory>)

Bases: object

These transforms are all using standard torchvision.transforms.v2 You can find out how these transformations affect images here: https://pytorch.org/vision/0.18/auto_examples/transforms/plot_transforms_illustrations.html We use a custom RandomSubsetApply container to sample them.

__init__(enable: bool = False, max_num_transforms: int = 3, random_order: bool = False, tfs: dict[str, ~opentau.datasets.transforms.ImageTransformConfig] = <factory>) → None

enable: bool = False

max_num_transforms: int = 3

random_order: bool = False

tfs: dict[str, ImageTransformConfig]

class opentau.datasets.transforms.RandomSubsetApply(transforms: Sequence[Callable], p: list[float] | None = None, n_subset: int | None = None, random_order: bool = False)

Bases: Transform

Apply a random subset of N transformations from a list of transformations.

Parameters:

• transforms – list of transformations.

• p – represents the multinomial probabilities (with no replacement) used for sampling the transform. If the sum of the weights is not 1, they will be normalized. If None (default), all transforms have the same probability.

• n_subset – number of transformations to apply. If None, all transforms are applied. Must be in [1, len(transforms)].

• random_order – apply transformations in a random order.

__init__(transforms: Sequence[Callable], p: list[float] | None = None, n_subset: int | None = None, random_order: bool = False) → None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

extra_repr() → str

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(*inputs: Any) → Any

Do not override this! Use transform() instead.

class opentau.datasets.transforms.SharpnessJitter(sharpness: float | Sequence[float])

Bases: Transform

Randomly change the sharpness of an image or video.

Similar to a v2.RandomAdjustSharpness with p=1 and a sharpness_factor sampled randomly. While v2.RandomAdjustSharpness applies — with a given probability — a fixed sharpness_factor to an image, SharpnessJitter applies a random sharpness_factor each time. This is to have a more diverse set of augmentations as a result.

A sharpness_factor of 0 gives a blurred image, 1 gives the original image while 2 increases the sharpness by a factor of 2.

If the input is a torch.Tensor, it is expected to have […, 1 or 3, H, W] shape, where … means an arbitrary number of leading dimensions.

Parameters:

sharpness – How much to jitter sharpness. sharpness_factor is chosen uniformly from [max(0, 1 - sharpness), 1 + sharpness] or the given [min, max]. Should be non negative numbers.

__init__(sharpness: float | Sequence[float]) → None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

make_params(flat_inputs: list[Any]) → dict[str, Any]

Generate random parameters for sharpness jitter.

Parameters:

flat_inputs – List of input tensors.

Returns:

Dictionary containing ‘sharpness_factor’ sampled uniformly from the configured sharpness range.

transform(inpt: Any, params: dict[str, Any]) → Any

Apply sharpness adjustment to input.

Parameters:

• inpt – Input image or video tensor.

• params – Dictionary containing ‘sharpness_factor’ from make_params.

Returns:

Transformed image or video with adjusted sharpness.

opentau.datasets.transforms.make_transform_from_config(cfg: ImageTransformConfig) → Transform

Create a transform instance from an ImageTransformConfig.

Parameters:

cfg – Configuration object specifying the transform type and parameters.

Returns:

Transform instance (Identity, ColorJitter, or SharpnessJitter).

Raises:

ValueError – If the transform type is not recognized.