fuse-augmentations

Fuse consecutive geometric augmentation transforms into a single interpolation pass -- fewer warps, better image quality.

Summary: fuse-augmentations is a framework-agnostic library that automatically groups consecutive fusible geometric transforms in your augmentation pipeline, then fuses their matrices into a single composed transform applied via one interpolation pass. Operations that are not yet fusible (blur, color jitter) pass through unchanged. Drop-in replacement for Kornia's AugmentationSequential, TorchVision, and Albumentations compose classes.

💡 Motivation

People pick specific transforms -- RandomRotation, RandomHorizontalFlip, RandomScale -- because they want intuitive, independent control over each one. That is why nobody just uses a single monolithic RandomAffine for everything: it does not let you set different probabilities per parameter (e.g. flip with p=0.5, rotation with p=0.8, scale with p=0.7, each drawn independently).

The problem is that chaining these individual transforms applies a separate interpolation for each one, compounding quality loss across your pipeline.

fuse-augmentations gives you the best of both worlds. You keep writing your pipeline with individual, independently-controlled transforms, and Compose is a drop-in replacement for your existing backend's compose class (AugmentationSequential, transforms.Compose, etc.) -- no pipeline rewrite needed. Under the hood, the library groups consecutive fusible geometric transforms and fuses their matrices, applying a single interpolation pass. The fusion is an implementation detail that gives you quality improvement for free.

🔍 Overview

Given a pipeline of transforms, fuse-augmentations performs two steps:

1. Grouping: consecutive fusible geometric transforms are identified and collected into segments. Operations that are not yet fusible (Gaussian blur, color jitter, normalization) act as natural segment boundaries and pass through via their native backend unchanged.
2. Fusing: within each segment, the individual affine (or projective) matrices are composed mathematically -- M_composed = M_n @ ... @ M_2 @ M_1 -- and a single interpolation pass applies the entire group.

A pipeline of three affine transforms saves two interpolation passes. At training time, with thousands of images and many augmentation steps, this translates to measurably better effective resolution in your augmented dataset.

✨ Features

• Automatic fusion of (consecutive) geometric transforms -- no manual configuration needed. Reordering with quality preservation is coming soon, so non-consecutive runs will also be fusible.
• All affine transforms from each supported backend (Kornia, TorchVision, Albumentations) are mapped and fusible -- not just a subset.
• Per-sample randomness: independent probability draws per image in the batch.
• Auxiliary target support: masks, bounding boxes (xyxy and xywh), and keypoints warped by the same composed matrix.
• Multi-backend: Kornia, TorchVision, and Albumentations transforms in the same pipeline.
• Backend-free mode: construct a pipeline from numeric parameter ranges with no framework imports.
• Reorder policy: optionally bubble color ops after geometric runs to extend fusion windows.
• Fusion introspection: inspect fusion_plan, n_warps_saved, and transform_matrix after each forward pass.
• Projective (perspective) transform fusion via full 3x3 homography matrices.
• Roadmap: future versions will allow passing transforms as meta-configuration and receiving output in any supported backend -- e.g. build a pipeline with Albumentations ops and switch to Kornia output when changing accelerator, without rewriting the pipeline.

📦 Installation

pip install fuse-augmentations

Backend extras are optional -- install only what your pipeline uses:

pip install "fuse-augmentations[kornia]"       # Kornia transforms
pip install "fuse-augmentations[torchvision]"  # TorchVision transforms
pip install "fuse-augmentations[albumentations]"  # Albumentations transforms
pip install "fuse-augmentations[all]"          # All backends

Requirements: Python 3.10+, PyTorch >= 2.2.

🚀 Quick Start

import torch
import kornia.augmentation as K
from fuse_aug import Compose  # or: from fuse_augmentations import Compose

pipe = Compose(
    [
        K.RandomRotation(degrees=30, p=0.8),
        K.RandomHorizontalFlip(p=0.5),
        K.RandomAffine(degrees=0, scale=(0.8, 1.2), p=0.7),
    ]
)

image = torch.rand(4, 3, 256, 256)  # (B, C, H, W)
out = pipe(image)  # one interpolation pass instead of three

print(pipe.fusion_plan)
# fused(RandomRotation, RandomHorizontalFlip, RandomAffine)

print(pipe.n_warps_saved)
# 2

The short import fuse_aug is a canonical alias for fuse_augmentations -- both expose the same public API. All affine transforms from each backend are supported, not just the ones shown in this example.

⚙️ How Fusion Works

Given a pipeline [Rotate, Scale, HFlip, GaussianBlur, Rotate]:

1. Grouping: [Rotate, Scale, HFlip] are consecutive geometric transforms and are collected into one segment. GaussianBlur is a spatial-kernel operation that is not yet fusible, so it acts as a segment boundary. The trailing Rotate forms its own segment.
2. Fusing: the first segment's affine matrices are composed: M = M_hflip @ M_scale @ M_rot. One interpolation pass applies all three. The trailing Rotate segment applies its own single pass.

All matrices are (B, 3, 3) homogeneous in pixel coordinates with align_corners=True. To apply the interpolation, the composed forward matrix is inverted once to yield backward (sampling) grid coordinates.

For flip-only chains, fuse-augmentations uses an ExactAffineSegment that applies tensor.flip directly -- zero interpolation error.

📖 API Reference

Core

Enums

Auxiliary Target Functions

🎯 Auxiliary Targets

Pass data_keys to route masks, boxes, or keypoints through the same fused transform:

from fuse_aug import Compose
import kornia.augmentation as K

pipe = Compose(
    [K.RandomRotation(degrees=30, p=0.8), K.RandomHorizontalFlip(p=0.5)],
    data_keys=["input", "mask", "bbox_xyxy", "keypoints"],
)

img_out, mask_out, bboxes_out, kpts_out = pipe(image, mask, bboxes, keypoints)
# mask warped with nearest-neighbour -- integer class labels preserved
# bboxes AABB-wrapped after rotation
# keypoints transformed exactly via homogeneous matrix

Supported data_keys values:

🔌 Backend-Free Pipelines

No Kornia or TorchVision import needed:

from fuse_aug import Compose, InterpolationMode, PaddingMode

pipe = Compose.from_params(
    rotation=(-30, 30),
    scale=(0.8, 1.2),
    hflip_p=0.5,
    vflip_p=0.3,
    interpolation=InterpolationMode.BICUBIC,
    padding_mode=PaddingMode.REFLECTION,
)

out = pipe(image)

from_params accepts: rotation, scale, scale_x, scale_y, shear_x, shear_y, translate_x, translate_y, hflip_p, vflip_p, interpolation, padding_mode, reorder, data_keys.

Future: from_params will accept a list of per-transform dicts/dataclasses, each covering native parameters and per-transform probability -- enabling complete parity with passing native backend transforms directly.

🔗 Multi-Backend Pipelines

Kornia, TorchVision, and Albumentations transforms can be mixed in the same Compose:

import torchvision.transforms.v2 as T
import kornia.augmentation as K
from fuse_aug import Compose

pipe = Compose(
    [
        K.RandomRotation(degrees=15),
        T.RandomHorizontalFlip(),
        K.ColorJitter(brightness=0.3),
    ]
)

out = pipe(image)
# fused(RandomRotation, RandomHorizontalFlip) -> passthrough(ColorJitter)

Each transform is resolved to the correct adapter at construction time. Framework-specific behavior (parameter sampling, matrix building, passthrough for operations not yet fusible) is handled by KorniaAdapter, TorchVisionAdapter, or AlbumentationsAdapter.

🔀 Reorder Policy

When a color operation sits between two geometric transforms, fusion is broken by default. ReorderPolicy.POINTWISE bubbles color ops to the end of each geometric stretch, extending the fusion window:

from fuse_aug import Compose, ReorderPolicy
import kornia.augmentation as K

pipe = Compose(
    [
        K.RandomRotation(degrees=15, p=0.8),
        K.ColorJitter(brightness=0.3, p=0.5),  # POINTWISE -- would break fusion
        K.RandomHorizontalFlip(p=0.5),
    ],
    reorder=ReorderPolicy.POINTWISE,
)

print(pipe.fusion_plan)
# fused(RandomRotation, RandomHorizontalFlip) -> passthrough(ColorJitter)

ReorderPolicy.NONE (default): preserves declared order, merges consecutive fusible transforms.

ReorderPolicy.AGGRESSIVE (coming soon): moves all pixel-wise ops (color jitter, normalization) to the beginning of each geometric stretch in their original relative order, then fuses the remaining geometric transforms into one segment -- extending the fusion window maximally.

🔬 Fusion Introspection

After any forward pass:

out = pipe(image)

print(pipe.fusion_plan)
# fused(RandomRotation, RandomAffine) -> passthrough(RandomGaussianBlur) -> fused(RandomHorizontalFlip)

print(pipe.n_warps_saved)
# 1  -- one interpolation pass saved

M = pipe.transform_matrix  # (B, 3, 3) composed forward matrix

transform_matrix gives the composed forward affine matrix for each sample in the batch. Use it to transform stored coordinates that were not passed as data_keys.

Future: fusion_plan will be exposed as a sequence of reusable, reversible transform descriptors (dicts/dataclasses), making pipelines inspectable, serialisable, and reproducible.

⚠️ Limitations

• Pixel-wise ops (ColorJitter, Normalize) are not yet fusible -- they are single-pixel operations and currently act as passthrough.
• Spatial-kernel ops (GaussianBlur, Sharpen) act as fusion barriers; transforms on either side of a barrier form separate segments. These are not yet fusible.
• Padding mode is segment-level: all transforms in a fused run share the same padding mode (the highest-quality setting among them).
• Gradients: image transforms are differentiable; mask sampling (mode='nearest') is not.

🤝 Contributing

Bug fixes are always welcome -- just open a pull request on GitHub. For new features or bigger ideas, open an issue first so we can discuss the direction -- all suggestions are genuinely appreciated.

📄 License

Apache-2.0. Copyright (c) 2025-2026 Jiri Borovec.