tracking on 2D for future multi-person synchronization

Pose2Sim/common.py

@@ -213,6 +213,52 @@ def reprojection(P_all, Q):
     return x_calc, y_calc
 
 
+def min_with_single_indices(L, T):
+    '''
+    Let L be a list (size s) with T associated tuple indices (size s).
+    Select the smallest values of L, considering that 
+    the next smallest value cannot have the same numbers 
+    in the associated tuple as any of the previous ones.
+
+    Example:
+    L = [  20,   27,  51,    33,   43,   23,   37,   24,   4,   68,   84,    3  ]
+    T = list(it.product(range(2),range(3)))
+      = [(0,0),(0,1),(0,2),(0,3),(1,0),(1,1),(1,2),(1,3),(2,0),(2,1),(2,2),(2,3)]
+
+    - 1st smallest value: 3 with tuple (2,3), index 11
+    - 2nd smallest value when excluding indices (2,.) and (.,3), i.e. [(0,0),(0,1),(0,2),X,(1,0),(1,1),(1,2),X,X,X,X,X]:
+    20 with tuple (0,0), index 0
+    - 3rd smallest value when excluding [X,X,X,X,X,(1,1),(1,2),X,X,X,X,X]:
+    23 with tuple (1,1), index 5
+    
+    INPUTS:
+    - L: list (size s)
+    - T: T associated tuple indices (size s)
+
+    OUTPUTS: 
+    - minL: list of smallest values of L, considering constraints on tuple indices
+    - argminL: list of indices of smallest values of L
+    - T_minL: list of tuples associated with smallest values of L
+    '''
+
+    minL = [np.nanmin(L)]
+    argminL = [np.nanargmin(L)]
+    T_minL = [T[argminL[0]]]
+    
+    mask_tokeep = np.array([True for t in T])
+    i=0
+    while mask_tokeep.any()==True:
+        mask_tokeep = mask_tokeep & np.array([t[0]!=T_minL[i][0] and t[1]!=T_minL[i][1] for t in T])
+        if mask_tokeep.any()==True:
+            indicesL_tokeep = np.where(mask_tokeep)[0]
+            minL += [np.nanmin(np.array(L)[indicesL_tokeep]) if not np.isnan(np.array(L)[indicesL_tokeep]).all() else np.nan]
+            argminL += [indicesL_tokeep[np.nanargmin(np.array(L)[indicesL_tokeep])] if not np.isnan(minL[-1]) else indicesL_tokeep[0]]
+            T_minL += (T[argminL[i+1]],)
+            i+=1
+    
+    return np.array(minL), np.array(argminL), np.array(T_minL)
+
+
 def euclidean_distance(q1, q2):
     '''
     Euclidean distance between 2 points (N-dim).

Pose2Sim/poseEstimation.py

@@ -36,12 +36,13 @@ import os
 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
+from Pose2Sim.common import natural_sort_key, min_with_single_indices, euclidean_distance
 
 
 ## AUTHORSHIP INFORMATION
@@ -99,35 +100,62 @@ def save_to_openpose(json_file_path, keypoints, scores):
         json.dump(json_output, json_file)
 
    
-def sort_people_rtmlib(pose_tracker, keypoints, scores):
+def sort_people_sports2d(keyptpre, keypt, scores):
     '''
-    Associate persons across frames (RTMLib method)
+    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:
-    - pose_tracker: PoseTracker. The initialized RTMLib pose tracker object
+    - keyptpre: array of shape K, L, M with K the number of detected persons,
-    - keypoints: array of shape K, L, M with K the number of detected persons,
     L the number of detected keypoints, M their 2D coordinates
-    - scores: array of shape K, L with K the number of detected persons,
+    - 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
     
-    OUTPUT:
+    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
     '''
     
-    try:
+    # Generate possible person correspondences across frames
-        desired_size = max(pose_tracker.track_ids_last_frame)+1
+    if len(keyptpre) < len(keypt):
-        sorted_keypoints = np.full((desired_size, keypoints.shape[1], 2), np.nan)
+        keyptpre = np.concatenate((keyptpre, np.full((len(keypt)-len(keyptpre), keypt.shape[1], 2), np.nan)))
-        sorted_keypoints[pose_tracker.track_ids_last_frame] = keypoints[:len(pose_tracker.track_ids_last_frame), :, :]
+    if len(keypt) < len(keyptpre):
-        sorted_scores = np.full((desired_size, scores.shape[1]), np.nan)
+        keypt = np.concatenate((keypt, np.full((len(keyptpre)-len(keypt), keypt.shape[1], 2), np.nan)))
-        sorted_scores[pose_tracker.track_ids_last_frame] = scores[:len(pose_tracker.track_ids_last_frame), :]
+        scores = np.concatenate((scores, np.full((len(keyptpre)-len(scores), scores.shape[1]), np.nan)))
-    except:
+    personsIDs_comb = sorted(list(it.product(range(len(keyptpre)), range(len(keypt)))))
-        sorted_keypoints, sorted_scores = keypoints, scores
     
-    return sorted_keypoints, sorted_scores
+    # 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]])]
+    
+    # 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):
+def process_video(video_path, pose_tracker, output_format, save_video, save_images, display_detection, frame_range, multi_person):
     '''
     Estimate pose from a video file
     
@@ -185,6 +213,11 @@ def process_video(video_path, pose_tracker, output_format, save_video, save_imag
                 # Perform pose estimation on the frame
                 keypoints, scores = pose_tracker(frame)
 
+                # 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)
+           
                 # Save to json
                 if 'openpose' in output_format:
                     json_file_path = os.path.join(json_output_dir, f'{video_name_wo_ext}_{frame_idx:06d}.json')
@@ -220,7 +253,7 @@ def process_video(video_path, pose_tracker, output_format, save_video, save_imag
         cv2.destroyAllWindows()
 
 
-def process_images(image_folder_path, vid_img_extension, pose_tracker, output_format, fps, save_video, save_images, display_detection, frame_range):
+def process_images(image_folder_path, vid_img_extension, pose_tracker, output_format, fps, save_video, save_images, display_detection, frame_range, multi_person):
     '''
     Estimate pose estimation from a folder of images
     
@@ -270,6 +303,11 @@ def process_images(image_folder_path, vid_img_extension, pose_tracker, output_fo
             # Perform pose estimation on the image
             keypoints, scores = pose_tracker(frame)
 
+            # 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)
+            
             # Extract frame number from the filename
             if 'openpose' in output_format:
                 json_file_path = os.path.join(json_output_dir, f"{os.path.splitext(os.path.basename(image_file))[0]}_{frame_idx:06d}.json")
@@ -332,6 +370,7 @@ def rtm_estimator(config_dict):
     # if single trial
     session_dir = session_dir if 'Config.toml' in os.listdir(session_dir) else os.getcwd()
     frame_range = config_dict.get('project').get('frame_range')
+    multi_person = config_dict.get('project').get('multi_person')
     video_dir = os.path.join(project_dir, 'videos')
     pose_dir = os.path.join(project_dir, 'pose')
 
@@ -432,7 +471,7 @@ def rtm_estimator(config_dict):
             logging.info(f'Found video files with extension {vid_img_extension}.')
             for video_path in video_files:
                 pose_tracker.reset()
-                process_video(video_path, pose_tracker, output_format, save_video, save_images, display_detection, frame_range)
+                process_video(video_path, pose_tracker, output_format, save_video, save_images, display_detection, frame_range, multi_person)
 
         else:
             # Process image folders
@@ -441,4 +480,4 @@ def rtm_estimator(config_dict):
             for image_folder in image_folders:
                 pose_tracker.reset()
                 image_folder_path = os.path.join(video_dir, image_folder)
-                process_images(image_folder_path, vid_img_extension, pose_tracker, output_format, frame_rate, save_video, save_images, display_detection, frame_range)
+                process_images(image_folder_path, vid_img_extension, pose_tracker, output_format, frame_rate, save_video, save_images, display_detection, frame_range, multi_person)

Pose2Sim/triangulation.py

@@ -51,7 +51,8 @@ from anytree.importer import DictImporter
 import logging
 
 from Pose2Sim.common import retrieve_calib_params, computeP, weighted_triangulation, \
-    reprojection, euclidean_distance, sort_stringlist_by_last_number, zup2yup, convert_to_c3d
+    reprojection, euclidean_distance, sort_stringlist_by_last_number, \
+    min_with_single_indices, zup2yup, convert_to_c3d
 from Pose2Sim.skeletons import *
 
 
@@ -119,52 +120,6 @@ def count_persons_in_json(file_path):
         return len(data.get('people', []))
     
 
-def min_with_single_indices(L, T):
-    '''
-    Let L be a list (size s) with T associated tuple indices (size s).
-    Select the smallest values of L, considering that 
-    the next smallest value cannot have the same numbers 
-    in the associated tuple as any of the previous ones.
-
-    Example:
-    L = [  20,   27,  51,    33,   43,   23,   37,   24,   4,   68,   84,    3  ]
-    T = list(it.product(range(2),range(3)))
-      = [(0,0),(0,1),(0,2),(0,3),(1,0),(1,1),(1,2),(1,3),(2,0),(2,1),(2,2),(2,3)]
-
-    - 1st smallest value: 3 with tuple (2,3), index 11
-    - 2nd smallest value when excluding indices (2,.) and (.,3), i.e. [(0,0),(0,1),(0,2),X,(1,0),(1,1),(1,2),X,X,X,X,X]:
-    20 with tuple (0,0), index 0
-    - 3rd smallest value when excluding [X,X,X,X,X,(1,1),(1,2),X,X,X,X,X]:
-    23 with tuple (1,1), index 5
-    
-    INPUTS:
-    - L: list (size s)
-    - T: T associated tuple indices (size s)
-
-    OUTPUTS: 
-    - minL: list of smallest values of L, considering constraints on tuple indices
-    - argminL: list of indices of smallest values of L
-    - T_minL: list of tuples associated with smallest values of L
-    '''
-
-    minL = [np.nanmin(L)]
-    argminL = [np.nanargmin(L)]
-    T_minL = [T[argminL[0]]]
-    
-    mask_tokeep = np.array([True for t in T])
-    i=0
-    while mask_tokeep.any()==True:
-        mask_tokeep = mask_tokeep & np.array([t[0]!=T_minL[i][0] and t[1]!=T_minL[i][1] for t in T])
-        if mask_tokeep.any()==True:
-            indicesL_tokeep = np.where(mask_tokeep)[0]
-            minL += [np.nanmin(np.array(L)[indicesL_tokeep]) if not np.isnan(np.array(L)[indicesL_tokeep]).all() else np.nan]
-            argminL += [indicesL_tokeep[np.nanargmin(np.array(L)[indicesL_tokeep])] if not np.isnan(minL[-1]) else indicesL_tokeep[0]]
-            T_minL += (T[argminL[i+1]],)
-            i+=1
-    
-    return np.array(minL), np.array(argminL), np.array(T_minL)
-
-
 def sort_people(Q_kpt_old, Q_kpt):
     '''
     Associate persons across frames

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = pose2sim
-version = 0.10.0
+version = 0.10.1
 author = David Pagnon
 author_email = contact@david-pagnon.com
 description = Perform a markerless kinematic analysis from multiple calibrated views as a unified workflow from an OpenPose input to an OpenSim result.