MultiviewHeatmapDataset

class lightning_pose.data.datasets.MultiviewHeatmapDataset

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()

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]

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], image_resize_height: int, image_resize_width: int, header_rows: list[int] | None = [0, 1, 2], imgaug_transform: Callable | None = None, downsample_factor: Literal[1, 2, 3] = 2, do_context: bool = False, resize: bool = True, uniform_heatmaps: bool = False) → None

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

• 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

__new__(**kwargs)