sports2d 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

Sports2D/Utilities/common.py

@@ -20,14 +20,18 @@ import sys
 import toml
 import subprocess
 from pathlib import Path
+import itertools as it
 import logging
+from anytree import PreOrderIter
 
 import numpy as np
 import pandas as pd
 from scipy import interpolate
 import imageio_ffmpeg as ffmpeg
 import cv2
+import c3d
 
+import matplotlib.pyplot as plt
 from PyQt5.QtWidgets import QMainWindow, QApplication, QWidget, QTabWidget, QVBoxLayout
 from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
 from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
@@ -81,6 +85,32 @@ angle_dict = { # lowercase!
     'left hand': [['LIndex', 'LWrist'], 'horizontal', 0, -1]
     }
 
+marker_Z_positions = {'right':
+                        {"RHip": 0.105, "RKnee": 0.0886, "RAnkle": 0.0972, "RBigToe":0.0766, "RHeel":0.0883, "RSmallToe": 0.1200, 
+                        "RShoulder": 0.2016, "RElbow": 0.1613, "RWrist": 0.120, "RThumb": 0.1625, "RIndex": 0.1735, "RPinky": 0.1740, "REye": 0.0311, 
+                        "LHip": -0.105, "LKnee": -0.0886, "LAnkle": -0.0972, "LBigToe": -0.0766, "LHeel": -0.0883, "LSmallToe": -0.1200, 
+                        "LShoulder": -0.2016, "LElbow": -0.1613, "LWrist": -0.120, "LThumb": -0.1625, "LIndex": -0.1735, "LPinky": -0.1740, "LEye": -0.0311, 
+                        "Hip": 0.0, "Neck": 0.0, "Head":0.0, "Nose": 0.0},
+                    'left':
+                        {"RHip": -0.105, "RKnee": -0.0886, "RAnkle": -0.0972, "RBigToe": -0.0766, "RHeel": -0.0883, "RSmallToe": -0.1200, 
+                        "RShoulder": -0.2016, "RElbow": -0.1613, "RWrist": -0.120, "RThumb": -0.1625, "RIndex": -0.1735, "RPinky": -0.1740, "REye": -0.0311, 
+                        "LHip": 0.105, "LKnee": 0.0886, "LAnkle": 0.0972, "LBigToe":0.0766, "LHeel":0.0883, "LSmallToe": 0.1200, 
+                        "LShoulder": 0.2016, "LElbow": 0.1613, "LWrist": 0.120, "LThumb": 0.1625, "LIndex": 0.1735, "LPinky": 0.1740, "LEye": 0.0311, 
+                        "Hip": 0.0, "Neck": 0.0, "Head":0.0, "Nose": 0.0},
+                    'front':
+                        {"RHip": 0.0301, "RKnee": 0.0179, "RAnkle": 0.0230, "RBigToe": 0.2179, "RHeel": -0.0119, "RSmallToe": 0.1804, 
+                        "RShoulder": -0.01275, "RElbow": 0.0119, "RWrist": 0.0002, "RThumb": 0.0106, "RIndex": -0.0004, "RPinky": -0.0009, "REye": 0.0702, 
+                        "LHip": -0.0301, "LKnee": -0.0179, "LAnkle": 0.0230, "LBigToe": 0.2179, "LHeel": -0.0119, "LSmallToe": 0.1804, 
+                        "LShoulder": 0.01275, "LElbow": -0.0119, "LWrist": -0.0002, "LThumb": -0.0106, "LIndex": 0.0004, "LPinky": 0.0009, "LEye": -0.0702, 
+                        "Hip": 0.0301, "Neck": -0.0008, "Head": 0.0655, "Nose": 0.1076},
+                    'back':
+                        {"RHip": -0.0301, "RKnee": -0.0179, "RAnkle": -0.0230, "RBigToe": -0.2179, "RHeel": 0.0119, "RSmallToe": -0.1804, 
+                        "RShoulder": 0.01275, "RElbow": -0.0119, "RWrist": -0.0002, "RThumb": -0.0106, "RIndex": 0.0004, "RPinky": 0.0009, "REye": -0.0702, 
+                        "LHip": 0.0301, "LKnee": 0.0179, "LAnkle": -0.0230, "LBigToe": -0.2179, "LHeel": 0.0119, "LSmallToe": -0.1804, 
+                        "LShoulder": -0.01275, "LElbow": 0.0119, "LWrist": 0.0002, "LThumb": 0.0106, "LIndex": -0.0004, "LPinky": -0.0009, "LEye": 0.0702, 
+                        "Hip": 0.0301, "Neck": -0.0008, "Head": -0.0655, "Nose": 0.1076},
+            }
+
 colors = [(255, 0, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (0, 0, 0), (255, 255, 255),
             (125, 0, 0), (0, 125, 0), (0, 0, 125), (125, 125, 0), (125, 0, 125), (0, 125, 125), 
             (255, 125, 125), (125, 255, 125), (125, 125, 255), (255, 255, 125), (255, 125, 255), (125, 255, 255), (125, 125, 125),
@@ -169,6 +199,85 @@ def read_trc(trc_path):
         raise ValueError(f"Error reading TRC file at {trc_path}: {e}")
 
 
+def extract_trc_data(trc_path):
+    '''
+    Extract marker names and coordinates from a trc file.
+
+    INPUTS:
+    - trc_path: Path to the trc file
+
+    OUTPUTS:
+    - marker_names: List of marker names
+    - marker_coords: Array of marker coordinates (n_frames, t+3*n_markers)
+    '''
+
+    # marker names
+    with open(trc_path, 'r') as file:
+        lines = file.readlines()
+        marker_names_line = lines[3]
+        marker_names = marker_names_line.strip().split('\t')[2::3]
+
+    # time and marker coordinates
+    trc_data_np = np.genfromtxt(trc_path, skip_header=5, delimiter = '\t')[:,1:] 
+
+    return marker_names, trc_data_np
+
+
+def create_c3d_file(c3d_path, marker_names, trc_data_np):
+    '''
+    Create a c3d file from the data extracted from a trc file.
+
+    INPUTS:
+    - c3d_path: Path to the c3d file
+    - marker_names: List of marker names
+    - trc_data_np: Array of marker coordinates (n_frames, t+3*n_markers)
+
+    OUTPUTS:
+    - c3d file
+    '''
+
+    # retrieve frame rate
+    times = trc_data_np[:,0]
+    frame_rate = round((len(times)-1) / (times[-1] - times[0]))
+
+    # write c3d file
+    writer = c3d.Writer(point_rate=frame_rate, analog_rate=0, point_scale=1.0, point_units='mm', gen_scale=-1.0)
+    writer.set_point_labels(marker_names)
+    writer.set_screen_axis(X='+Z', Y='+Y')
+    
+    for frame in trc_data_np:
+        residuals = np.full((len(marker_names), 1), 0.0)
+        cameras = np.zeros((len(marker_names), 1))
+        coords = frame[1:].reshape(-1,3)*1000
+        points = np.hstack((coords, residuals, cameras))
+        writer.add_frames([(points, np.array([]))])
+
+    writer.set_start_frame(0)
+    writer._set_last_frame(len(trc_data_np)-1)
+
+    with open(c3d_path, 'wb') as handle:
+        writer.write(handle)
+
+
+def convert_to_c3d(trc_path):
+    '''
+    Make Visual3D compatible c3d files from a trc path
+
+    INPUT:
+    - trc_path: string, trc file to convert
+
+    OUTPUT:
+    - c3d file
+    '''
+
+    trc_path = str(trc_path)
+    c3d_path = trc_path.replace('.trc', '.c3d')
+    marker_names, trc_data_np = extract_trc_data(trc_path)
+    create_c3d_file(c3d_path, marker_names, trc_data_np)
+
+    return c3d_path
+
+
 def interpolate_zeros_nans(col, *args):
     '''
     Interpolate missing points (of value zero),
@@ -466,7 +575,7 @@ def add_neck_hip_coords(kpt_name, p_X, p_Y, p_scores, kpt_ids, kpt_names):
     return p_X, p_Y, p_scores
 
 
-def best_coords_for_measurements(trc_data, keypoints_names, fastest_frames_to_remove_percent=0.2, close_to_zero_speed=0.2, large_hip_knee_angles=45):
+def best_coords_for_measurements(Q_coords, keypoints_names, fastest_frames_to_remove_percent=0.2, close_to_zero_speed=0.2, large_hip_knee_angles=45):
     '''
     Compute the best coordinates for measurements, after removing:
     - 20% fastest frames (may be outliers)
@@ -474,7 +583,7 @@ def best_coords_for_measurements(trc_data, keypoints_names, fastest_frames_to_re
     - frames when hip and knee angle below 45° (imprecise coordinates when person is crouching)
     
     INPUTS:
-    - trc_data: pd.DataFrame. The XYZ coordinates of each marker
+    - Q_coords: pd.DataFrame. The XYZ coordinates of each marker
     - keypoints_names: list. The list of marker names
     - fastest_frames_to_remove_percent: float
     - close_to_zero_speed: float (sum for all keypoints: about 50 px/frame or 0.2 m/frame)
@@ -482,44 +591,46 @@ def best_coords_for_measurements(trc_data, keypoints_names, fastest_frames_to_re
     - trimmed_extrema_percent
 
     OUTPUT:
-    - trc_data_low_speeds_low_angles: pd.DataFrame. The best coordinates for measurements
+    - Q_coords_low_speeds_low_angles: pd.DataFrame. The best coordinates for measurements
     '''
 
     # Add MidShoulder column
-    df_MidShoulder = pd.DataFrame((trc_data['RShoulder'].values + trc_data['LShoulder'].values) /2)
+    df_MidShoulder = pd.DataFrame((Q_coords['RShoulder'].values + Q_coords['LShoulder'].values) /2)
     df_MidShoulder.columns = ['MidShoulder']*3
-    trc_data = pd.concat((trc_data.reset_index(drop=True), df_MidShoulder), axis=1)
+    Q_coords = pd.concat((Q_coords.reset_index(drop=True), df_MidShoulder), axis=1)
 
     # Add Hip column if not present
     n_markers_init = len(keypoints_names)
     if 'Hip' not in keypoints_names:
-        df_Hip = pd.DataFrame((trc_data['RHip'].values + trc_data['LHip'].values) /2)
+        df_Hip = pd.DataFrame((Q_coords['RHip'].values + Q_coords['LHip'].values) /2)
         df_Hip.columns = ['Hip']*3
-        trc_data = pd.concat((trc_data.reset_index(drop=True), df_Hip), axis=1)
+        Q_coords = pd.concat((Q_coords.reset_index(drop=True), df_Hip), axis=1)
     n_markers = len(keypoints_names)
 
     # Using 80% slowest frames
-    sum_speeds = pd.Series(np.nansum([np.linalg.norm(trc_data.iloc[:,kpt:kpt+3].diff(), axis=1) for kpt in range(n_markers)], axis=0))
+    sum_speeds = pd.Series(np.nansum([np.linalg.norm(Q_coords.iloc[:,kpt:kpt+3].diff(), axis=1) for kpt in range(n_markers)], axis=0))
     sum_speeds = sum_speeds[sum_speeds>close_to_zero_speed] # Removing when speeds close to zero (out of frame)
     if len(sum_speeds)==0:
-        raise ValueError('All frames have speed close to zero. Make sure the person is moving and correctly detected, or change close_to_zero_speed to a lower value.')
-    min_speed_indices = sum_speeds.abs().nsmallest(int(len(sum_speeds) * (1-fastest_frames_to_remove_percent))).index
-    trc_data_low_speeds = trc_data.iloc[min_speed_indices].reset_index(drop=True)
+        logging.warning('All frames have speed close to zero. Make sure the person is moving and correctly detected, or change close_to_zero_speed to a lower value. Not restricting the speeds to be above any threshold.')
+        Q_coords_low_speeds = Q_coords
+    else:
+        min_speed_indices = sum_speeds.abs().nsmallest(int(len(sum_speeds) * (1-fastest_frames_to_remove_percent))).index
+        Q_coords_low_speeds = Q_coords.iloc[min_speed_indices].reset_index(drop=True)
     
     # Only keep frames with hip and knee flexion angles below 45% 
     # (if more than 50 of them, else take 50 smallest values)
     try:
-        ang_mean = mean_angles(trc_data_low_speeds, ang_to_consider = ['right knee', 'left knee', 'right hip', 'left hip'])
-        trc_data_low_speeds_low_angles = trc_data_low_speeds[ang_mean < large_hip_knee_angles]
-        if len(trc_data_low_speeds_low_angles) < 50:
-            trc_data_low_speeds_low_angles = trc_data_low_speeds.iloc[pd.Series(ang_mean).nsmallest(50).index]
+        ang_mean = mean_angles(Q_coords_low_speeds, ang_to_consider = ['right knee', 'left knee', 'right hip', 'left hip'])
+        Q_coords_low_speeds_low_angles = Q_coords_low_speeds[ang_mean < large_hip_knee_angles]
+        if len(Q_coords_low_speeds_low_angles) < 50:
+            Q_coords_low_speeds_low_angles = Q_coords_low_speeds.iloc[pd.Series(ang_mean).nsmallest(50).index]
     except:
-        logging.warning(f"At least one among the RAnkle, RKnee, RHip, RShoulder, LAnkle, LKnee, LHip, LShoulder markers is missing for computing the knee and hip angles. Not restricting these agles to be below {large_hip_knee_angles}°.")
+        logging.warning(f"At least one among the RAnkle, RKnee, RHip, RShoulder, LAnkle, LKnee, LHip, LShoulder markers is missing for computing the knee and hip angles. Not restricting these angles to be below {large_hip_knee_angles}°.")
 
     if n_markers_init < n_markers:
-        trc_data_low_speeds_low_angles = trc_data_low_speeds_low_angles.iloc[:,:-3]
+        Q_coords_low_speeds_low_angles = Q_coords_low_speeds_low_angles.iloc[:,:-3]
 
-    return trc_data_low_speeds_low_angles
+    return Q_coords_low_speeds_low_angles
 
 
 def compute_height(trc_data, keypoints_names, fastest_frames_to_remove_percent=0.1, close_to_zero_speed=50, large_hip_knee_angles=45, trimmed_extrema_percent=0.5):
@@ -547,7 +658,7 @@ def compute_height(trc_data, keypoints_names, fastest_frames_to_remove_percent=0
     try:
         rfoot, lfoot = [euclidean_distance(trc_data_low_speeds_low_angles[pair[0]],trc_data_low_speeds_low_angles[pair[1]]) for pair in feet_pairs]
     except:
-        rfoot, lfoot = 10, 10
+        rfoot, lfoot = 0.10, 0.10
         logging.warning('The Heel marker is missing from your model. Considering Foot to Heel size as 10 cm.')
 
     ankle_to_shoulder_pairs =  [['RAnkle', 'RKnee'], ['RKnee', 'RHip'], ['RHip', 'RShoulder'],
@@ -688,4 +799,350 @@ def write_calibration(calib_params, toml_path):
             fish_str = f'fisheye = false\n\n'
             cal_f.write(cam_str + name_str + size_str + mat_str + dist_str + rot_str + tran_str + fish_str)
         meta = '[metadata]\nadjusted = false\nerror = 0.0\n'
-        cal_f.write(meta)
+        cal_f.write(meta)
+
+    
+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 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 (indices of best combinations)
+    - 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_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 sort_people_rtmlib(pose_tracker, keypoints, scores):
+    '''
+    Associate persons across frames (RTMLib method)
+
+    INPUTS:
+    - pose_tracker: PoseTracker. The initialized RTMLib pose tracker object
+    - 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,
+    L the confidence of detected keypoints
+
+    OUTPUT:
+    - sorted_keypoints: array with reordered persons
+    - sorted_scores: array with reordered scores
+    '''
+    
+    try:
+        desired_size = max(pose_tracker.track_ids_last_frame)+1
+        sorted_keypoints = np.full((desired_size, keypoints.shape[1], 2), np.nan)
+        sorted_keypoints[pose_tracker.track_ids_last_frame] = keypoints[:len(pose_tracker.track_ids_last_frame), :, :]
+        sorted_scores = np.full((desired_size, scores.shape[1]), np.nan)
+        sorted_scores[pose_tracker.track_ids_last_frame] = scores[:len(pose_tracker.track_ids_last_frame), :]
+    except:
+        sorted_keypoints, sorted_scores = keypoints, scores
+
+    return sorted_keypoints, sorted_scores
+
+
+def sort_people_deepsort(keypoints, scores, deepsort_tracker, frame,frame_count):
+    '''
+    Associate persons across frames (DeepSort method)
+
+    INPUTS:
+    - 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,
+    L the confidence of detected keypoints
+    - deepsort_tracker: The initialized DeepSort tracker object
+    - frame: np.array. The current image opened with cv2.imread
+
+    OUTPUT:
+    - sorted_keypoints: array with reordered persons
+    - sorted_scores: array with reordered scores
+    '''
+
+    try:
+        # Compute bboxes from keypoints and create detections (bboxes, scores, class_ids)
+        bboxes_ltwh = bbox_ltwh_compute(keypoints, padding=20)
+        bbox_scores = np.mean(scores, axis=1)
+        class_ids = np.array(['person']*len(bboxes_ltwh))
+        detections = list(zip(bboxes_ltwh, bbox_scores, class_ids))
+
+        # Estimates the tracks and retrieve indexes of the original detections
+        det_ids = [i for i in range(len(detections))]
+        tracks = deepsort_tracker.update_tracks(detections, frame=frame, others=det_ids)
+        track_ids_frame, orig_det_ids = [], []
+        for track in tracks:
+            if not track.is_confirmed():
+                continue
+            track_ids_frame.append(int(track.track_id)-1)       # ID of people
+            orig_det_ids.append(track.get_det_supplementary())  # ID of detections
+
+        # Correspondence between person IDs and original detection IDs
+        desired_size = max(track_ids_frame) + 1
+        sorted_keypoints = np.full((desired_size, keypoints.shape[1], 2), np.nan)
+        sorted_scores = np.full((desired_size, scores.shape[1]), np.nan)
+        for i,v in enumerate(track_ids_frame):
+            if orig_det_ids[i] is not None:    
+                sorted_keypoints[v] = keypoints[orig_det_ids[i]]
+                sorted_scores[v] = scores[orig_det_ids[i]]
+
+    except Exception as e:
+        sorted_keypoints, sorted_scores = keypoints, scores
+        if frame_count > deepsort_tracker.tracker.n_init:
+            logging.warning(f"Tracking error: {e}. Sorting persons with DeepSort method failed for this frame.")
+
+    return sorted_keypoints, sorted_scores
+
+
+def bbox_ltwh_compute(keypoints, padding=0):
+    '''
+    Compute bounding boxes in (x_min, y_min, width, height) format
+    Optionally add padding to the bounding boxes 
+    as a percentage of the bounding box size (+padding% horizontally, +padding/2% vertically)
+
+    INPUTS:
+    - keypoints: array of shape K, L, M with K the number of detected persons,
+                    L the number of detected keypoints, M their 2D coordinates
+    - padding: int. The padding to add to the bounding boxes, in perceptage
+    '''
+
+    x_coords = keypoints[:, :, 0]
+    y_coords = keypoints[:, :, 1]
+
+    x_min, x_max = np.min(x_coords, axis=1), np.max(x_coords, axis=1)
+    y_min, y_max = np.min(y_coords, axis=1), np.max(y_coords, axis=1)
+    width = x_max - x_min
+    height = y_max - y_min
+
+    if padding > 0:
+        x_min = x_min - width*padding/100
+        y_min = y_min - height/2*padding/100
+        width = width + 2*width*padding/100
+        height = height + height*padding/100
+
+    bbox_ltwh = np.stack((x_min, y_min, width, height), axis=1)
+
+    return bbox_ltwh
+
+
+def draw_bounding_box(img, X, Y, colors=[(255, 0, 0), (0, 255, 0), (0, 0, 255)], fontSize=0.3, thickness=1):
+    '''
+    Draw bounding boxes and person ID around list of lists of X and Y coordinates.
+    Bounding boxes have a different color for each person.
+    
+    INPUTS:
+    - img: opencv image
+    - X: list of list of x coordinates
+    - Y: list of list of y coordinates
+    - colors: list of colors to cycle through
+    
+    OUTPUT:
+    - img: image with rectangles and person IDs
+    '''
+   
+    color_cycle = it.cycle(colors)
+
+    for i,(x,y) in enumerate(zip(X,Y)):
+        color = next(color_cycle)
+        if not np.isnan(x).all():
+            x_min, y_min = np.nanmin(x).astype(int), np.nanmin(y).astype(int)
+            x_max, y_max = np.nanmax(x).astype(int), np.nanmax(y).astype(int)
+            if x_min < 0: x_min = 0
+            if x_max > img.shape[1]: x_max = img.shape[1]
+            if y_min < 0: y_min = 0
+            if y_max > img.shape[0]: y_max = img.shape[0]
+
+            # Draw rectangles
+            cv2.rectangle(img, (x_min-25, y_min-25), (x_max+25, y_max+25), color, thickness) 
+        
+            # Write person ID
+            cv2.putText(img, str(i), (x_min-30, y_min-30), cv2.FONT_HERSHEY_SIMPLEX, fontSize, color, 2, cv2.LINE_AA) 
+    
+    return img
+
+
+def draw_skel(img, X, Y, model):
+    '''
+    Draws keypoints and skeleton for each person.
+    Skeletons have a different color for each person.
+
+    INPUTS:
+    - img: opencv image
+    - X: list of list of x coordinates
+    - Y: list of list of y coordinates
+    - model: skeleton model (from skeletons.py)
+    - colors: list of colors to cycle through
+    
+    OUTPUT:
+    - img: image with keypoints and skeleton
+    '''
+    
+    # Get (unique) pairs between which to draw a line
+    id_pairs, name_pairs = [], []
+    for data_i in PreOrderIter(model.root, filter_=lambda node: node.is_leaf):
+        node_branch_ids = [node_i.id for node_i in data_i.path]
+        node_branch_names = [node_i.name for node_i in data_i.path]
+        id_pairs += [[node_branch_ids[i],node_branch_ids[i+1]] for i in range(len(node_branch_ids)-1)]
+        name_pairs += [[node_branch_names[i],node_branch_names[i+1]] for i in range(len(node_branch_names)-1)]
+    node_pairs = {tuple(name_pair): id_pair for (name_pair,id_pair) in zip(name_pairs,id_pairs)}
+
+    
+    # Draw lines
+    for (x,y) in zip(X,Y):
+        if not np.isnan(x).all():
+            for names, ids in node_pairs.items():
+                if not None in ids and not (np.isnan(x[ids[0]]) or np.isnan(y[ids[0]]) or np.isnan(x[ids[1]]) or np.isnan(y[ids[1]])):
+                    if any(n.startswith('R') for n in names) and not any(n.startswith('L') for n in names):
+                        c = (255,128,0)
+                    elif any(n.startswith('L') for n in names) and not any(n.startswith('R') for n in names):
+                        c = (0,255,0)
+                    else:
+                        c = (51, 153, 255)
+                    cv2.line(img, (int(x[ids[0]]), int(y[ids[0]])), (int(x[ids[1]]), int(y[ids[1]])), c, thickness)
+
+    return img
+
+
+def draw_keypts(img, X, Y, scores, cmap_str='RdYlGn'):
+    '''
+    Draws keypoints and skeleton for each person.
+    Keypoints' colors depend on their score.
+
+    INPUTS:
+    - img: opencv image
+    - X: list of list of x coordinates
+    - Y: list of list of y coordinates
+    - scores: list of list of scores
+    - cmap_str: colormap name
+    
+    OUTPUT:
+    - img: image with keypoints and skeleton
+    '''
+    
+    scores = np.where(np.isnan(scores), 0, scores)
+    # scores = (scores - 0.4) / (1-0.4) # to get a red color for scores lower than 0.4
+    scores = np.where(scores>0.99, 0.99, scores)
+    scores = np.where(scores<0, 0, scores)
+    
+    cmap = plt.get_cmap(cmap_str)
+    for (x,y,s) in zip(X,Y,scores):
+        c_k = np.array(cmap(s))[:,:-1]*255
+        [cv2.circle(img, (int(x[i]), int(y[i])), thickness+4, c_k[i][::-1], -1)
+            for i in range(len(x))
+            if not (np.isnan(x[i]) or np.isnan(y[i]))]
+
+    return img

Sports2D/Utilities/skeletons.py

@@ -85,10 +85,10 @@ HALPE_26 = Node("Hip", id=19, children=[
 ])
 
 
-'''COCO_133_wrist (full-body with hands and face, from AlphaPose, MMPose, etc.)
+'''COCO_133_WRIST (full-body with hands and face, from AlphaPose, MMPose, etc.)
 https://github.com/MVIG-SJTU/AlphaPose/blob/master/docs/MODEL_ZOO.md
 https://github.com/open-mmlab/mmpose/tree/main/projects/rtmpose'''
-COCO_133_wrist = Node("CHip", id=None, children=[
+COCO_133_WRIST = Node("Hip", id=None, children=[
     Node("RHip", id=12, children=[
         Node("RKnee", id=14, children=[
             Node("RAnkle", id=16, children=[
@@ -139,7 +139,7 @@ COCO_133_wrist = Node("CHip", id=None, children=[
 '''COCO_133 (full-body with hands and face, from AlphaPose, MMPose, etc.)
 https://github.com/MVIG-SJTU/AlphaPose/blob/master/docs/MODEL_ZOO.md
 https://github.com/open-mmlab/mmpose/tree/main/projects/rtmpose'''
-COCO_133 = Node("CHip", id=None, children=[
+COCO_133 = Node("Hip", id=None, children=[
     Node("RHip", id=12, children=[
         Node("RKnee", id=14, children=[
             Node("RAnkle", id=16, children=[
@@ -359,9 +359,7 @@ COCO_133 = Node("CHip", id=None, children=[
                                                                         Node("Mouth17", id=87, children=[
                                                                             Node("Mouth18", id=88, children=[
                                                                                 Node("Mouth19", id=89, children=[
-                                                                                    Node("Mouth20", id=90, children=[
-                                                                                        Node("Mouth21", id=91)
-                                                                                    ]),
+                                                                                    Node("Mouth20", id=90)
                                                                                 ]),
                                                                             ]),
                                                                         ]),
@@ -387,7 +385,7 @@ COCO_133 = Node("CHip", id=None, children=[
 
 '''COCO_17 (full-body without hands and feet, from OpenPose, AlphaPose, OpenPifPaf, YOLO-pose, MMPose, etc.)
 https://github.com/open-mmlab/mmpose/tree/main/projects/rtmpose'''
-COCO_17 = Node("CHip", id=None, children=[
+COCO_17 = Node("Hip", id=None, children=[
     Node("RHip", id=12, children=[
         Node("RKnee", id=14, children=[
             Node("RAnkle", id=16),
@@ -645,9 +643,10 @@ FACE_106 = Node("root", id=None, children=[
     ]),
 ])
 
+
 '''ANIMAL2D_17 (full-body animal)
 https://github.com/AlexTheBad/AP-10K/'''
-ANIMAL2D_17 = Node("CHip", id=4, children=[
+ANIMAL2D_17 = Node("Hip", id=4, children=[
     Node("RHip", id=14, children=[
         Node("RKnee", id=15, children=[
             Node("RAnkle", id=16),

Sports2D/Utilities/tests.py

@@ -56,11 +56,11 @@ def test_workflow():
 
     # Default
     demo_cmd = ["sports2d", "--show_realtime_results", "False", "--show_graphs", "False"]
-    subprocess.run(demo_cmd, check=True, capture_output=True, text=True)
+    subprocess.run(demo_cmd, check=True, capture_output=True, text=True, encoding='utf-8')
 
     # With no pixels to meters conversion, no multiperson, lightweight mode, detection frequency, time range and slowmo factor
     demo_cmd2 = ["sports2d", "--to_meters", "False", "--multiperson", "False", "--mode", "lightweight", "--det_frequency", "50", "--time_range", "1.2", "2.7",  "--slowmo_factor", "4", "--show_realtime_results", "False", "--show_graphs", "False"]
-    subprocess.run(demo_cmd2, check=True, capture_output=True, text=True)
+    subprocess.run(demo_cmd2, check=True, capture_output=True, text=True, encoding='utf-8')
     
     # With inverse kinematics, body pose_model and custom RTMO mode
     # demo_cmd3 = ["sports2d", "--do_ik", "--person_orientation", "front none left", "--pose_model", "body", "--mode", "{'pose_class':'RTMO', 'pose_model':'https://download.openmmlab.com/mmpose/v1/projects/rtmo/onnx_sdk/rtmo-m_16xb16-600e_body7-640x640-39e78cc4_20231211.zip', 'pose_input_size':[640, 640]}", "--show_realtime_results", "False", "--show_graphs", "False"]
@@ -74,4 +74,4 @@ def test_workflow():
     with open(cli_config_path, 'w') as f: toml.dump(config_dict, f)
 
     demo_cmd4 = ["sports2d", "--config", str(cli_config_path), "--show_realtime_results", "False", "--show_graphs", "False"]
-    subprocess.run(demo_cmd4, check=True, capture_output=True, text=True)
+    subprocess.run(demo_cmd4, check=True, capture_output=True, text=True, encoding='utf-8')