BaseTrackingDataset

class lightning_pose.data.datasets.BaseTrackingDataset(root_directory: str, csv_path: str, header_rows: list[int] | None = [0, 1, 2], imgaug_transform: Callable | None = None, do_context: bool = False)

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, header_rows: list[int] | None = [0, 1, 2], imgaug_transform: Callable | None = None, do_context: bool = False) → None

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)

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

   within the csv must still be relative to root_directory

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

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