Core API

Compose is the main public entry point. It accepts a list of recognized backend transform objects and builds a segmented torch.nn.Module around the supported (B, C, H, W) tensor contract.

Compose and FusedCompose refer to the same class. AugmentationSequential is another naming alias; these aliases do not imply complete behavioral compatibility with every native container option.

Choose a construction style

import torch
import torchvision.transforms.v2 as transforms

from fuse_augmentations import Compose

pipe = Compose(
    [
        transforms.RandomRotation(15),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomAffine(degrees=0, scale=(0.9, 1.1)),
    ]
)

image = torch.rand(8, 3, 224, 224)
output = pipe(image)
print(pipe.fusion_plan)

Fusion plan for the native TorchVision pipeline

fused(RandomRotation, RandomHorizontalFlip, RandomAffine)

Use Compose.from_params for an augmentation-backend-free pipeline and Compose.from_config for declarative backend-specific construction. Both are documented in Configuration API.

Input and output forms

The standard call accepts a BCHW tensor. Set data_keys when positional auxiliary targets are present:

Use only explicitly supported spatial transforms with auxiliary targets

The runtime refusal policy is a finite class-name list. An unknown crop,
resize, or custom spatial transform can modify the image while leaving
masks, boxes, or keypoints stale. Treat every
`Unknown ... SPATIAL_KERNEL barrier` warning as unsafe with `data_keys`.
Read [Auxiliary targets](../guides/auxiliary-targets.md) before using this
path for training data.

import torch
import torchvision.transforms.v2 as transforms

from fuse_augmentations import Compose

safe_transforms = [
    transforms.RandomRotation(15),
    transforms.RandomHorizontalFlip(p=0.5),
]

pipe = Compose(
    safe_transforms,
    data_keys=["input", "mask", "bbox_xyxy", "keypoints"],
)

image = torch.rand(2, 3, 64, 64)
mask = torch.zeros(2, 1, 64, 64, dtype=torch.long)
boxes = torch.tensor([[[8.0, 8.0, 32.0, 32.0]]] * 2)
keypoints = torch.tensor([[[16.0, 16.0]]] * 2)

image_out, mask_out, boxes_out, keypoints_out = pipe(
    image,
    mask,
    boxes,
    keypoints,
)

output_backend="numpy" converts images and masks to channel-last NumPy arrays. Coordinate outputs remain tensors.

An image-only Albumentations pipeline additionally accepts pipe(image=hwc_array) and returns an Albumentations-style dictionary. Auxiliary keyword keys are not supported on that native NumPy path; use tensor inputs and data_keys instead.

Inspect a pipeline

for descriptor in pipe.fusion_plan_descriptors:
    print(descriptor.kind, descriptor.transforms, descriptor.n_warps_saved)

output, last_matrix = pipe(image, return_matrix=True)

Segment descriptor and saved warp count for the pipeline

fused ('RandomRotation', 'RandomHorizontalFlip', 'RandomAffine') 2

fusion_plan, fusion_plan_descriptors, and n_warps_saved describe the constructed plan. transform_matrix and return_matrix=True report only the most recent call's last matrix-producing segment, not a whole-pipeline composition.

Constructor options that change semantics

Option	Effect
reorder	Preserves order by default; POINTWISE may move color operations to extend geometric runs
randomness	Preserves backend sampling by default; per_sample requests independent draws where supported
execution	Selects cv2 or torch execution for fused Albumentations geometry only
interpolation, padding_mode	Segment-level sampling choices for fused geometry
clip_policy	Controls boundaries/clamping in fused color runs
mask_interpolation	nearest for hard labels or bilinear for floating soft masks
compile	Opts eligible non-CPU torch warp cores into torch.compile
antialias	Enables supported downscale prefiltering
substitute_passthrough	Opts into behavior-changing registered passthrough substitutions
pipeline_dtype	"bfloat16" or "float16" low-precision execution for fused warp/color/LUT cores; default None

Extra backend_kwargs are currently reserved and unused; do not rely on them as an extension mechanism.

Compose / FusedCompose

FusedCompose

FusedCompose(transforms: list[object], reorder: ReorderPolicy = NONE, interpolation: InterpolationStr | None = None, padding_mode: ComposePaddingModeStr | None = None, data_keys: list[str] | None = None, output_backend: Literal['numpy', 'numpy_hwc', 'torch'] | None = None, randomness: RandomnessPolicy | Literal['backend', 'per_sample'] = BACKEND, execution: ExecutionStr = 'cv2', compile: bool = False, antialias: bool = False, substitute_passthrough: bool = False, clip_policy: ClipPolicyStr = 'final', mask_interpolation: MaskInterpolationStr = 'nearest', pipeline_dtype: PipelineDtypeStr | None = None, **backend_kwargs: object)

Bases: FactoriesMixin, IntrospectionMixin, Module

Fused augmentation pipeline that replaces the backend's native Compose.

Segments the transform list into fused geometric segments and passthrough transforms, then executes them sequentially. Consecutive geometric ops are grouped and executed as either:

• A :class:~fuse_augmentations.affine.segment.FusedAffineSegment - when the run contains at least one GEOMETRIC_INTERP op. Matrices are composed and a single grid_sample call is used, eliminating redundant interpolation passes.
• An :class:~fuse_augmentations.affine.segment.ExactAffineSegment - when the run contains only GEOMETRIC_EXACT ops (HFlip, VFlip). Transforms are applied via tensor.flip with zero interpolation error.

SPATIAL_KERNEL and nonlinear POINTWISE transforms are passed through to the backend adapter unchanged. Supported POINTWISE_LINEAR color transforms, including standard Normalize, are folded into a color segment.

ReorderPolicy.POINTWISE is fully implemented: before segmentation, POINTWISE and POINTWISE_LINEAR ops are bubbled past geometric ops within each SPATIAL_KERNEL-bounded stretch, maximising the geometric run length available for fusion.

Parameters:

Name	Type	Description	Default
transforms	list[object]	List of augmentation transform objects.	required
reorder	ReorderPolicy	Reorder policy applied before segmentation. NONE (default) preserves the original order. POINTWISE reorders pointwise ops after geometric chains. AGGRESSIVE currently aliases POINTWISE and is kept for API compatibility with future stronger reorder semantics.	NONE
interpolation	InterpolationStr | None	Interpolation mode override for fused segments ("bilinear", "nearest", "bicubic"). Defaults to "bilinear" when None.	None
padding_mode	ComposePaddingModeStr | None	Padding mode override for fused segments ("zeros", "border", "reflection"), or the opt-in "per_transform" policy. The latter honors compatible transform modes and keeps opaque modes as native boundaries. Defaults to "zeros" when None.	None
data_keys	list[str] | None	List of key names describing positional arguments to :meth:forward. The first key should be "input" (the image). Auxiliary keys ("mask", "bbox_xyxy", "bbox_xywh", "keypoints") are routed through segments and transformed alongside the image. Unknown keys are passed through unchanged with a UserWarning. None preserves backward-compatible single-tensor input/output. Albumentations fused segments route auxiliary targets through the composed pixel matrix, matching the Kornia/TorchVision coordinate convention, so multi-target data_keys are supported for every backend.	None
output_backend	Literal['numpy', 'numpy_hwc', 'torch'] | None	Target output format. "numpy" (or its alias "numpy_hwc") converts the primary image output to a NumPy ndarray with channel-last layout: batched inputs of shape (batch_size, channels, height, width) become (batch_size, height, width, channels), while unbatched inputs of shape (channels, height, width) become (height, width, channels) (i.e. the batch dimension is implicit/squeezed). Variable-batch trap: a batch of size one, (1, channels, height, width), is also squeezed and returns (height, width, channels) -- the leading batch axis is dropped rather than kept as (1, height, width, channels). Callers that loop over a variable batch size must special-case batch_size == 1 (or re-insert the axis with np.expand_dims(out, 0)) to avoid a rank mismatch. "torch" or None keeps the native torch.Tensor output. For multi-target data_keys the conversion is applied per target: image and mask outputs are converted, while coordinate targets (boxes, keypoints) stay tensors.	None
randomness	RandomnessPolicy | Literal['backend', 'per_sample']	Batch randomness policy. "backend" preserves native backend semantics, including TorchVision v2 batch-shared sampling. "per_sample" asks adapters with canonical samplers to draw independent parameters per batch item.	BACKEND
execution	ExecutionStr	Warp strategy for the Albumentations fused segments. "cv2" (default) warps each sample with OpenCV -- bit-exact with the native cv2 backend and fastest on CPU at small batch sizes. "torch" composes the same per-sample matrices (identical sampling and RNG) but applies one batched grid_sample for the whole batch, giving batch-size-independent throughput and a native GPU/MPS warp; its border/bilinear numerics differ slightly from cv2. Only affects Albumentations pipelines; the Kornia/TorchVision backends already run the torch engine.	'cv2'
compile	bool	Opt-in torch.compile of the warp core (matrix normalize -> affine_grid -> grid_sample) of the torch-backed fused segments. Off by default. When True, the compiled region is used only on non-CPU tensors and only when the installed torch is new enough (>= 2.2); otherwise the eager path runs and the flag is a no-op, so the output is unchanged. Probability masking stays outside the compiled region, so the graph has no data-dependent breaks. The first non-CPU call pays a one-time compilation cost; dynamic=True avoids per-shape recompiles.	False
antialias	bool	Opt-in antialiasing for aggressive downscales. Off by default. When True, a crop-resize (or fused geo+crop) segment whose worst-axis scale drops below 0.5 prefilters the input with a Gaussian (mipmap-rule sigma) before the single warp, so the downscale no longer aliases. A no-op when the scale is safe or the flag is off, keeping the default output bit-identical; requires kornia for the blur (falls back to the un-filtered warp when kornia is absent).	False
substitute_passthrough	bool	When True (default False), non-fusible passthrough ops that have a registered torch-native equivalent in an already-installed backend are replaced with that equivalent, so the pipeline can stay on-device instead of paying a device-to-host round-trip. The initial registry maps Albumentations GaussianBlur to Kornia RandomGaussianBlur. This is behaviour-changing: the substitute uses a different kernel implementation, border handling, and per-call random stream, so outputs and RNG differ from the original op — every substitution emits a :class:UserWarning. Substitution happens only when the target backend is importable; otherwise the original op is kept silently and runs on the normal passthrough path. The default (False) leaves outputs unchanged.	False
clip_policy	ClipPolicyStr	Clamp policy for fused color segments. "final" (default) applies one clamp after the composed color matmul — most precise. "per_op_parity" splits the fused color run at any op where a composed intermediate would leave [0, 1] and clamps in between, approximating the native per-op clamped chain on gamut-escaping pipelines. Known gap: when a gamut-escaping op immediately precedes a mean-relative contrast op, the contrast midpoint is taken from the pre-clamp mean, so output can differ from native by ~1e-2.	'final'
mask_interpolation	MaskInterpolationStr	Auxiliary mask sampling mode. "nearest" (default) preserves hard labels; "bilinear" provides differentiable soft-mask sampling and requires floating-point mask input.	'nearest'
pipeline_dtype	PipelineDtypeStr | None	Optional "bfloat16" or "float16" GPU execution dtype for fused affine/projective/crop warps and fused color/LUT applies. Matrix composition and inversion remain float32 or float64, and the returned image keeps its input dtype. CPU ignores this option and uses the existing float32/float64 path. MPS supports both requested dtypes for affine_grid and grid_sample on PyTorch 2.10.	None
**backend_kwargs	object	Reserved for backend-specific options (currently unused).	{}

Initialize FusedCompose.

forward

forward(*args: Tensor, return_matrix: bool = False) -> object

Apply the augmentation pipeline to an image batch and optional auxiliary targets.

Single-tensor mode (data_keys=None, default): accepts one image tensor and returns one tensor or a NumPy ndarray when output_backend="numpy". This is the backward-compatible path.

Multi-target mode (data_keys is set): accepts positional arguments in the same order as data_keys. The first key must map to the image ("input"); subsequent keys are auxiliary targets ("mask", "bbox_xyxy", "bbox_xywh", "keypoints"). Returns a tuple of tensors in data_keys order, or a bare tensor when data_keys contains exactly one entry.

Parameters:

Name	Type	Description	Default
*args	Tensor	Positional tensors matching the data_keys list. args[0] is always the image (B, C, H, W) float32 channel-first. Auxiliary args follow in data_keys[1:] order: "mask" as (B, C, H, W) float/int; "bbox_xyxy" / "bbox_xywh" as (B, N, 4) float32; "keypoints" as (B, N, 2) float32.	()
return_matrix	bool	When True, return the output and its last fused geometric pixel matrix.	False

Returns:

Type	Description
object	Single Tensor or NumPy ndarray when data_keys is
object	None or has one entry. NumPy output is returned only when
object	output_backend="numpy". tuple[Tensor, ...] in
object	data_keys order otherwise.
object	When return_matrix=True, returns (output, matrix) where
object	matrix is the image-dtype pixel matrix from the last fused
object	geometric segment, or None when no such segment ran.

Raises:

Type	Description
TypeError	If the number of positional arguments does not match the number of data_keys entries (when data_keys is set), or if more than one argument is passed when data_keys is None.

Note

Passthrough (non-fused) transforms in the pipeline apply to the image only. Auxiliary targets skip passthrough segments and retain their values from the preceding fused segment. This is by design - passthrough backends do not expose a target-routing API.

output_backend conversion is applied per target in multi-target mode: image and mask outputs are converted to the requested backend, while coordinate targets (boxes, keypoints) remain tensors because the channel-last image layout does not apply to them.

inverse

inverse(image: Tensor, *auxiliary_targets: Tensor, matrix: Tensor | None = None) -> object

Map a paired augmented tensor back to its original geometric frame.

Pass the matrix returned by the exact paired forward(..., return_matrix=True) call. This avoids reading :attr:transform_matrix, which is mutable compatibility state and therefore unsafe to pair with an output from another concurrent call. The method supports one fused affine or projective segment, including a chain already fused into that segment. It applies one grid_sample to the image and routes declared masks, boxes, and keypoints through the inverse pixel matrix.

Keypoints and masks recover to sampling precision. Bounding boxes are axis-aligned (AABB), so a forward-then-inverse box is exact only for axis-aligned transforms (flip, scale, translation) and inflates under a rotation, shear, or projective warp. The matrix is not validated against the paired image: a matrix from a different call yields silently wrong geometry, so always pass the matrix returned by the same forward call.

Parameters:

Name	Type	Description	Default
image	Tensor	Augmented (B, C, H, W) floating-point image tensor.	required
*auxiliary_targets	Tensor	Augmented targets in data_keys[1:] order.	()
matrix	Tensor | None	Forward pixel matrix returned by the paired forward call.	None

Returns:

Type	Description
object	The de-augmented image, or a tuple in data_keys order when
object	auxiliary targets are supplied.

Raises:

Type	Description
TypeError	If the image or matrix has an unsupported tensor shape or dtype.
ValueError	If the pipeline has a non-invertible segment, no paired matrix, or targets that do not match data_keys.

Examples:

>>> import torch
>>> from fuse_augmentations import Compose
>>> pipe = Compose.from_params(translate_x=(2.0, 2.0))
>>> image = torch.rand(1, 3, 8, 8)
>>> augmented, matrix = pipe(image, return_matrix=True)
>>> recovered = pipe.inverse(augmented, matrix=matrix)
>>> recovered.shape
torch.Size([1, 3, 8, 8])

NumPy converters

The converters are useful when conversion should be explicit rather than attached to a pipeline.

NumpyToTorchConverter

Convert NumPy HWC/BHWC arrays to (batch_size, channels, height, width) torch tensors.

All pipeline outputs are (batch_size, channels, height, width) torch.Tensor regardless of backend. uint8 inputs are normalised to float32 in [0, 1] (÷255); uint16 likewise (÷65535); float32 inputs are passed through unchanged. Any other dtype (float64, float16, signed integers) is cast to float32 WITHOUT rescaling — values keep their original range.

convert

convert(array: NDArray[Any]) -> Tensor

Convert a NumPy array (HW/HWC/BHWC) to a float32 torch tensor (BCHW).

Parameters:

Name	Type	Description	Default
array	NDArray[Any]	NumPy array in channel-last layout: (height, width), (height, width, channels), or (batch_size, height, width, channels).	required

Returns:

Type	Description
Tensor	torch.float32 tensor of shape (1, 1, height, width) for 2-D input,
Tensor	(1, channels, height, width) for 3-D input, or (batch_size, channels, height, width) for 4-D input.

Raises:

Type	Description
ValueError	If array is not 2-D/3-D/4-D, or if a 3-D input is ambiguous and does not look like (H, W, C).

TorchToNumpyConverter

Convert (batch_size, channels, height, width) torch tensors to NumPy HWC/BHWC arrays.

Single-image batches (1, channels, height, width) are squeezed to (height, width, channels) for convenience. Multi-image batches produce (batch_size, height, width, channels).

convert

convert(tensor: Tensor) -> NDArray[Any]

Convert a torch tensor (BCHW) to a NumPy array (HWC or BHWC).

Parameters:

Name	Type	Description	Default
tensor	Tensor	torch.Tensor of shape (batch_size, channels, height, width).	required

Returns:

Type	Description
NDArray[Any]	NumPy ndarray of shape (height, width, channels) when batch_size == 1, or
NDArray[Any]	(batch_size, height, width, channels) otherwise. Dtype is preserved (typically float32).

Raises:

Type	Description
ValueError	If tensor is not a 4-D (batch_size, channels, height, width) tensor.