BaseTrackingDataset

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

Bases: Dataset

Base dataset that contains images and keypoints as (x, y) pairs.

Attributes Summary

height

width

Attributes Documentation

height
width
__init__(root_directory: str, csv_path: str, image_resize_height: int, image_resize_width: int, header_rows: list[int] | None = [0, 1, 2], imgaug_transform: Callable | None = None, do_context: bool = False, resize: bool = True, bbox_path: str | None = None) → None[source]

Initialize a dataset for regression (rather than heatmap) models.

The csv file of labels will be searched for in the following order: 1. assume csv is located at root_directory/csv_path (i.e. csv_path

argument is a path relative to root_directory)

  1. if not found, assume csv_path is absolute. Note the image paths

    within the csv must still be relative to root_directory

  2. if not found, assume dlc directory structure: root_directory/training-data/iteration-0/csv_path (csv_path argument will look like β€œCollectedData_<scorer>.csv”)

Parameters:

  • root_directory – path to data directory

  • csv_path – path to CSV file (within root_directory). CSV file should be in the form (image_path, bodypart_1_x, bodypart_1_y, …, bodypart_n_y) Note: image_path is relative to the given root_directory

  • resize_height – height to resize images before sending to network

  • 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

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

  • bbox_path – path to csv file that contains bounding box information; rows must be in same order as csv file

__new__(**kwargs)