#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""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)
"""
import collections
from dataclasses import dataclass, field
from typing import Any, Callable, Sequence
import torch
from torchvision.transforms import v2
from torchvision.transforms.v2 import Transform
from torchvision.transforms.v2 import functional as F # noqa: N812
[docs]
class RandomSubsetApply(Transform):
"""Apply a random subset of N transformations from a list of transformations.
Args:
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.
"""
[docs]
def __init__(
self,
transforms: Sequence[Callable],
p: list[float] | None = None,
n_subset: int | None = None,
random_order: bool = False,
) -> None:
super().__init__()
if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables")
if p is None:
p = [1] * len(transforms)
elif len(p) != len(transforms):
raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
)
if n_subset is None:
n_subset = len(transforms)
elif not isinstance(n_subset, int):
raise TypeError("n_subset should be an int or None")
elif not (1 <= n_subset <= len(transforms)):
raise ValueError(f"n_subset should be in the interval [1, {len(transforms)}]")
self.transforms = transforms
total = sum(p)
self.p = [prob / total for prob in p]
self.n_subset = n_subset
self.random_order = random_order
self.selected_transforms = None
[docs]
def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1
selected_indices = torch.multinomial(torch.tensor(self.p), self.n_subset)
if not self.random_order:
selected_indices = selected_indices.sort().values
self.selected_transforms = [self.transforms[i] for i in selected_indices]
for transform in self.selected_transforms:
outputs = transform(*inputs)
inputs = outputs if needs_unpacking else (outputs,)
return outputs
[docs]
def extra_repr(self) -> str:
return (
f"transforms={self.transforms}, "
f"p={self.p}, "
f"n_subset={self.n_subset}, "
f"random_order={self.random_order}"
)
[docs]
class SharpnessJitter(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 :class:`torch.Tensor`,
it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
Args:
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.
"""
[docs]
def __init__(self, sharpness: float | Sequence[float]) -> None:
super().__init__()
self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness):
if isinstance(sharpness, (int, float)):
if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.")
sharpness = [1.0 - sharpness, 1.0 + sharpness]
sharpness[0] = max(sharpness[0], 0.0)
elif isinstance(sharpness, collections.abc.Sequence) and len(sharpness) == 2:
sharpness = [float(v) for v in sharpness]
else:
raise TypeError(f"{sharpness=} should be a single number or a sequence with length 2.")
if not 0.0 <= sharpness[0] <= sharpness[1]:
raise ValueError(f"sharpnesss values should be between (0., inf), but got {sharpness}.")
return float(sharpness[0]), float(sharpness[1])
[docs]
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
"""Generate random parameters for sharpness jitter.
Args:
flat_inputs: List of input tensors.
Returns:
Dictionary containing 'sharpness_factor' sampled uniformly from
the configured sharpness range.
"""
sharpness_factor = torch.empty(1).uniform_(self.sharpness[0], self.sharpness[1]).item()
return {"sharpness_factor": sharpness_factor}