Skip to content

utils

sleap_nn.tracking.utils

Helper functions for Tracker module.

Functions:

Name Description
compute_cosine_sim

Return cosine simalirity between a and b vectors.
compute_euclidean_distance

Return the negative euclidean distance between a and b points.
compute_iou

Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax].
get_bbox

Return the bounding box coordinates for the PredictedInstance object.
get_centroid

Return the centroid of the PredictedInstance object.
get_keypoints

Return keypoints as np.array from the PredictedInstance object.
greedy_matching

Match new instances to existing tracks using greedy bipartite matching.
hungarian_matching

Match new instances to existing tracks using Hungarian matching.

compute_cosine_sim(a, b)

Return cosine simalirity between a and b vectors.

Source code in sleap_nn/tracking/utils.py 
 95
 96
 97
 98
 99
100
def compute_cosine_sim(a, b):
    """Return cosine simalirity between a and b vectors."""
    number = np.dot(a, b)
    denom = np.linalg.norm(a) * np.linalg.norm(b)
    cosine_sim = number / denom
    return cosine_sim

compute_euclidean_distance(a, b)

Return the negative euclidean distance between a and b points.

Source code in sleap_nn/tracking/utils.py 
70
71
72
def compute_euclidean_distance(a, b):
    """Return the negative euclidean distance between a and b points."""
    return -np.linalg.norm(a - b)

compute_iou(a, b)

Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax].

Source code in sleap_nn/tracking/utils.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def compute_iou(a, b):
    """Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax]."""
    (xmin1, ymin1, xmax1, ymax1), (xmin2, ymin2, xmax2, ymax2) = a, b

    xmin_intersection = max(xmin1, xmin2)
    ymin_intersection = max(ymin1, ymin2)
    xmax_intersection = min(xmax1, xmax2)
    ymax_intersection = min(ymax1, ymax2)

    intersection_area = max(0, xmax_intersection - xmin_intersection + 1) * max(
        0, ymax_intersection - ymin_intersection + 1
    )
    bbox1_area = (xmax1 - xmin1 + 1) * (ymax1 - ymin1 + 1)
    bbox2_area = (xmax2 - xmin2 + 1) * (ymax2 - ymin2 + 1)
    union_area = bbox1_area + bbox2_area - intersection_area

    iou = intersection_area / union_area
    return iou

get_bbox(pred_instance)

Return the bounding box coordinates for the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def get_bbox(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return the bounding box coordinates for the `PredictedInstance` object."""
    points = (
        pred_instance.numpy()
        if not isinstance(pred_instance, np.ndarray)
        else pred_instance
    )
    bbox = np.concatenate(
        [
            np.nanmin(points, axis=0),
            np.nanmax(points, axis=0),
        ]  # [xmin, ymin, xmax, ymax]
    )
    return bbox

get_centroid(pred_instance)

Return the centroid of the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py 
45
46
47
48
49
50
51
def get_centroid(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return the centroid of the `PredictedInstance` object."""
    pts = pred_instance
    if not isinstance(pred_instance, np.ndarray):
        pts = pred_instance.numpy()
    centroid = np.nanmedian(pts, axis=0)
    return centroid

get_keypoints(pred_instance)

Return keypoints as np.array from the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py 
38
39
40
41
42
def get_keypoints(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return keypoints as np.array from the `PredictedInstance` object."""
    if isinstance(pred_instance, np.ndarray):
        return pred_instance
    return pred_instance.numpy()

greedy_matching(cost_matrix)

Match new instances to existing tracks using greedy bipartite matching.

Source code in sleap_nn/tracking/utils.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def greedy_matching(cost_matrix: np.ndarray) -> List[Tuple[int, int]]:
    """Match new instances to existing tracks using greedy bipartite matching."""
    # Sort edges by ascending cost.
    rows, cols = np.unravel_index(np.argsort(cost_matrix, axis=None), cost_matrix.shape)
    unassigned_edges = list(zip(rows, cols))

    # Greedily assign edges.
    row_inds, col_inds = [], []
    while len(unassigned_edges) > 0:
        # Assign the lowest cost edge.
        row_ind, col_ind = unassigned_edges.pop(0)
        row_inds.append(row_ind)
        col_inds.append(col_ind)

        # Remove all other edges that contain either node (in reverse order).
        for i in range(len(unassigned_edges) - 1, -1, -1):
            if unassigned_edges[i][0] == row_ind or unassigned_edges[i][1] == col_ind:
                del unassigned_edges[i]

    return row_inds, col_inds

hungarian_matching(cost_matrix)

Match new instances to existing tracks using Hungarian matching.

Source code in sleap_nn/tracking/utils.py 
10
11
12
13
def hungarian_matching(cost_matrix: np.ndarray) -> List[Tuple[int, int]]:
    """Match new instances to existing tracks using Hungarian matching."""
    row_ids, col_ids = linear_sum_assignment(cost_matrix)
    return row_ids, col_ids