lightning_pose.data

lightning_pose.data.augmentations Module

Functions to build augmentation pipelines from imgaug transform names or preset strings.

imgaug_transform() builds an iaa.Sequential pipeline from a params dict (transform name -> {p, args, kwargs}); see its docstring for the dict format.

expand_imgaug_str_to_dict() expands a shorthand preset string ("dlc", "dlc-mv", etc.) into that same params-dict format.

Adding a new preset: add its string to _allowed_imgaug_strs and a branch in expand_imgaug_str_to_dict() that builds params_dict the same way the existing "dlc" branches do – one params_dict[<TransformName>] = {"p": ..., "kwargs": {...}} entry per imgaug transform, using if not params.endswith("mv") to skip transforms that don’t make sense for multiview data (e.g. per-view crop/rotation).

lightning_pose.data.cameras Module

Camera geometry utilities for multi-view 2D-to-3D projection and triangulation.

lightning_pose.data.video.dali Module

Data pipelines based on efficient video reading by nvidia dali package.

Import warning

nvidia-dali-cuda110 is not installed on CPU-only machines (macOS, GPU-less Linux), so importing this module at the top level of any other module will raise ImportError on those platforms and break import lightning_pose entirely. Always import from this module lazily, inside the function or method body that uses it:

def my_function(...):
    from lightning_pose.data.video.dali import PrepareDALI  # lazy: avoids cpu-only ImportError
    ...

Architecture overview

PrepareDALI is the entry point. Its __init__ validates inputs and pre-computes pipeline arguments for all four combinations of stage ("train", "predict") and model type ("base", "context"). Calling the instance (__call__) builds the DALI pipe for the requested combination and returns a ready-to-iterate LitDaliWrapper.

LitDaliWrapper extends DALIGenericIterator and converts raw DALI output into typed UnlabeledBatchDict or MultiviewUnlabeledBatchDict instances on every __next__.

Two prediction modes

Standard mode (no bbox_df):

DALI resizes frames to resize_dims on the GPU. The bbox field of each returned batch covers the full frame (x=0, y=0, h=H, w=W).

Bbox-crop mode (bbox_df supplied to PrepareDALI):

DALI delivers full-resolution frames (resize_dims=None in the predict pipe) so that LitDaliWrapper._apply_bbox_crop can crop each frame to its per-frame bounding box and resize to the original resize_dims using torch.nn.functional.interpolate. The bbox field of each batch contains the actual crop coordinates, so downstream code can remap predictions back to the original coordinate space.

lightning_pose.data.video.factory Module

Reader-selection dispatch: turns a backend name (or None) into a built dataloader.

Single entry point: build_video_reader(). Resolves which backend to use – either the caller’s explicit choice, validated against this machine/video, or an auto-select fallback chain (pynvvc -> dali -> opencv, richest/fastest first, most portable last) – then constructs and returns that backend’s ready-to-iterate loader.

Adding a new video reader

Every backend module (dali.py, pynvvc.py, opencv.py, …) implements the same contract. To add one:

1. Two-phase prepare class: Prepare<Name>.__init__(model_type, filenames, resize_dims, dali_config, bbox_df=None, ...) validates inputs and precomputes windowing parameters (raise early, before any GPU/decoder allocation); __call__() builds and returns a Lit<Name>Wrapper. (The dali_config parameter name is a pre-existing wart – PreparePynvvc already reuses cfg.dali’s windowing section rather than inventing a per-backend config block, since the windowing semantics are decoder-agnostic; new backends should follow the same precedent rather than fixing it in isolation.)

2. Iterator wrapper class: Lit<Name>Wrapper is iterable, defines __len__, and yields UnlabeledBatchDict / MultiviewUnlabeledBatchDict – FCHW float frames, ImageNet-normalized, plus transforms and bbox. Match LitDaliWrapper/LitPynvvcWrapper’s shape exactly so downstream code (PredictionHandler, trainer.predict) doesn’t need to know which backend produced a batch.

3. Availability probe (optional): is_<name>_available(video_path, ...) -> bool, only needed if the backend can be installed but still unusable for a given machine/video (wrong hardware/driver generation, as with pynvvc). Skip it for a backend that’s simply present or absent as a plain package dependency.

4. Import discipline: if the backend’s package is platform-gated or proprietary (not a guaranteed cross-platform install), every import of it must be lazy – inside the function or method body that uses it, never at module top level (see the # lazy: avoids ImportError on cpu-only installs convention in dali.py/pynvvc.py). A top-level import of such a package anywhere in the package or tests will break import lightning_pose on machines that don’t have it. If the package is an unconditional, cross-platform dependency already declared in pyproject.toml (like opencv-python-headless), top-level imports are fine.

5. Wire it in:

   • add the name to _Reader in lightning_pose/utils/inference_types.py (the single source for the CLI’s --reader choices and API type hints);

   • add a branch – and, if step 3 applies, a fallback rung – in build_video_reader() (this file);

   • update the --reader help text in lightning_pose/cli/commands/predict.py;

   • add this module to the package map in lightning_pose.data.video’s docstring.

6. Tests: add tests/data/video/test_<name>.py mirroring a sibling reader’s test module one-for-one, plus a case in tests/data/video/test_factory.py covering its fallback/validation branch; mark any test that needs real hardware with @pytest.mark.gpu so the CPU CI workflow (pytest -m "not gpu") still exercises everything else.

lightning_pose.data.video.opencv Module

Data pipeline for video prediction based on OpenCV (cv2.VideoCapture).

Import discipline

Unlike dali.py/pynvvc.py, cv2 (opencv-python-headless) is an unconditional, cross-platform dependency already declared in pyproject.toml and already imported at module top level elsewhere in this package (e.g. lightning_pose.data.cameras, lightning_pose.utils.predictions). So, unlike those two backends, this module imports cv2 eagerly; no lazy-import discipline is needed here (see lightning_pose.data.video’s package docstring for when lazy imports are required).

Architecture overview

Mirrors lightning_pose.data.video.pynvvc’s two-phase construction: PrepareOpenCV.__init__ validates inputs and precomputes windowing parameters; calling the instance (__call__) builds and returns a ready-to-iterate LitOpenCVWrapper.

Predict-only, like pynvvc: this backend never runs random_shuffle or the imgaug augmentation pipeline, since litpose predict already runs with both off. Training continues to go through DALI exclusively (lightning_pose.data.video.dali).

Sequential decode, not seek-based

Context (MHCRNN) models read overlapping windows: consecutive sequence_length-frame windows advance by step = sequence_length - 4, so the last 4 frames of window i are the same physical frames as the first 4 of window i+1. cv2.VideoCapture’s CAP_PROP_POS_FRAMES seeking is not reliably frame-exact on containers with B-frames or a variable/estimated frame rate – a seek-per-window implementation risks silently misaligning that overlap: no crash, no shape change, just a temporally-shifted context stack feeding a plausible-but-wrong prediction.

To avoid that, LitOpenCVWrapper never seeks: it opens each view’s cv2.VideoCapture once and reads strictly forward via cap.read(), caching the trailing 4 frames of each window in self._tail to prepend to the next one. Every physical frame is decoded exactly once (aside from the cached tail).

lightning_pose.data.video.pynvvc Module

Data pipeline for video prediction based on PyNvVideoCodec (direct NVDEC access).

Import warning

Same discipline as lightning_pose.data.video.dali: pynvvideocodec is an unconditional but platform-gated dependency (Linux x86_64 only), so importing this module at the top level of any other module could raise ImportError/ModuleNotFoundError on other platforms. Always import from this module lazily, inside the function or method body that uses it:

def my_function(...):
    from lightning_pose.data.video.pynvvc import PreparePynvvc  # lazy
    ...

The PyNvVideoCodec package itself is additionally deferred to inside LitPynvvcWrapper.__init__ and is_pynvvc_available specifically, rather than this module’s top level, since import PyNvVideoCodec succeeds regardless of GPU generation or driver age (Turing/Ampere/Ada/Hopper/Blackwell only, driver >= 530.41.03 on Linux) – the failure only surfaces when a decoder is actually constructed.

Architecture overview

Mirrors lightning_pose.data.video.dali’s two-phase construction: PreparePynvvc.__init__ validates inputs and precomputes windowing parameters; calling the instance (__call__) builds and returns a ready-to-iterate LitPynvvcWrapper.

Predict-only, unlike PrepareDALI: this backend never runs random_shuffle or the imgaug augmentation pipeline, since litpose predict already runs with both off. Training continues to go through DALI exclusively (lightning_pose.data.video.dali).

CUDA stream synchronization

PyNvVideoCodec.SimpleDecoder writes decoded frames asynchronously on its own internal CUDA stream and hands them back as DLPack buffers. If a consumer (this module’s own resize/normalize step, or eventually the model’s forward pass) reads those buffers on a different stream before the decode is actually finished, the read happens on stale/partial data – silently. No crash, no NaN, just quietly wrong pixels feeding a plausible-looking keypoint prediction. LitPynvvcWrapper.__next__ guards against this with an explicit torch.cuda.current_stream().synchronize() after building each batch’s frame tensors and before returning them. This is the “safe but not necessarily optimal” fix flagged in issue #476 – passing an explicit stream handle into the decoder constructor (if supported) would avoid the sync cost entirely, but that needs checking against the installed PyNvVideoCodec version and is tracked as a follow-up, not done here.

lightning_pose.data.datamodules Module

Data modules split a dataset into train, val, and test modules.

BaseDataModule wraps a labeled dataset and splits it via _setup. UnlabeledDataModule (extends BaseDataModule) additionally builds a DALI video dataloader for semi-supervised training.

Split logic (BaseDataModule._setup): the imgaug pipeline always contains at least one element (a final resize transform appended by BaseTrackingDataset.__init__), so len(imgaug_transform) == 1 means “resize only, no augmentations.” When that’s true and imgaug_hflip is False, all three splits share one underlying dataset object (cheap path). Otherwise three deep-copied datasets are created, with the val/test copies’ pipelines reset to resize-only and imgaug_hflip reset to False. Any new augmentation applied outside the imgaug pipeline (like imgaug_hflip) must be added to this condition and explicitly stripped from val/test in the else branch – see datasets.py’s module docstring for the full “adding a keypoint-affecting augmentation” recipe.

lightning_pose.data.datasets Module

Dataset objects store images, labels, and functions for manipulation.

Three classes, in a hierarchy:

• BaseTrackingDataset – base class. Loads images and (x, y) keypoints, applies the imgaug pipeline, and handles imgaug_hflip. Returns a BaseLabeledExampleDict.

• HeatmapDataset (extends BaseTrackingDataset) – adds compute_heatmap to convert keypoints to (K, H, W) Gaussian heatmap targets, and synthesizes self.visibility from NaN positions when the CSV has no visible column. Returns a HeatmapLabeledExampleDict.

• MultiviewHeatmapDataset – does not inherit from BaseTrackingDataset. Holds a dict[str, HeatmapDataset] at self.dataset (keyed by view name); __getitem__ delegates to each child and stacks results. Shares imgaug_transform/imgaug_hflip attributes with child datasets so the data module can update them in one pass.

__getitem__ branching: both BaseTrackingDataset and HeatmapDataset branch on self.do_context at the top of __getitem__ – the non-context branch loads a single PIL image, the context branch loads a sequence of frames – and all augmentation logic (imgaug pipeline + hflip) is duplicated in both branches, so a change to one must be mirrored in the other.

Adding a keypoint-affecting augmentation: add any per-sample random state at the top of each __getitem__ branch (not outside them, since the two paths are structurally independent); apply it after the imgaug pipeline so keypoints are already in resized coordinate space; for context mode, draw the random decision once and reuse it for every frame in the sequence. See CLAUDE.md’s “Adding an augmentation that affects keypoints” section for the full step-by-step, including the imgaug_hflip-reset requirements in datamodules.py’s _setup and in api/model.py’s _build_datamodule_pred.

lightning_pose.data.datatypes Module

Classes to streamline data typechecking.

Classes

BaseLabeledExampleDict	Return type when calling __getitem__() on BaseTrackingDataset.
HeatmapLabeledExampleDict	Return type when calling __getitem__() on HeatmapTrackingDataset.
MultiviewLabeledExampleDict	Return type when calling __getitem__() on MultiviewDataset.
MultiviewHeatmapLabeledExampleDict	Return type when calling __getitem__() on MultiviewHeatmapDataset.
BaseLabeledBatchDict	Batch type for base labeled data.
HeatmapLabeledBatchDict	Batch type for heatmap labeled data.
MultiviewLabeledBatchDict	Batch type for multiview labeled data.
MultiviewHeatmapLabeledBatchDict	Batch type for multiview heatmap labeled data.
UnlabeledBatchDict	Batch type for unlabeled data.
MultiviewUnlabeledBatchDict	Batch type for multiview unlabeled data.
SemiSupervisedBatchDict	Batch type for base labeled+unlabeled data.
SemiSupervisedHeatmapBatchDict	Batch type for heatmap labeled+unlabeled data.
SemiSupervisedDataLoaderDict	Return type when calling train/val/test_dataloader() on semi-supervised models.

lightning_pose.data.extractor Module

Helper class to extract labeled data from a data module.

lightning_pose.data.factory Module

Factory functions to build data pipeline components from a Hydra config.

Three public functions, typically called in order:

1. get_imgaug_transform() — builds an imgaug augmentation pipeline from cfg.training.imgaug.

2. get_dataset() — wraps the labeled CSV data in the appropriate dataset class (regression, single-view heatmap, or multiview heatmap).

3. get_data_module() — wraps a dataset in a data module that handles train/val/test splitting; selects UnlabeledDataModule for semi-supervised training (adds DALI video loader) or BaseDataModule for supervised-only training.

Adding a new model type (data-side changes only — see models/factory.py for the model-side steps):

1. If the new type can reuse an existing dataset class (e.g. it is a heatmap variant), extend the appropriate elif branch in get_dataset() to match the new cfg.model.model_type string. If it needs a new dataset class, define that class in datasets.py, import it here, and add a new elif branch.

2. If the new type needs a different BaseDataModule subclass, add a branch in get_data_module(); otherwise no change is needed there.

lightning_pose.data.utils Module

Dataset/data module utilities.