MultiviewHeatmapDataset

class lightning_pose.data.datasets.MultiviewHeatmapDataset[source]

Bases: Dataset

Heatmap dataset that contains the images and keypoints in 2D arrays from all the cameras.

Attributes Summary

height

Image height in pixels after resizing.

num_views

Number of camera views in this multiview dataset.

output_shape

Spatial shape of the heatmap output (height, width) after downsampling.

width

Image width in pixels after resizing.

Methods Summary

apply_3d_transforms(data_dict,Β camgroup[,Β ...])

Apply 3D transforms to keypoint and image data (scale, translate).

check_data_images_names(data_length_by_view)

Data checking Each object in self.datasets will have the attribute image_names (i.e. self.datasets['top'].image_names) since each values is a HeatmapDataset.

fusion(datadict)

Merge images, heatmaps, keypoints, and bboxes across views.

Attributes Documentation

height

Image height in pixels after resizing.

num_views

Number of camera views in this multiview dataset.

output_shape

Spatial shape of the heatmap output (height, width) after downsampling.

Returns:

Tuple of (heatmap_height, heatmap_width).

width

Image width in pixels after resizing.

Methods Documentation

apply_3d_transforms(data_dict: dict, camgroup: CameraGroup, scale_params: tuple = (0.8, 1.2), shift_param: float = 0.25) → tuple[source]

Apply 3D transforms to keypoint and image data (scale, translate).

check_data_images_names(data_length_by_view: dict[str, int]) → None[source]

Data checking Each object in self.datasets will have the attribute image_names (i.e. self.datasets[β€˜top’].image_names) since each values is a HeatmapDataset. Include a check to make sure that the image names are the same across all views, so that when it loads element n from each individual view we know these are properly matched.

Parameters:

data_length_by_view – number of labeled frames per view

fusion(datadict: dict) → tuple[Float[Tensor, 'num_views RGB image_height image_width'] | Float[Tensor, 'num_views frames RGB image_height image_width'], Float[Tensor, 'keypoints'], Float[Tensor, 'num_views heatmap_height heatmap_width'], Float[Tensor, 'num_views_x_xyhw'], list][source]

Merge images, heatmaps, keypoints, and bboxes across views.

Parameters:

datadict – this comes from HeatmapDataset.__getItems__(idx) for each view.

Returns:

tuple

  • images

  • keypoints

  • heatmaps

  • bboxes

  • concat order

__init__(root_directory: str | Path, csv_paths: list[str], view_names: list[str], image_resize_height: int, image_resize_width: int, header_rows: list[int] | None = [0, 1, 2], imgaug_transform: Sequential | None = None, downsample_factor: Literal[1, 2, 3] = 2, do_context: bool = False, resize: bool = False, uniform_heatmaps: bool = False, camera_params_path: str | None = None, bbox_paths: list[str] | None = None) → None[source]

Initialize the MultiViewHeatmap Dataset.

Parameters:

  • root_directory – path to data directory

  • csv_paths – paths to CSV files (within root_directory). CSV files should be in this form (image_path, bodypart_1_x, bodypart_1_y, …, bodypart_n_y) these should match in all CSV files Note: image_path is relative to the given root_directory

  • view_names – a list of strings with the view names

  • image_resize_height – height to resize images before sending to network

  • image_resize_width – height to resize images before sending to network

  • header_rows – which rows in the csv are header rows

  • imgaug_transform – imgaug transform pipeline to apply to images

  • downsample_factor – factor by which to downsample original image dims to have a smaller heatmap

  • do_context – include additional frames of context if possible

  • resize – True to add final resizing augmentation before sending data to network. This can be set to False if inheritors of this class need to implement more sophisticated augmentations before resizing (e.g. 3d augmentations). Note that when this is False, it is up to the child class to perform this resizing on both images and keypoints before returning a batch of data.

  • uniform_heatmaps – True to force the model to output uniform heatmaps for missing data; False will output all-zero heatmaps

  • camera_params_path – path to toml file with camera calibration parameters in format output by anipose

  • bbox_paths – paths to csv files of the form (image_path, x, h, height, width) where (x, y) correspond to upper left corner of bbox. These files should be in the same view order as the csv paths

__new__(**kwargs)