lisbet.datasets.common

Common code for selecting windows from a dataset of records.

Classes

AnnotatedWindowSelector(records, window_size)	WindowSelector with annotation extraction.
WindowSelector(records, window_size[, ...])	Selects windows from a dataset of records.

class lisbet.datasets.common.WindowSelector(records, window_size, window_offset=0, fps_scaling=1.0, engine='xarray')

Selects windows from a dataset of records.

This class provides methods to extract temporal windows from a list of records, handling padding and interpolation as needed. It supports mapping between global and local frame indices and can scale windows according to a frames-per-second (fps) scaling factor.

__init__(records, window_size, window_offset=0, fps_scaling=1.0, engine='xarray')

Initialize the WindowSelector.

Parameters:

• records (list) – List of records containing pose tracking data.

• window_size (int) – Size of the window in frames.

• window_offset (int, optional) – Offset for the window in frames (default is 0).

• fps_scaling (float, optional) – Scaling factor for the frames per second (default is 1.0).

• engine (Literal['xarray', 'numpy']) – Output engine. The default "xarray" returns an xarray Dataset with coordinates and data variables intact. "numpy" returns an owned, writable array with shape (time, individuals, keypoints, space).

Raises:

ValueError – If no records are provided or if any record contains fewer than 2 individuals, or if engine is unsupported.

global_to_local(global_idx)

Map a global frame index to a local (record_index, local_frame_index) pair.

Parameters:

global_idx (int) – Global frame index (0 ≤ global_idx < total_n_frames).

Returns:

• rec_idx (int) – Index of the record containing the frame.

• local_idx (int) – Local frame index within the selected record.

select(rec_idx, frame_idx, fps_scaling=None)

Select a window from the dataset, applying padding and interpolation as needed.

The selected window is returned as an independent xarray Dataset or NumPy array to avoid unintentional changes to source records (for example, by a self-supervised task or augmentation).

Parameters:

• rec_idx (int) – Index of the record from which to select the window.

• frame_idx (int) – Index of the central frame within the record.

• fps_scaling (float, optional) – Override the default fps scaling factor for this selection. If None, uses the default fps_scaling set during initialization.

Returns:

x – The selected and interpolated window. NumPy output has shape (time, individuals, keypoints, space).

Return type:

xarray.Dataset or numpy.ndarray

Notes

1. The interpolation is done here, and not directly on the records, to avoid resampling at the original fps before returning the output during inference. Furthermore, even during training, it is useful to only iterate over the original frames, rather than artificially inflating or deflating the dataset.

class lisbet.datasets.common.AnnotatedWindowSelector(records, window_size, window_offset=0, fps_scaling=1.0, annot_format='multiclass', engine='xarray')

WindowSelector with annotation extraction.

Extends WindowSelector to also extract annotation targets for each selected window, supporting binary, multiclass, and multilabel annotation formats.

__init__(records, window_size, window_offset=0, fps_scaling=1.0, annot_format='multiclass', engine='xarray')

Initialize the AnnotatedWindowSelector.

Parameters:

• records (list) – List of records containing the data and annotations.

• window_size (int) – Size of the window in frames.

• window_offset (int, optional) – Offset for the window in frames (default is 0).

• fps_scaling (float, optional) – Scaling factor for the frames per second (default is 1.0).

• annot_format (str, optional) – Format of the annotations (‘binary’, ‘multiclass’, or ‘multilabel’).

• engine (Literal['xarray', 'numpy']) – Output engine (default is "xarray").

Raises:

ValueError – If annot_format is not one of "binary", "multiclass", or "multilabel", or if engine is unsupported.

select(rec_idx, frame_idx, fps_scaling=None)

Select a window and its corresponding annotation target.

Parameters:

• rec_idx (int) – Index of the record from which to select the window.

• frame_idx (int) – Index of the central frame within the record.

• fps_scaling (float, optional) – Override the default fps scaling factor for this selection. If None, uses the default fps_scaling set during initialization.

Returns:

• x (xarray.Dataset or numpy.ndarray) – The selected and interpolated window. NumPy output has shape (time, individuals, keypoints, space).

• y (numpy.ndarray) – The annotation target(s) for the selected window, format depends on annot_format.