Unified sports2d and pose2Sim sorting

Pose2Sim/common.py

@@ -18,6 +18,7 @@
 import cv2
 import c3d
 import sys
+import itertools as it
 
 import matplotlib as mpl
 mpl.use('qt5agg')
@@ -291,6 +292,87 @@ def euclidean_distance(q1, q2):
     return euc_dist
 
 
+def pad_shape(arr, target_len, fill_value=np.nan):
+    '''
+    Pads an array to the target length with specified fill values
+    
+    INPUTS:
+    - arr: Input array to be padded.
+    - target_len: The target length of the first dimension after padding.
+    - fill_value: The value to use for padding (default: np.nan).
+    
+    OUTPUTS:
+    - Padded array with shape (target_len, ...) matching the input dimensions.
+    '''
+
+    if len(arr) < target_len:
+        pad_shape = (target_len - len(arr),) + arr.shape[1:]
+        padding = np.full(pad_shape, fill_value)
+        return np.concatenate((arr, padding))
+
+    return arr
+
+
+def sort_people_sports2d(keyptpre, keypt, scores=None):
+    '''
+    Associate persons across frames (Sports2D method)
+    Persons' indices are sometimes swapped when changing frame
+    A person is associated to another in the next frame when they are at a small distance
+    
+    N.B.: Requires min_with_single_indices and euclidian_distance function (see common.py)
+
+    INPUTS:
+    - keyptpre: (K, L, M) array of 2D coordinates for K persons in the previous frame, L keypoints, M 2D coordinates
+    - keypt: idem keyptpre, for current frame
+    - score: (K, L) array of confidence scores for K persons, L keypoints (optional) 
+    
+    OUTPUTS:
+    - sorted_prev_keypoints: array with reordered persons with values of previous frame if current is empty
+    - sorted_keypoints: array with reordered persons --> if scores is not None
+    - sorted_scores: array with reordered scores     --> if scores is not None
+    - associated_tuples: list of tuples with correspondences between persons across frames --> if scores is None (for Pose2Sim.triangulation())
+    '''
+
+    # Generate possible person correspondences across frames
+    max_len = max(len(keyptpre), len(keypt))
+    keyptpre = pad_shape(keyptpre, max_len, fill_value=np.nan)
+    keypt = pad_shape(keypt, max_len, fill_value=np.nan)
+    if scores is not None:
+        scores = pad_shape(scores, max_len, fill_value=np.nan)
+
+    # Compute distance between persons from one frame to another
+    personsIDs_comb = sorted(list(it.product(range(len(keyptpre)), range(len(keypt)))))
+    frame_by_frame_dist = [euclidean_distance(keyptpre[comb[0]],keypt[comb[1]]) for comb in personsIDs_comb]
+    frame_by_frame_dist = np.mean(frame_by_frame_dist, axis=1)
+
+    # Sort correspondences by distance
+    _, _, associated_tuples = min_with_single_indices(frame_by_frame_dist, personsIDs_comb)
+
+    # Associate points to same index across frames, nan if no correspondence
+    sorted_keypoints = []
+    for i in range(len(keyptpre)):
+        id_in_old =  associated_tuples[:,1][associated_tuples[:,0] == i].tolist()
+        if len(id_in_old) > 0:      sorted_keypoints += [keypt[id_in_old[0]]]
+        else:                       sorted_keypoints += [keypt[i]]
+    sorted_keypoints = np.array(sorted_keypoints)
+
+    if scores is not None:
+        sorted_scores = []
+        for i in range(len(keyptpre)):
+            id_in_old =  associated_tuples[:,1][associated_tuples[:,0] == i].tolist()
+            if len(id_in_old) > 0:  sorted_scores += [scores[id_in_old[0]]]
+            else:                   sorted_scores += [scores[i]]
+        sorted_scores = np.array(sorted_scores)
+
+    # Keep track of previous values even when missing for more than one frame
+    sorted_prev_keypoints = np.where(np.isnan(sorted_keypoints) & ~np.isnan(keyptpre), keyptpre, sorted_keypoints)
+
+    if scores is not None:
+        return sorted_prev_keypoints, sorted_keypoints, sorted_scores
+    else: # For Pose2Sim.triangulation()
+        return sorted_keypoints, associated_tuples
+
+
 def trimmed_mean(arr, trimmed_extrema_percent=0.5):
     '''
     Trimmed mean calculation for an array.

Pose2Sim/poseEstimation.py

@@ -36,13 +36,11 @@
 import glob
 import json
 import logging
-import itertools as it
 from tqdm import tqdm
-import numpy as np
 import cv2
 
 from rtmlib import PoseTracker, Body, Wholebody, BodyWithFeet, draw_skeleton
-from Pose2Sim.common import natural_sort_key, min_with_single_indices, euclidean_distance
+from Pose2Sim.common import natural_sort_key, sort_people_sports2d
 
 
 ## AUTHORSHIP INFORMATION
@@ -99,62 +97,6 @@ def save_to_openpose(json_file_path, keypoints, scores):
     with open(json_file_path, 'w') as json_file:
         json.dump(json_output, json_file)
 
-
-def sort_people_sports2d(keyptpre, keypt, scores):
-    '''
-    Associate persons across frames (Pose2Sim method)
-    Persons' indices are sometimes swapped when changing frame
-    A person is associated to another in the next frame when they are at a small distance
-    
-    N.B.: Requires min_with_single_indices and euclidian_distance function (see common.py)
-
-    INPUTS:
-    - keyptpre: array of shape K, L, M with K the number of detected persons,
-    L the number of detected keypoints, M their 2D coordinates
-    - keypt: idem keyptpre, for current frame
-    - score: array of shape K, L with K the number of detected persons,
-    L the confidence of detected keypoints
-    
-    OUTPUTS:
-    - sorted_prev_keypoints: array with reordered persons with values of previous frame if current is empty
-    - sorted_keypoints: array with reordered persons
-    - sorted_scores: array with reordered scores
-    '''
-
-    # Generate possible person correspondences across frames
-    if len(keyptpre) < len(keypt):
-        keyptpre = np.concatenate((keyptpre, np.full((len(keypt)-len(keyptpre), keypt.shape[1], 2), np.nan)))
-    if len(keypt) < len(keyptpre):
-        keypt = np.concatenate((keypt, np.full((len(keyptpre)-len(keypt), keypt.shape[1], 2), np.nan)))
-        scores = np.concatenate((scores, np.full((len(keyptpre)-len(scores), scores.shape[1]), np.nan)))
-    personsIDs_comb = sorted(list(it.product(range(len(keyptpre)), range(len(keypt)))))
-
-    # Compute distance between persons from one frame to another
-    frame_by_frame_dist = []
-    for comb in personsIDs_comb:
-        frame_by_frame_dist += [euclidean_distance(keyptpre[comb[0]],keypt[comb[1]])]
-    frame_by_frame_dist = np.mean(frame_by_frame_dist, axis=1)
-
-    # Sort correspondences by distance
-    _, _, associated_tuples = min_with_single_indices(frame_by_frame_dist, personsIDs_comb)
-
-    # Associate points to same index across frames, nan if no correspondence
-    sorted_keypoints, sorted_scores = [], []
-    for i in range(len(keyptpre)):
-        id_in_old =  associated_tuples[:,1][associated_tuples[:,0] == i].tolist()
-        if len(id_in_old) > 0:
-            sorted_keypoints += [keypt[id_in_old[0]]]
-            sorted_scores += [scores[id_in_old[0]]]
-        else:
-            sorted_keypoints += [keypt[i]]
-            sorted_scores += [scores[i]]
-    sorted_keypoints, sorted_scores = np.array(sorted_keypoints), np.array(sorted_scores)
-
-    # Keep track of previous values even when missing for more than one frame
-    sorted_prev_keypoints = np.where(np.isnan(sorted_keypoints) & ~np.isnan(keyptpre), keyptpre, sorted_keypoints)
-
-    return sorted_prev_keypoints, sorted_keypoints, sorted_scores
-
 
 def process_video(video_path, pose_tracker, output_format, save_video, save_images, display_detection, frame_range, multi_person):
     '''
@@ -217,7 +159,7 @@ def process_video(video_path, pose_tracker, output_format, save_video, save_imag
                 # Tracking people IDs across frames
                 if multi_person:
                     if 'prev_keypoints' not in locals(): prev_keypoints = keypoints
-                    prev_keypoints, keypoints, scores = sort_people_sports2d(prev_keypoints, keypoints, scores)
+                    prev_keypoints, keypoints, scores = sort_people_sports2d(prev_keypoints, keypoints, scores=scores)
 
                 # Save to json
                 if 'openpose' in output_format:

Pose2Sim/triangulation.py

@@ -51,8 +51,8 @@
 import logging
 
 from Pose2Sim.common import retrieve_calib_params, computeP, weighted_triangulation, \
-    reprojection, euclidean_distance, sort_stringlist_by_last_number, \
-    min_with_single_indices, zup2yup, convert_to_c3d
+    reprojection, euclidean_distance, sort_people_sports2d, \
+    sort_stringlist_by_last_number, min_with_single_indices, zup2yup, convert_to_c3d
 from Pose2Sim.skeletons import *
 
 
@@ -120,53 +120,6 @@ def count_persons_in_json(file_path):
         return len(data.get('people', []))
 
 
-def sort_people(Q_kpt_old, Q_kpt):
-    '''
-    Associate persons across frames
-    Persons' indices are sometimes swapped when changing frame
-    A person is associated to another in the next frame when they are at a small distance
-    
-    INPUTS:
-    - Q_kpt_old: list of arrays of 3D coordinates [X, Y, Z, 1.] for the previous frame
-    - Q_kpt: idem Q_kpt_old, for current frame
-    
-    OUTPUT:
-    - Q_kpt_new: array with reordered persons
-    - personsIDs_sorted: index of reordered persons
-    '''
-
-    # Generate possible person correspondences across frames
-    if len(Q_kpt_old) < len(Q_kpt):
-        Q_kpt_old = np.concatenate((Q_kpt_old, [[0., 0., 0., 1.]]*(len(Q_kpt)-len(Q_kpt_old))))
-    if len(Q_kpt) < len(Q_kpt_old):
-        Q_kpt = np.concatenate((Q_kpt, [[0., 0., 0., 1.]]*(len(Q_kpt_old)-len(Q_kpt))))
-    personsIDs_comb = sorted(list(it.product(range(len(Q_kpt_old)),range(len(Q_kpt)))))
-
-    # Compute distance between persons from one frame to another
-    frame_by_frame_dist = []
-    for comb in personsIDs_comb:
-        frame_by_frame_dist += [euclidean_distance(Q_kpt_old[comb[0]],Q_kpt[comb[1]])]
-    frame_by_frame_dist = np.mean(frame_by_frame_dist, axis=1)
-
-    # sort correspondences by distance
-    minL, _, associated_tuples = min_with_single_indices(frame_by_frame_dist, personsIDs_comb)
-    # print('Distances :', minL)
-
-    # associate 3D points to same index across frames, nan if no correspondence
-    Q_kpt_new, personsIDs_sorted = [], []
-    for i in range(len(Q_kpt_old)):
-        id_in_old =  associated_tuples[:,1][associated_tuples[:,0] == i].tolist()
-        # print('id_in_old ', i, id_in_old)
-        if len(id_in_old) > 0:
-            personsIDs_sorted += id_in_old
-            Q_kpt_new += [Q_kpt[id_in_old[0]]]
-        else:
-            personsIDs_sorted += [-1]
-            Q_kpt_new += [Q_kpt_old[i]]
-
-    return Q_kpt_new, personsIDs_sorted, associated_tuples
-
-
 def make_trc(config_dict, Q, keypoints_names, f_range, id_person=-1):
     '''
     Make Opensim compatible trc file from a dataframe with 3D coordinates
@@ -840,19 +793,20 @@ def triangulate_all(config_dict):
             # reID persons across frames by checking the distance from one frame to another
             # print('Q before ordering ', np.array(Q)[:,:2])
             if f !=0:
-                Q, personsIDs_sorted, associated_tuples = sort_people(Q_old, Q)
+                Q, associated_tuples = sort_people_sports2d(Q_old, Q)
+                # Q, personsIDs_sorted, associated_tuples = sort_people(Q_old, Q)
                 # print('Q after ordering ', personsIDs_sorted, associated_tuples, np.array(Q)[:,:2])
 
                 error_sorted, nb_cams_excluded_sorted, id_excluded_cams_sorted = [], [], []
                 for i in range(len(Q)):
                     id_in_old =  associated_tuples[:,1][associated_tuples[:,0] == i].tolist()
                     if len(id_in_old) > 0:
-                        personsIDs_sorted += id_in_old
+                        # personsIDs_sorted += id_in_old
                         error_sorted += [error[id_in_old[0]]]
                         nb_cams_excluded_sorted += [nb_cams_excluded[id_in_old[0]]]
                         id_excluded_cams_sorted += [id_excluded_cams[id_in_old[0]]]
                     else:
-                        personsIDs_sorted += [-1]
+                        # personsIDs_sorted += [-1]
                         error_sorted += [error[i]]
                         nb_cams_excluded_sorted += [nb_cams_excluded[i]]
                         id_excluded_cams_sorted += [id_excluded_cams[i]]