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

With auxiliary data_keys, an unknown or unclassified spatial transform is rejected before any segment executes. Image-only calls may use native passthrough, so review every Unknown ... SPATIAL_KERNEL barrier warning. Read Auxiliary targets 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.

For detector batches whose per-image annotations are ragged, use the package-root augment_detection_batch boundary rather than padding targets in application code.

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 the actual forward pixel-centre matrix from the most recent call's last supported matrix-producing segment. Fused affine/projective, exact D4/flip/quarter-turn, and direct deterministic letterbox segments publish it; it is 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, generator: Generator | None = None, fill: FillValue | None = None, keypoint_flip_index: Sequence[int] | None = None, mask_fill: MaskFillValue = 0, **backend_kwargs: object)

Bases: FactoriesMixin, IntrospectionMixin, GeneratorPicklingMixin, 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", "rboxes", "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. Throughput depends on the batch, device, and workload; 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. Each active image is filtered separately, preventing a batch neighbour from changing its filter support. A no-op when the scale is safe or the flag is off, keeping the default output bit-identical; antialias=True requires the optional kornia dependency at construction.	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'
mask_fill	MaskFillValue	Scalar auxiliary-mask border value, independent of image fill. Defaults to zero; use a dtype-compatible ignore label such as 255 or -1. Bilinear soft masks interpolate toward this value outside the source canvas.	0
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
keypoint_flip_index	Sequence[int] | None	Keypoint pair permutation applied to the "keypoints" target whenever the composed transform reverses orientation: slot i takes its value from slot flip_index[i]. A mirrored image has its left and right anatomy swapped, so the identity of each slot has to follow the coordinates. Which slots pair with which is dataset schema and stays with the caller; this package only decides when to apply it, and decides that from the sign of the composed matrix's determinant rather than from the presence of a flip transform — after fusion the mirror is no longer a discrete op, and two mirrors compose back to a rotation that must not swap. None (default) leaves the keypoint axis in input order.	None
fill	FillValue | None	Constant value written outside the source canvas of every resampling warp, in the image's own value range (114 for a uint8 grey, 114 / 255 for the same grey as float). A scalar fills every channel; a sequence fills one channel each. None (default) keeps the plain zero border, so existing pipelines are unchanged. The fill applies to the image only; mask_fill independently controls auxiliary-mask borders. Requires padding_mode="zeros": the other modes have no constant to replace, and the combination raises.	None
generator	Generator | None	Caller-owned torch.Generator driving every pipeline-owned draw. None (default) keeps the global torch stream, leaving existing pipelines bit-for-bit unchanged. Backend transform objects sample through their own libraries (Kornia's samplers, Albumentations' py_random and np.random, TorchVision's parameter hooks), none of which accept a torch.Generator; passing one together with backend transforms raises rather than silently falling back to the global stream. Use :meth:from_params (or backend="native") for a fully caller-seeded pipeline.	None
**backend_kwargs	object	Reserved for backend-specific options (currently unused).	{}

Initialize FusedCompose.

forward

forward(*args: Tensor, return_matrix: bool = False, **kwargs: object) -> 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; "rboxes" as (B, N, 5) float32 (cx, cy, w, h, theta), theta in radians, returned un-canonicalized (this package imposes no long-edge or angle-range convention -- see :func:~fuse_augmentations.targets.transform_rboxes).	()
return_matrix	bool	When True, return the output and its last supported geometric pixel-centre matrix.	False
**kwargs	object	Declared multi-target keyword inputs. Tensor keyword calls return a dict and retain the same public module hooks as positional tensor calls.	{}

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 full-precision pixel-centre matrix from the last supported
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

Only known coordinate-preserving passthrough transforms may run with auxiliary targets. Unclassified or coordinate-changing passthrough is refused before any segment executes, since it cannot route the targets.

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.

Ragged detection adapter

augment_detection_batch

augment_detection_batch(pipeline: FusedCompose, images: Tensor, targets: Sequence[Mapping[str, Tensor]], *, min_size: float = 0.0, min_visibility: float = 0.0) -> tuple[Tensor, list[dict[str, Tensor]]]

Augment a detector batch while keeping supported ragged target fields aligned.

The pipeline must declare exactly ["input", "bbox_xyxy"]. Each target requires floating boxes of shape (N, 4) and int64 labels of shape (N,) on the image device; area, iscrowd, and scalar image_id are optional. The returned list contains new dictionaries and tensors: boxes are clipped to the output pixel-edge canvas, positive-area instances meeting the explicit thresholds survive, labels/iscrowd follow the same mask, and supplied area is recomputed from the clipped boxes.

Parameters:

Name	Type	Description	Default
pipeline	FusedCompose	A :class:FusedCompose declaring input and bbox_xyxy.	required
images	Tensor	Floating image batch with shape (B, C, H, W).	required
targets	Sequence[Mapping[str, Tensor]]	One supported detector target mapping per image.	required
min_size	float	Minimum clipped width and height in pixels.	0.0
min_visibility	float	Minimum clipped/original box-area ratio in [0, 1].	0.0

Returns:

Type	Description
Tensor	The augmented image batch and one ragged, detector-ready target mapping
list[dict[str, Tensor]]	per input image.

Raises:

Type	Description
TypeError	If the pipeline, images, targets, or thresholds have an incompatible type.
ValueError	If target schemas, devices, dtypes, fields, or the pipeline declaration are incompatible with this boundary.