lightning_pose.data

lightning_pose.data.augmentations Module

Functions to build augmentation pipeline.

Functions

imgaug_transform(params_dict)

Create simple and flexible data transform pipeline that augments images and keypoints.

lightning_pose.data.dali Module

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

Functions

video_pipe(filenames[, resize_dims, ...])

Generic video reader pipeline that loads videos, resizes, augments, and normalizes.

Classes

LitDaliWrapper	wrapper around a DALI pipeline to get batches for ptl.
PrepareDALI	All the DALI stuff in one place.

lightning_pose.data.datamodules Module

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

Classes

BaseDataModule	Splits a labeled dataset into train, val, and test data loaders.
UnlabeledDataModule	Data module that contains labeled and unlabled data loaders.

lightning_pose.data.datasets Module

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

Classes

BaseTrackingDataset	Base dataset that contains images and keypoints as (x, y) pairs.
HeatmapDataset	Heatmap dataset that contains the images and keypoints in 2D arrays.
MultiviewHeatmapDataset	Heatmap dataset that contains the images and keypoints in 2D arrays from all the cameras.

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.utils Module

Dataset/data module utilities.

Functions

split_sizes_from_probabilities(total_number, ...)	Returns the number of examples for train, val and test given split probs.
clean_any_nans(data, dim)	Remove samples from a data array that contain nans.
count_frames(video_file)	Simple function to count the number of frames in a video.
compute_num_train_frames(len_train_dataset)	Quickly compute number of training frames for a given dataset.
generate_heatmaps(keypoints, height, width, ...)	Generate 2D Gaussian heatmaps from mean and sigma.
evaluate_heatmaps_at_location(heatmaps, locs)	Evaluate 4D heatmaps using a 3D location tensor (last dim is x, y coords).
undo_affine_transform(keypoints, transform)	Undo an affine transform given a tensor of keypoints and the tranform matrix.
undo_affine_transform_batch(...[, is_multiview])	Potentially undo an affine transform given a tensor of keypoints and the tranform matrix.
normalized_to_bbox(keypoints, bbox)
convert_bbox_coords(batch_dict, ...)	Transform keypoints from bbox coordinates to absolute frame coordinates.

Classes

DataExtractor

Helper class to extract all data from a data module.