deeplabcut.pose_estimation_pytorch.post_processing.identity

Functions to assign identity to predictions from an identity head.

Functions:

Name	Description
assign_identity	Args:

assign_identity

assign_identity(predictions: ndarray, identity_scores: ndarray) -> np.ndarray

Parameters:

Name	Type	Description	Default
predictions	ndarray	Pose predictions for an image, with shape (num_individuals, num_bodyparts, 3)	required
identity_scores	ndarray	Identity predictions for keypoints in an image, of shape (num_individuals, num_bodyparts, num_individuals).	required

Returns:

Type	Description
ndarray	The ordering to use to match predictions to identities.

Source code in deeplabcut/pose_estimation_pytorch/post_processing/identity.py

def assign_identity(predictions: np.ndarray, identity_scores: np.ndarray) -> np.ndarray:
    """
    Args:
        predictions: Pose predictions for an image, with shape (num_individuals,
            num_bodyparts, 3)
        identity_scores: Identity predictions for keypoints in an image, of shape
            (num_individuals, num_bodyparts, num_individuals).

    Returns:
        The ordering to use to match predictions to identities.
    """
    if not len(predictions) == len(identity_scores):
        raise ValueError(
            "There are not the same number of predictions as identity scores"
            f" ({len(predictions)} != {len(identity_scores)}"
        )

    # average of ID scores, weighted by keypoint confidence
    pose_conf = predictions[:, :, 2:3]
    cost_matrix = np.mean(pose_conf * identity_scores, axis=1)

    row_ind, col_ind = linear_sum_assignment(cost_matrix, maximize=True)
    new_order = np.zeros_like(row_ind)
    for old_pos, new_pos in zip(row_ind, col_ind, strict=True):
        new_order[new_pos] = old_pos

    return new_order