lightning_pose.data

lightning_pose.data.augmentations Module

Functions to build augmentation pipeline.

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

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

`project_camera_pairs_to_3d`(points, ...)	Project 2D keypoints from each pair of cameras into 3D world space.
`project_3d_to_2d`(points_3d, intrinsics, ...)	Project 3D keypoints to 2D using camera parameters.

Inherit Anipose camera group and add new non-jitted triangulation method for dataloaders.

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

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

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

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

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

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

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

`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.

Classes to streamline data typechecking.

`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.

Helper class to extract labeled data from a data module.

Helper class to extract all data from a data module.

Factory functions to build data pipeline components from config.

`get_imgaug_transform`(cfg)	Create simple and flexible data transform pipeline that augments images and keypoints.
`get_dataset`(cfg, data_dir, imgaug_transform)	Create a dataset that contains labeled data.
`get_data_module`(cfg, dataset[, video_dir])	Create a data module that splits a dataset into train/val/test iterators.

Dataset/data module utilities.

`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)	Transform keypoints from normalized coordinates to bbox coordinates
`convert_bbox_coords`(batch_dict, ...[, in_place])	Transform keypoints from bbox coordinates to absolute frame coordinates.
`convert_original_to_model_coords`(batch_dict, ...)	Transform keypoints from original frame coordinates to model input coordinates.
`original_to_model`(keypoints, bbox, ...)	Convert keypoints from original image coordinates to model input coordinates.