Augmentation Pipeline

pipeline

Composable augmentation pipeline for ML training.

Designed to be called inside a PyTorch Dataset's __getitem__ or any data loading loop.

AugmentationStep

Bases: NamedTuple

One configured pipeline step.

A :class:tuple subclass, so a step still unpacks as (fn, prob, kwargs) and indexes as step[2] — the named fields are what reads at the call site when introspecting a pipeline (pipeline.augmentations[0].kwargs["angle"]).

AugmentationPipeline(augmentations: list[tuple[Callable, float, dict]], cache_quats: bool = True, *, representation: str | None = None, euler_orders: list[str] | None = None)

Composable sequence of augmentations with per-step probabilities.

Each augmentation is a tuple of (fn, probability, kwargs) where fn has signature fn(arrays, **kwargs) -> MotionArrays — :class:~pybvh_ml.MotionArrays in, a new one out. Custom steps read arrays.root_pos / arrays.joint_rot and typically return arrays.replace(joint_rot=...).

Kwargs values may be callables of the form lambda rng: value, which are resolved at each invocation using the pipeline's rng. This enables random parameter sampling per sample (e.g., random rotation angles). Pass return_params=True to get back what each call actually drew — which steps fired and with which sampled values (see __call__).

The pipeline automatically forwards its rng to augmentation functions that accept an rng parameter (detected via signature inspection). This ensures reproducibility without requiring "rng": lambda rng: rng in kwargs.

Parameters:

Name	Type	Description	Default
augmentations	list of (callable, float, dict)	Each entry is (fn, probability, kwargs). probability is in [0, 1]; the augmentation is applied when a uniform draw is below this threshold. Entries are stored as :class:AugmentationStep named tuples, so pipeline.augmentations[i].kwargs and the plain pipeline.augmentations[i][2] both work.	required
representation	str	Pipeline-level default for the representation kwarg. A pipeline is homogeneous in practice, and repeating the token on every step is where one step in five ends up disagreeing with the rest. Steps that declare their own representation keep it — the default only fills in for those that don't (and only for functions that name the parameter; a **kwargs catch-all does not count, so custom steps taking neither are called with exactly their own kwargs). Also satisfies the cache_quats=True requirement that something declare what joint_data is in. Built-in steps must agree on the resulting representation: each step's output is the next one's input, so two built-ins declaring different tokens with nothing between them to convert raises at construction. A custom step in between lifts the restriction, since it may legitimately convert.	None
euler_orders	list of str	Pipeline-level default for the euler_orders kwarg, with the same per-step-override semantics. Needed only when representation="euler".	None
cache_quats	bool	Share a quaternion cache across pybvh-ml's built-in augmentations. Functions like :func:add_joint_rotation_noise and :func:speed_perturbation_arrays always operate in quaternion space internally; when a pipeline strings several of them together with representation="axisangle" or "euler", this flag eliminates all but the first and last conversion — typically a 2–3× speedup on non-6d pipelines, 1.5× on 6d. User-defined augmentations not registered in the internal staging table are supported transparently: the cache is flushed around them and they receive joint_rot in their declared representation kwarg — or, when they declare none, in the pipeline's current declared representation (the most recent step carrying a representation kwarg), exactly as on the cache_quats=False path. Set to False for historical bit-exact behavior.	True

Notes

Composition order matters. Steps run left-to-right on the output of the previous step. The mathematically interesting interactions to be aware of:

• Mirror vs. vertical rotation. mirror_* reflects the lateral axis; rotate_*_vertical rotates around the up axis. The two commute up to a sign flip on the rotation angle (mirror ∘ rotate(θ) == rotate(-θ) ∘ mirror). Either order is correct, but if you stack them with probabilities and expect the resulting distribution to be symmetric, keep that sign flip in mind.
• Speed perturbation changes F. speed_perturbation_arrays resamples time, so downstream steps receive arrays with a different F. Frame-count-sensitive steps (e.g. dropout_arrays with a fixed keep-mask) should run before speed perturbation or be written to tolerate variable lengths.
• Noise + re-augmentation. add_joint_rotation_noise perturbs rotations in place (via quaternion space); following it with a second rotation-space step is fine, but a subsequent deterministic check (e.g. equality to the input) will naturally fail.
• A re-derivation discards the position stream's own history. On a sample carrying positions, add_joint_rotation_noise replaces them with forward kinematics of the noised rotations rather than transforming the incoming ones (see :func:~pybvh_ml.handles_streams). So on a rig with asymmetric rest offsets, [mirror, add_joint_rotation_noise] ends with FK of locally-mirrored rotations — throwing away the world-exact reflection the position stream held — while [add_joint_rotation_noise, mirror] keeps it. Both are defensible and neither is a bug, but the two do not produce the same positions.

Every step must handle every stream the sample carries. The check runs once at __call__ entry, for every configured step, before any of them fires — a p=0.1 step with the wrong stream support or a missing fk_topology would otherwise raise on one sample in ten. A custom step that declares nothing is assumed to handle {"root_pos", "joint_rot"}; decorate it with :func:~pybvh_ml.handles_streams once it transforms positions too.

Examples:

>>> from pybvh_ml.augmentation import rotate_vertical, mirror
>>> pipeline = AugmentationPipeline([
...     (rotate_vertical, 1.0, {
...         "angle": lambda rng: rng.uniform(-np.pi, np.pi),
...         "up_axis": bvh.world_up,
...     }),
...     (mirror, 0.5, {
...         "lr_joint_pairs": pairs,
...         "lateral_axis": "+x",
...     }),
... ], representation="6d")
>>> out = pipeline(MotionArrays(root_pos=root_pos,
...                                 joint_rot=joint_rot6d), rng=rng)
>>> out.joint_rot.shape
(120, 31, 6)

standard(skeleton_info: dict, *, representation: str | None = '6d', up_axis: str = '+y', lateral_axis: str = '+x', rotate_angle_range: tuple[float, float] | None = (-np.pi, np.pi), mirror_prob: float = 0.5, noise_sigma: float | None = np.radians(1.0), position_noise_sigma: float | None = None, position_space: str | None = None, speed_factor_range: tuple[float, float] | None = (0.8, 1.2), degrees: bool = False, cache_quats: bool = True) -> 'AugmentationPipeline' classmethod

Build the canonical rotate + mirror + noise + speed pipeline.

Convenience factory that wires the four common augmentation steps from a skeleton_info dict (as returned by :func:pybvh_ml.skeleton.get_skeleton_info or :func:pybvh_ml.preprocessing.load_preprocessed) so callers don't reassemble the boilerplate for every project.

Each step is optional: pass None (or 0 for mirror_prob) to skip it. For anything beyond what these kwargs expose, build the pipeline directly with the (fn, prob, kwargs) constructor — this factory is the opinionated common case, not a wrapper around every knob.

Parameters:

Name	Type	Description	Default
skeleton_info	dict	Supplies lr_pairs / node_lr_pairs (required for mirror), euler_orders (required when representation="euler"), and — for a positions-carrying dataset — fk_topology, world_up and position_space.	required
representation	str or None	Rotation representation threaded through every step. One of "quat", "6d", "axisangle", "rotmat", "euler". None builds a positions-only pipeline and skips the rotation-noise step, which would otherwise be configured by default and refuse every sample: noising rotations is meaningless on a clip that has none. (A direct :func:~pybvh_ml.add_joint_rotation_noise call on such a sample still raises — a factory declining to configure a meaningless step and a function refusing a meaningless call are different questions.)	'6d'
up_axis	str	Signed-axis strings (e.g. "+y", "+x"). The defaults assume a +y-up, +x-lateral skeleton; set from bvh.world_up and the dataset's lateral convention otherwise.	'+y'
lateral_axis	str	Signed-axis strings (e.g. "+y", "+x"). The defaults assume a +y-up, +x-lateral skeleton; set from bvh.world_up and the dataset's lateral convention otherwise.	'+y'
rotate_angle_range	(float, float) or None	Random yaw range in radians (degrees when degrees=True); None skips rotation.	(-pi, pi)
degrees	bool	Interpret rotate_angle_range and noise_sigma in degrees. Default False (radians). One flag serves both because both are angles — which is exactly why root-position noise is not a knob on this factory: its sigma is a length, and a single flag could not have covered it. Use :func:~pybvh_ml.add_root_position_noise as an explicit step for that.	False
mirror_prob	float	Probability of left/right mirror. 0 skips it. Silently skipped when the pair list this configuration needs is empty (skeleton_info["lr_pairs"] for the joint-space streams, node_lr_pairs for node_pos) — no pairs were detected on this skeleton.	0.5
noise_sigma	float or None	Per-joint rotation noise standard deviation in radians (default one degree); None skips noise. On a dataset carrying positions this step also refreshes them by forward kinematics, so the factory wires fk_topology and world_up from skeleton_info.	radians(1.0)
position_noise_sigma	float or None	Per-vertex keypoint jitter, in the data's positional units; None (default) skips it. Joint-space and node-space jitter are different functions with different stream declarations, and the pipeline is built before any sample is seen, so the index space is resolved here — from position_space when given, else skeleton_info["position_space"]. Wiring one unconditionally would make the pipeline refuse every sample of a dataset stored in the other space.	None
position_space	('joint', 'node')	Explicit override for that resolution. None (default) reads skeleton_info["position_space"].	"joint"
speed_factor_range	(float, float) or None	Random speed factor range; None skips speed perturbation. Runs last because it changes F.	(0.8, 1.2)
cache_quats	bool	Passed through to the pipeline constructor.	True

__call__(arrays: MotionArrays, *, rng: np.random.Generator | None = None, return_params: bool = False) -> MotionArrays | tuple[MotionArrays, list[dict]]

Apply augmentations with their configured probabilities.

Parameters:

Name	Type	Description	Default
arrays	MotionArrays	The clip to augment. Positional because it is a distinct type — every other argument stays keyword-only.	required
rng	numpy Generator	Random number generator. Defaults to a new unseeded one.	None
return_params	bool	Also return what this call drew (see params below). Purely additive: the random stream is untouched, so a given rng produces identical arrays either way. This is the only thing that changes the return arity.	False

Returns:

Name	Type	Description
	MotionArrays	Always freshly allocated — the outputs never share storage with the input arrays, even when no augmentation fires. (The container's fields are read-only either way; take np.array(out.joint_rot) for a writable working array.) Each stream comes back in the dtype it went in as, with the math done in float64 regardless: the dtype must not depend on which steps this sample's probability draws happened to fire, and both call paths have to agree bit for bit.
params	list of dict	Only when return_params=True. One record per configured step, in pipeline order (index-aligned with the augmentations list), each shaped {"name": str, "applied": bool, "params": dict}. applied is the outcome of the step's probability draw. params holds the kwargs this call sampled — those whose spec is a callable — resolved to the values the augmentation received. Static kwargs are pipeline configuration, readable from augmentations (or repr), and rng is machinery rather than a parameter; neither appears here. A step that did not fire reports {}: its callables are never invoked, which is what keeps the random stream identical.

Examples:

>>> out, steps = pipeline(arrays, rng=rng, return_params=True)
>>> [(s["name"], s["applied"]) for s in steps]
[('rotate_vertical', True), ('mirror', False)]
>>> steps[0]["params"]["angle"]
1.8721...