MultiviewHeatmapDataset

class lightning_pose.data.datasets.MultiviewHeatmapDataset(root_directory: str, csv_paths: list[str], view_names: list[str], header_rows: list[int] | None = [0, 1, 2], downsample_factor: Literal[1, 2, 3] = 2, uniform_heatmaps: bool = False, do_context: bool = False, imgaug_transform: Callable | None = None)[source]

Bases: Dataset

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

Attributes Summary

height

num_views

output_shape

width

Methods Summary

check_data_images_names()

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
num_views
output_shape
width

Methods Documentation

check_data_images_names()[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.

fusion(datadict: dict) → 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, csv_paths: list[str], view_names: list[str], header_rows: list[int] | None = [0, 1, 2], downsample_factor: Literal[1, 2, 3] = 2, uniform_heatmaps: bool = False, do_context: bool = False, imgaug_transform: Callable | 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 we suggest that these CSV files start with the view numbers

  • view_names – a list of integers with the view numbers

  • 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