PredictionHandler

class lightning_pose.utils.predictions.PredictionHandler

Bases: object

Convert batches of model outputs into a prediction dataframe.

Attributes Summary

do_context	Whether the model/loader uses 5-frame context.
frame_count	Returns the number of frames in the video or the labeled dataset
keypoint_names	List of keypoint name strings from the config.

Methods Summary

__call__()	Call this function to get a pandas dataframe of the predictions for a single video.
add_split_indices_to_df(df)	Add split indices to the dataframe.
fix_context_preds_confs(stacked_preds[, ...])	In the context model, ind=0 is associated with image[2], and ind=1 is associated with image[3], so we need to shift the predictions and confidences by two and eliminate the edges.
make_dlc_pandas_index([keypoint_names])	Build a DLC-style pandas MultiIndex for labelling prediction DataFrames.
make_pred_arr_undo_resize(keypoints_np, ...)	Resize keypoints and add confidences into one numpy array.
unpack_preds(preds)	unpack list of preds coming out from pl.trainer.predict, confs tuples into tensors.

Attributes Documentation

do_context

Whether the model/loader uses 5-frame context.

Returns:

True if context frames are used, otherwise False.

frame_count

Returns the number of frames in the video or the labeled dataset

keypoint_names

List of keypoint name strings from the config.

Returns:

List of keypoint names.

Methods Documentation

__call__(preds: list[tuple[Float[Tensor, 'batch two_times_num_keypoints'], Float[Tensor, 'batch num_keypoints']]], is_multiview_video: Literal[False] = False) → DataFrame

__call__(preds: list[tuple[Float[Tensor, 'batch two_times_num_keypoints'], Float[Tensor, 'batch num_keypoints']]], is_multiview_video: Literal[True]) → dict[str, DataFrame]

Call this function to get a pandas dataframe of the predictions for a single video. Assuming you’ve already run trainer.predict(), and have a list of Tuple predictions.

Parameters:

• preds – list of tuples of (predictions, confidences)

• is_multiview_video – specify True when you are using multiview video prediction dataloader, i.e. for heatmap_multiview.

Returns:

index is (frame, bodypart, x, y, likelihood)

Return type:

pd.DataFrame

add_split_indices_to_df(df: DataFrame) → DataFrame

Add split indices to the dataframe.

fix_context_preds_confs(stacked_preds: Tensor, zero_pad_confidence: bool = False) → Tensor

In the context model, ind=0 is associated with image[2], and ind=1 is associated with image[3], so we need to shift the predictions and confidences by two and eliminate the edges. NOTE: confidences are not zero in the first and last two images, they are instead replicas of images[-2] and images[-3]

make_dlc_pandas_index(keypoint_names: list | None = None) → MultiIndex

Build a DLC-style pandas MultiIndex for labelling prediction DataFrames.

Parameters:

keypoint_names – optional override for the list of keypoint names; defaults to self.keypoint_names.

Returns:

pd.MultiIndex with levels ["scorer", "bodyparts", "coords"].

static make_pred_arr_undo_resize(keypoints_np: ndarray, confidence_np: ndarray) → ndarray

Resize keypoints and add confidences into one numpy array.

Parameters:

• keypoints_np – shape (n_frames, n_keypoints * 2)

• confidence_np – shape (n_frames, n_keypoints)

Returns:

cols are (bp0_x, bp0_y, bp0_likelihood, bp1_x, bp1_y, …)

Return type:

np.ndarray

unpack_preds(preds: list[tuple[Float[Tensor, 'batch two_times_num_keypoints'], Float[Tensor, 'batch num_keypoints']]]) → tuple[Float[Tensor, 'num_frames two_times_num_keypoints'], Float[Tensor, 'num_frames num_keypoints']]

unpack list of preds coming out from pl.trainer.predict, confs tuples into tensors. It still returns unnecessary final rows, which should be discarded at the dataframe stage. This works for the output of predict_loader, suitable for batch_size=1, sequence_length=16, step=16

__init__(cfg: DictConfig | ListConfig, data_module: BaseDataModule | UnlabeledDataModule | None = None, video_file: str | None = None) → None

Args

cfg data_module: Only required for prediction of CSV files. video_file: For prediction on video, path to the video file.

Used to get frame_count.

__new__(**kwargs)