DeepLabCut · n-poulsen · Nov 14, 2024 · Nov 11, 2024 · Nov 12, 2024 · Nov 13, 2024
diff --git a/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py b/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py
@@ -198,12 +198,15 @@ def make_pytorch_test_config(
         The test configuration file.
     """
     bodyparts = model_config["metadata"]["bodyparts"]
+    unique_bodyparts = model_config["metadata"]["unique_bodyparts"]
+    all_joint_names = bodyparts + unique_bodyparts
+
     test_config = dict(
         dataset=model_config["metadata"]["project_path"],
         dataset_type="multi-animal-imgaug",  # required for downstream tracking
-        num_joints=len(bodyparts),
-        all_joints=[[i] for i in range(len(bodyparts))],
-        all_joints_names=bodyparts,
+        num_joints=len(all_joint_names),
+        all_joints=[[i] for i in range(len(all_joint_names))],
+        all_joints_names=all_joint_names,
         net_type=model_config["net_type"],
         global_scale=1,
         scoremap_dir="test",

diff --git a/deeplabcut/pose_estimation_pytorch/data/postprocessor.py b/deeplabcut/pose_estimation_pytorch/data/postprocessor.py
@@ -67,7 +67,6 @@ def build_bottom_up_postprocessor(
         keys_to_rescale.append("unique_bodyparts")
 
     if with_identity:
-        # TODO: do we really want to return the heatmaps?
         keys_to_concatenate["identity_heatmap"] = ("identity", "heatmap")
         empty_shapes["identity_heatmap"] = (1, 1, max_individuals)
 
@@ -85,6 +84,7 @@ def build_bottom_up_postprocessor(
                 identity_key="identity_scores",
                 identity_map_key="identity_heatmap",
                 pose_key="bodyparts",
+                keep_id_maps=False,
             )
         )
 
@@ -96,8 +96,6 @@ def build_bottom_up_postprocessor(
         PadOutputs(
             max_individuals={
                 "bodyparts": max_individuals,
-                "unique_bodyparts": 0,  # no need to pad
-                "identity_heatmap": 0,  # no need to pad
                 "identity_scores": max_individuals,
             },
             pad_value=-1,
@@ -146,7 +144,6 @@ def build_top_down_postprocessor(
                     "bodyparts": max_individuals,
                     "bboxes": max_individuals,
                     "bbox_scores": max_individuals,
-                    "unique_bodyparts": 0,  # no need to pad
                 },
                 pad_value=-1,
             ),
@@ -257,7 +254,10 @@ def __call__(
     ) -> tuple[dict[str, np.ndarray], Context]:
         for name in predictions:
             output = predictions[name]
-            if len(output) < self.max_individuals[name]:
+            if (
+                name in self.max_individuals
+                and len(output) < self.max_individuals[name]
+            ):
                 pad_size = self.max_individuals[name] - len(output)
                 tail_shape = output.shape[1:]
                 padding = self.pad_value * np.ones((pad_size, *tail_shape))
@@ -404,39 +404,52 @@ def __call__(
 class PredictKeypointIdentities(Postprocessor):
     """Assigns predicted identities to keypoints
 
+    The identity maps have shape (h, w, num_ids).
+
     Attributes:
-        identity_key:
-        identity_map_key: shape (h, w, num_ids)
-        pose_key:
+        identity_key: Key with which to add predicted identities in the predictions dict
+        identity_map_key: Key for the identity maps in the predictions dict
+        pose_key: Key for the bodyparts in the predictions dict
+        keep_id_maps: Whether to keep identity heatmaps in the output dictionary.
+            Setting this value to True can be useful for debugging, but can lead to
+            memory issues when running video analysis on long videos.
     """
 
     def __init__(
         self,
         identity_key: str,
         identity_map_key: str,
         pose_key: str,
+        keep_id_maps: bool = False,
     ) -> None:
         self.identity_key = identity_key
         self.identity_map_key = identity_map_key
         self.pose_key = pose_key
+        self.keep_id_maps = keep_id_maps
 
     def __call__(
         self, predictions: dict[str, np.ndarray], context: Context
     ) -> tuple[dict[str, np.ndarray], Context]:
-        individuals = predictions[self.pose_key]
+        pose = predictions[self.pose_key]
+        num_preds, num_keypoints, _ = pose.shape
+
         identity_heatmap = predictions[self.identity_map_key]  # (h, w, num_ids)
         h, w, num_ids = identity_heatmap.shape
-        num_individuals, num_keypoints, _ = individuals.shape
 
-        assembly_id_scores = []
-        for individual_keypoints in individuals:
+        id_score_matrix = np.zeros((num_preds, num_keypoints, num_ids))
+        for pred_idx, individual_keypoints in enumerate(pose):
             heatmap_indices = np.rint(individual_keypoints).astype(int)
             xs = np.clip(heatmap_indices[:, 0], 0, w - 1)
             ys = np.clip(heatmap_indices[:, 1], 0, h - 1)
-            id_scores = []
-            for x, y in zip(xs, ys):
-                id_scores.append(identity_heatmap[y, x, :])
-            assembly_id_scores.append(np.stack(id_scores))
 
-        predictions[self.identity_key] = np.stack(assembly_id_scores)
+            # get the score from each identity heatmap at each predicted keypoint
+            for kpt_idx, (x, y) in enumerate(zip(xs, ys)):
+                id_score_matrix[pred_idx, kpt_idx] = identity_heatmap[y, x, :]
+
+        predictions[self.identity_key] = id_score_matrix
+        if not self.keep_id_maps:
+            # delete the heatmaps as this saves memory
+            id_heatmaps = predictions.pop(self.identity_map_key)
+            del id_heatmaps
+
         return predictions, context
diff --git a/deeplabcut/pose_estimation_pytorch/post_processing/identity.py b/deeplabcut/pose_estimation_pytorch/post_processing/identity.py
@@ -35,7 +35,10 @@ def assign_identity(
 
     predictions_with_identity = []
     for pred, scores in zip(predictions, identity_scores):
-        cost_matrix = np.product(scores, axis=1)
+        # average of ID scores, weighted by keypoint confidence
+        pose_conf = pred[:, :, 2:3]
+        cost_matrix = np.mean(pose_conf * 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):

diff --git a/tests/pose_estimation_pytorch/data/test_postprocessor.py b/tests/pose_estimation_pytorch/data/test_postprocessor.py
@@ -203,6 +203,7 @@ def test_assign_id_scores(data):
         identity_key="keypoint_identity",
         identity_map_key="identity_map",
         pose_key="bodyparts",
+        keep_id_maps=True,
     )
     bodyparts = np.array(data["bodyparts"])
     id_heatmap = np.array(data["id_heatmap"])