Speed up pair creation and matching.

scripts/test_skelda_dataset.py

@@ -319,6 +319,7 @@ def main():
     all_ids = []
     all_paths = []
     times = []
+    last_poses_3d = np.array([])
     for label in tqdm.tqdm(labels):
         images_2d = []
 
@@ -377,7 +378,7 @@ def main():
             poses2D = np.zeros([len(images_2d), 1, len(joint_names_3d), 3])
         else:
             poses3D = triangulate_poses.get_3d_pose(
-                poses_2d, label["cameras"], roomparams, joint_names_2d, minscore
+                poses_2d, label["cameras"], roomparams, joint_names_2d, last_poses_3d, minscore
             )
             poses2D = []
             for cam in label["cameras"]:
@@ -392,6 +393,7 @@ def main():
                 drop_few_limbs=(dataset_use != "mvor"),
             )
             poses3D = add_missing_joints(poses3D, joint_names_3d)
+        last_poses_3d = poses3D
 
         time_3d = time.time() - start
         print("3D time:", time_3d)

scripts/triangulate_poses.py

@@ -1,6 +1,5 @@
 import copy
 import math
-import time
 
 import cv2
 import numpy as np
@@ -66,6 +65,29 @@ def get_camera_P(cam):
 # ==================================================================================================
 
 
+def calc_pose_score(pose1, pose2, dist1, cam1, joint_names, use_joints):
+    """Calculates the score between two poses"""
+
+    # Select core joints
+    jids = [joint_names.index(j) for j in use_joints]
+    pose1 = pose1[jids]
+    pose2 = pose2[jids]
+    dist1 = dist1[jids]
+
+    mask = (pose1[:, 2] > 0.1) & (pose2[:, 2] > 0.1)
+    if np.sum(mask) < 3:
+        return 0.0
+
+    iscale = (cam1["width"] + cam1["height"]) / 2
+    scores = score_projection(pose1, pose2, dist1, mask, iscale)
+    score = np.mean(scores)
+
+    return score
+
+
+# ==================================================================================================
+
+
 def calc_pair_score(pair, poses_2d, camparams, roomparams, joint_names_2d, use_joints):
     """Triangulates a pair of persons and scores them based on the reprojection error"""
 
@@ -79,14 +101,38 @@ def calc_pair_score(pair, poses_2d, camparams, roomparams, joint_names_2d, use_j
     pose1 = pose1[jids]
     pose2 = pose2[jids]
 
-    poses_3d, score = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
+    poses_3d, score = triangulate_and_score(pose1, pose2, cam1, cam2, roomparams)
     return poses_3d, score
 
 
 # ==================================================================================================
 
 
-def calc_pose_scored(pose1, pose2, cam1, cam2, roomparams):
+def score_projection(pose1, repro1, dists1, mask, iscale):
+
+    min_score = 0.1
+    error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[mask, 0:2], axis=1)
+
+    # Set errors of invisible reprojections to a high value
+    penalty = iscale
+    mask1b = (repro1[:, 2] < min_score)[mask]
+    error1[mask1b] = penalty
+
+    # Scale error by image size and distance to the camera
+    error1 = error1.clip(0, iscale / 4) / iscale
+    dscale1 = np.sqrt(np.mean(dists1[mask]) / 3.5)
+    error1 = error1 * dscale1
+
+    # Convert errors to a score
+    score1 = 1.0 / (1.0 + error1 * 10)
+
+    return score1
+
+
+# ==================================================================================================
+
+
+def triangulate_and_score(pose1, pose2, cam1, cam2, roomparams):
     """Triangulates a pair of persons and scores them based on the reprojection error"""
 
     # Mask out invisible joints
@@ -129,27 +175,16 @@ def calc_pose_scored(pose1, pose2, cam1, cam2, roomparams):
     poses_3d = np.expand_dims(pose3d, axis=0)
     repro1, dists1 = utils_pose.project_poses(poses_3d, cam1, calc_dists=True)
     repro2, dists2 = utils_pose.project_poses(poses_3d, cam2, calc_dists=True)
-    error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[0, mask, 0:2], axis=1)
+    repro1, dists1 = repro1[0], dists1[0]
-    error2 = np.linalg.norm(pose2[mask, 0:2] - repro2[0, mask, 0:2], axis=1)
+    repro2, dists2 = repro2[0], dists2[0]
 
-    # Set errors of invisible reprojections to a high value
+    # Calculate scores for each view
-    penalty = (cam1["width"] + cam1["height"]) / 2
-    mask1b = (repro1[0, :, 2] < min_score)[mask]
-    mask2b = (repro2[0, :, 2] < min_score)[mask]
-    error1[mask1b] = penalty
-    error2[mask2b] = penalty
-
-    # Convert errors to a score
-    # Scale by image size and distance to the camera
     iscale = (cam1["width"] + cam1["height"]) / 2
-    error1 = error1.clip(0, iscale / 4) / iscale
+    score1 = score_projection(pose1, repro1, dists1, mask, iscale)
-    error2 = error2.clip(0, iscale / 4) / iscale
+    score2 = score_projection(pose2, repro2, dists2, mask, iscale)
-    dscale1 = np.sqrt(np.mean(dists1[0, mask]) / 3.5)
+
-    dscale2 = np.sqrt(np.mean(dists2[0, mask]) / 3.5)
+    # Combine scores
-    error1 = error1 * dscale1
+    scores = (score1 + score2) / 2
-    error2 = error2 * dscale2
-    error = (error1 + error2) / 2
-    scores = 1.0 / (1.0 + error * 10)
 
     # Drop lowest scores
     drop_k = math.floor(len(pose1) * 0.2)
@@ -283,7 +318,14 @@ def merge_group(poses_3d: np.ndarray, min_score: float):
 # ==================================================================================================
 
 
-def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95):
+def get_3d_pose(
+    poses_2d,
+    camparams,
+    roomparams,
+    joint_names_2d,
+    last_poses_3d=np.array([]),
+    min_score=0.95,
+):
     """Triangulates 3D poses from 2D poses of multiple views"""
 
     # Convert poses and camparams to numpy arrays
@@ -314,10 +356,42 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
         cam["P"] = get_camera_P(cam)
         camparams[i] = cam
 
+    # Project last 3D poses to 2D
+    last_poses_2d = []
+    last_poses_3d = np.asarray(last_poses_3d)
+    if last_poses_3d.size > 0:
+        for i in range(len(camparams)):
+            poses2d, dists = utils_pose.project_poses(last_poses_3d, camparams[i])
+            last_poses_2d.append((poses2d, dists))
+
+    # Check matches to old poses
+    threshold = min_score - 0.2
+    scored_pasts = {}
+    if last_poses_3d.size > 0:
+        for i in range(len(camparams)):
+            scored_pasts[i] = {}
+            poses = poses_2d[i]
+            last_poses, dists = last_poses_2d[i]
+            for j in range(len(last_poses)):
+                scored_pasts[i][j] = []
+                for k in range(len(poses)):
+                    score = calc_pose_score(
+                        poses[k],
+                        last_poses[j],
+                        dists[j],
+                        camparams[i],
+                        joint_names_2d,
+                        core_joints,
+                    )
+                    if score > threshold:
+                        scored_pasts[i][j].append(k)
+
     # Create pairs of persons
+    # Checks if the person was already matched to the last frame and if so only creates pairs with those
+    # Else it creates all possible pairs
     num_persons = [len(p) for p in poses_2d]
     all_pairs = []
-    for i in range(len(poses_2d)):
+    for i in range(len(camparams)):
         poses = poses_2d[i]
         for j in range(i + 1, len(poses_2d)):
             poses2 = poses_2d[j]
@@ -325,6 +399,15 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
                 for l in range(len(poses2)):
                     pid1 = sum(num_persons[:i]) + k
                     pid2 = sum(num_persons[:j]) + l
+                    match = False
+                    if last_poses_3d.size > 0:
+                        for m in range(len(last_poses_3d)):
+                            if k in scored_pasts[i][m] and l in scored_pasts[j][m]:
+                                match = True
+                                all_pairs.append([(i, j, k, l), (pid1, pid2)])
+                            elif k in scored_pasts[i][m] or l in scored_pasts[j][m]:
+                                match = True
+                    if not match:
                         all_pairs.append([(i, j, k, l), (pid1, pid2)])
 
     # Calculate pair scores
@@ -356,7 +439,7 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
         pose1 = poses_2d[pair[0][0]][pair[0][2]]
         pose2 = poses_2d[pair[0][1]][pair[0][3]]
 
-        pose_3d, _ = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
+        pose_3d, _ = triangulate_and_score(pose1, pose2, cam1, cam2, roomparams)
         pair.append(pose_3d)
 
     # Merge groups