diplomat.processing.TrackingData

class diplomat.processing.TrackingData(scmap: ndarray, locref: ndarray | None = None, stride: float = 8)

Bases: object

Represents tracking data received from a neural network. Includes a source map of probabilities, the predicted location offsets within the source map, stride info and ect. Also provides many convenience methods for working with and getting info from the DLC neural network data.

__init__(scmap: ndarray, locref: ndarray | None = None, stride: float = 8)

Create an new tracking data object to store DLC neural network data for one frame or a batch of frames.

Parameters:

• scmap – The probability maps produced by the neural network, a 4-dimensional numpy array containing the dimensions: [frame, y location, x location, body part].

• locref – The “offsets” produced by the neural network, stored in a 5-dimensional numpy array containing the dimensions: [frame, y location, x location, bodypart, 0 for x offset or 1 for y offset]

• stride – Float which stores the down scaling of the probability map relative to the size of the original video. This value defaults to 8, meaning the original video is 8 times the size of the probability map.

Methods

__init__(scmap[, locref, stride])	Create an new tracking data object to store DLC neural network data for one frame or a batch of frames.
empty_tracking_data(frame_amt, part_count, ...)	Create a new empty tracking data object with space allocated to fit the specified sizes of data.
get_bodypart_count()	Get the number of body parts stored in this TrackingData object per frame.
get_down_scaling()	Get the down scaling performed on this source map, as an integer.
get_frame_count()	Get the number of frames stored in this TrackingData object.
get_frame_height()	Return the height of each probability map in this TrackingData object.
get_frame_width()	Return the width of each probability map in this TrackingData object.
get_max_of_frame(frame[, num_outputs])	Get the locations of the highest probabilities for a single frame in the array.
get_max_scmap_points([num_max])	Get the locations with the max probabilities for each frame in this TrackingData.
get_offset_map()	Returns the offset prediction map representing offset predictions for each location in the probability map.
get_poses_for(points)	Return a pose object for the "maximum" predicted indexes passed in.
get_prob_table(frame, bodypart)	Get the probability map for a selection of frames and body parts or a single frame and body part.
get_source_map()	Gets the raw probability source map of this tracking data.
set_down_scaling(scale)	Set the down scaling performed on the probability map, as an integer
set_offset_map(locref)	Set the offset prediction map representing offset predictions for each location in the probability map.
set_pose_at(frame, bodypart, scmap_x, ...[, ...])	Set a pose in the specified Pose object to the specified x and y coordinate for a provided body part and frame.
set_prob_table(frame, bodypart, values)	Set the probability table for a selection of frames and body parts or a single frame and body part.
set_source_map(scmap)	Set the raw probability map of this tracking data.

Attributes

DEFAULT_SCALE

static _get_count_of(val: int | slice | Sequence[int], length: int) → int

Internal private method to get length of an index selection(as in how many indexes it selects…)

classmethod empty_tracking_data(frame_amt: int, part_count: int, width: int, height: int, stride: float = 8, allocate_offsets: bool = False) → TrackingData

Create a new empty tracking data object with space allocated to fit the specified sizes of data.

Parameters:

• frame_amt – The amount of probability map frames to allocate space for, an Integer.

• part_count – The amount of body parts per frame to allocate space for, an Integer.

• width – The width of each probability frame, an Integer.

• height – The height of each probability frame, an Integer.

• stride – The downscaling of the probability frame relative to the original video, defaults to 8, meaning the original video is 8 times the size of the probability map.

• allocate_offsets – A boolean, determines if the offset array is setup as an array of zeros. Default value is False, meaning the offsets array is just set to None.

Returns:

A tracking data object full of zeroes.

get_bodypart_count() → int

Get the number of body parts stored in this TrackingData object per frame.

Returns:

The number of body parts per frame as an integer.

get_down_scaling() → float

Get the down scaling performed on this source map, as an integer.

Returns:

An integer representing the downscaling of the probability map compared to the original video file, in terms of what the dimensions of the probability map need to multiplied by to match the dimensions of the original video file.

get_frame_count() → int

Get the number of frames stored in this TrackingData object.

Returns:

The number of frames stored in this tracking data object.

get_frame_height() → int

Return the height of each probability map in this TrackingData object.

Returns:

The height of each probability map as an integer.

get_frame_width() → int

Return the width of each probability map in this TrackingData object.

Returns:

The width of each probability map as an integer.

get_max_of_frame(frame: int, num_outputs: int = 1) → Tuple[ndarray, ndarray]

Get the locations of the highest probabilities for a single frame in the array.

Parameters:

• frame – The index of the frame to get the maximum of, in form of an integer.

• num_outputs – Specifies the number of maximums to grab for each body part from the frame. Defaults to 1.

Returns:

A tuple of numpy arrays, the first numpy array being the y index of the max probability for each body part in the frame, the second being the x index of the max probability for each body part in the frame Indexing: [1 for x or 0 for y index] -> [bodypart 1, bodypart 1 prediction 2, …, bodypart 2, …]

get_max_scmap_points(num_max: int = 1) → Tuple[ndarray, ndarray]

Get the locations with the max probabilities for each frame in this TrackingData.

Parameters:

num_max – Specifies the number of maximums to grab for each body part from each frame. Defaults to 1.

Returns:

A tuple of numpy arrays, the first numpy array being the y index for the max of each frame, the second being the x index for the max of each frame. Dimensions: [1 for x and 0 for y index, frame, body part] -> index

get_offset_map() → ndarray | None

Returns the offset prediction map representing offset predictions for each location in the probability map.

Returns:

A numpy array representing the predicted offsets for each location within the probability map, or None if this TrackingData doesn’t have offset predictions… Indexing: [frame, y location, x location, bodypart, 0 for x offset or 1 for y offset]

get_poses_for(points: Tuple[ndarray, ndarray])

Return a pose object for the “maximum” predicted indexes passed in.

Parameters:

points – A tuple of 2 numpy arrays, one representing the y indexes for each frame and body part, the other being the x indexes represented the same way. (Note ‘get_max_scmap_points’ returns maximum predictions in this exact format).

Returns:

The Pose object representing all predicted maximum locations for selected points…

NOTE: This method detects when multiple predictions(num_outputs > 1) have been made and will still work

correctly…

get_prob_table(frame: int | slice | Sequence[int], bodypart: int | slice | Sequence[int]) → ndarray

Get the probability map for a selection of frames and body parts or a single frame and body part.

Parameters:

• frame – The frame index, as an integer or slice.

• bodypart – The body part index, as an integer or slice.

Returns:

The probability map(s) for a single frame or selection of frames based on indexes, as a numpy array… Dimensions: [frame, body part, y location, x location] -> probability

get_source_map() → ndarray

Gets the raw probability source map of this tracking data.

Returns:

A numpy array representing the source probability map of this tracking data. It is a 4-dimensional array containing the dimensions: [frame, y location, x location, body part] -> probability

set_down_scaling(scale: float)

Set the down scaling performed on the probability map, as an integer

Parameters:

scale – An float representing the downscaling of the probability map compared to the original video file, in terms of what the dimensions of the probability map need to multiplied by to match the dimensions of the original video file.

set_offset_map(locref: ndarray | None)

Set the offset prediction map representing offset predictions for each location in the probability map.

Parameters:

locref – A numpy array representing the predicted offsets within the probability map. Can also be set to None to indicate this TrackingData doesn’t have or support higher precision predictions. Dimensions: [frame, y location, x location, bodypart, 0 for x offset or 1 for y offset]

set_pose_at(frame: int, bodypart: int, scmap_x: int, scmap_y: int, pose_object: Pose, output_num: int = 0)

Set a pose in the specified Pose object to the specified x and y coordinate for a provided body part and frame. This method will use data from this TrackingData object to correctly set the information in the Pose object.

Parameters:

• frame – The specified frame to copy from this TrackingData to the Pose object and set.

• bodypart – The specified body part to copy from this TrackingData to the pose object and set

• scmap_x – The x index of this TrackingData to set the Pose prediction to.

• scmap_y – The y index of this TrackingData to set the Pose prediction to.

• pose_object – The pose object to be modified/copied to.

• output_num – The output number to set in the pose object (Which prediction for this bodypart?). Should only be needed if num_outputs > 1. Defaults to 0, meaning the first prediction.

Returns:

Nothing, changes stored in pose_object…

set_prob_table(frame: int | slice | Sequence[int], bodypart: int | slice | Sequence[int], values: ndarray)

Set the probability table for a selection of frames and body parts or a single frame and body part.

Parameters:

• frame – The frame index, as an integer or slice.

• bodypart – The body part index, as an integer or slice.

• values – The probability map(s) to set in this TrackingData object based on the frame and body parts specified, as a numpy array… Dimensions of values: [frame, body part, y location, x location] -> probability

set_source_map(scmap: ndarray)

Set the raw probability map of this tracking data.

Parameters:

scmap – A numpy array representing the probability map of this tracking data. It is a 4-dimensional array containing the dimensions: [frame, y location, x location, body part]