sports2d 0.6.1__tar.gz → 0.6.2__tar.gz

{sports2d-0.6.1 → sports2d-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sports2d
-Version: 0.6.1
+Version: 0.6.2
 Summary: Detect pose and compute 2D joint angles from a video.
 Home-page: https://github.com/davidpagnon/Sports2D
 Author: David Pagnon
@@ -33,10 +33,11 @@ Requires-Dist: opencv-python
 Requires-Dist: matplotlib
 Requires-Dist: PyQt5
 Requires-Dist: statsmodels
-Requires-Dist: rtmlib_pose2sim
+Requires-Dist: rtmlib
 Requires-Dist: openvino
 Requires-Dist: tqdm
 Requires-Dist: imageio_ffmpeg
+Requires-Dist: deep-sort-realtime
 
 
 [![Continuous integration](https://github.com/davidpagnon/sports2d/actions/workflows/continuous-integration.yml/badge.svg?branch=main)](https://github.com/davidpagnon/sports2d/actions/workflows/continuous-integration.yml)
@@ -212,6 +213,9 @@ Note that it does not take distortions into account, and that it will be less ac
   ``` cmd
   sports2d --multiperson false --pose_model Body --mode lightweight --det_frequency 50 
   ```
+  ``` cmd
+  sports2d --tracking_mode deepsort --deepsort_params """{'max_age':30, 'n_init':3, 'nms_max_overlap':0.8, 'max_cosine_distance':0.3, 'nn_budget':200, 'max_iou_distance':0.8, 'embedder_gpu': True}"""
+  ```
 <br>
 
 #### Run with a toml configuration file: 
@@ -249,6 +253,7 @@ Note that any detection and pose models can be used (first [deploy them with MMP
 - Use `--det_frequency 50`: Will detect poses only every 50 frames, and track keypoints in between, which is faster.
 - Use `--multiperson false`: Can be used if one single person is present in the video. Otherwise, persons' IDs may be mixed up.
 - Use `--load_trc <path_to_file_px.trc>`: Will use pose estimation results from a file. Useful if you want to use different parameters for pixel to meter conversion or angle calculation without running detection and pose estimation all over.
+- Use `--tracking_mode sports2d`: Will use the default Sports2D tracker. Unlike DeepSort, it is faster, does not require any parametrization, and is as good in non-crowded scenes. 
 
 <br> 
 
@@ -369,7 +374,7 @@ sports2d --time_range 1.2 2.7 --ik true --person_orientation front none left
 
 ### All the parameters
 
-Have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type for a full list of the available parameters:
+For a full list of the available parameters, have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type:
 
 ``` cmd
 sports2d --help
@@ -414,7 +419,10 @@ sports2d --help
 'osim_setup_path': ["", "path to OpenSim setup. '../OpenSim_setup' if not specified"],
 'person_orientation': ["", "front, back, left, right, auto, or none. 'front none left' if not specified. If 'auto', will be either left or right depending on the direction of the motion."],
 'close_to_zero_speed_m': ["","Sum for all keypoints: about 50 px/frame or 0.2 m/frame"], 
-'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],                'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],
+'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'deepsort_params': ["", 'Deepsort tracking parameters: """{dictionary between 3 double quotes}""". \n\
+                    More information there: https://github.com/levan92/deep_sort_realtime/blob/master/deep_sort_realtime/deepsort_tracker.py#L51'],     
 'input_size': ["", "width, height. 1280, 720 if not specified. Lower resolution will be faster but less precise"],
 'keypoint_likelihood_threshold': ["", "detected keypoints are not retained if likelihood is below this threshold. 0.3 if not specified"],
 'average_likelihood_threshold': ["", "detected persons are not retained if average keypoint likelihood is below this threshold. 0.5 if not specified"],
@@ -459,7 +467,7 @@ Sports2D:
 
 2. **Sets up pose estimation with RTMLib.** It can be run in lightweight, balanced, or performance mode, and for faster inference, keypoints can be tracked instead of detected for a certain number of frames. Any RTMPose model can be used. 
 
-3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. This carefully crafted `sports2d` tracker runs at a comparable speed as the RTMlib one but is much more robust. The user can still choose the RTMLib method if they need it by specifying it in the Config.toml file.
+3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. We crafted a 'sports2D' tracker which gives good results and runs in real time, but it is also possible to use `deepsort` in particularly challenging situations. 
 
 4. **Chooses the right persons to keep.** In single-person mode, only keeps the person with the highest average scores over the sequence. In multi-person mode, only retrieves the keypoints with high enough confidence, and only keeps the persons with high enough average confidence over each frame.
 

{sports2d-0.6.1 → sports2d-0.6.2}/README.md

@@ -172,6 +172,9 @@ Note that it does not take distortions into account, and that it will be less ac
   ``` cmd
   sports2d --multiperson false --pose_model Body --mode lightweight --det_frequency 50 
   ```
+  ``` cmd
+  sports2d --tracking_mode deepsort --deepsort_params """{'max_age':30, 'n_init':3, 'nms_max_overlap':0.8, 'max_cosine_distance':0.3, 'nn_budget':200, 'max_iou_distance':0.8, 'embedder_gpu': True}"""
+  ```
 <br>
 
 #### Run with a toml configuration file: 
@@ -209,6 +212,7 @@ Note that any detection and pose models can be used (first [deploy them with MMP
 - Use `--det_frequency 50`: Will detect poses only every 50 frames, and track keypoints in between, which is faster.
 - Use `--multiperson false`: Can be used if one single person is present in the video. Otherwise, persons' IDs may be mixed up.
 - Use `--load_trc <path_to_file_px.trc>`: Will use pose estimation results from a file. Useful if you want to use different parameters for pixel to meter conversion or angle calculation without running detection and pose estimation all over.
+- Use `--tracking_mode sports2d`: Will use the default Sports2D tracker. Unlike DeepSort, it is faster, does not require any parametrization, and is as good in non-crowded scenes. 
 
 <br> 
 
@@ -329,7 +333,7 @@ sports2d --time_range 1.2 2.7 --ik true --person_orientation front none left
 
 ### All the parameters
 
-Have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type for a full list of the available parameters:
+For a full list of the available parameters, have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type:
 
 ``` cmd
 sports2d --help
@@ -374,7 +378,10 @@ sports2d --help
 'osim_setup_path': ["", "path to OpenSim setup. '../OpenSim_setup' if not specified"],
 'person_orientation': ["", "front, back, left, right, auto, or none. 'front none left' if not specified. If 'auto', will be either left or right depending on the direction of the motion."],
 'close_to_zero_speed_m': ["","Sum for all keypoints: about 50 px/frame or 0.2 m/frame"], 
-'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],                'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],
+'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'deepsort_params': ["", 'Deepsort tracking parameters: """{dictionary between 3 double quotes}""". \n\
+                    More information there: https://github.com/levan92/deep_sort_realtime/blob/master/deep_sort_realtime/deepsort_tracker.py#L51'],     
 'input_size': ["", "width, height. 1280, 720 if not specified. Lower resolution will be faster but less precise"],
 'keypoint_likelihood_threshold': ["", "detected keypoints are not retained if likelihood is below this threshold. 0.3 if not specified"],
 'average_likelihood_threshold': ["", "detected persons are not retained if average keypoint likelihood is below this threshold. 0.5 if not specified"],
@@ -419,7 +426,7 @@ Sports2D:
 
 2. **Sets up pose estimation with RTMLib.** It can be run in lightweight, balanced, or performance mode, and for faster inference, keypoints can be tracked instead of detected for a certain number of frames. Any RTMPose model can be used. 
 
-3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. This carefully crafted `sports2d` tracker runs at a comparable speed as the RTMlib one but is much more robust. The user can still choose the RTMLib method if they need it by specifying it in the Config.toml file.
+3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. We crafted a 'sports2D' tracker which gives good results and runs in real time, but it is also possible to use `deepsort` in particularly challenging situations. 
 
 4. **Chooses the right persons to keep.** In single-person mode, only keeps the person with the highest average scores over the sequence. In multi-person mode, only retrieves the keypoints with high enough confidence, and only keeps the persons with high enough average confidence over each frame.
 

{sports2d-0.6.1 → sports2d-0.6.2}/Sports2D/Demo/Config_demo.toml

@@ -54,7 +54,7 @@ mode = 'balanced' # 'lightweight', 'balanced', 'performance', or """{dictionary}
 
 # A dictionary (WITHIN THREE DOUBLE QUOTES) allows you to manually select the person detection (if top_down approach) and/or pose estimation models (see https://github.com/Tau-J/rtmlib). 
 # Models can be local paths or URLs.
-# Make sure the input_sizes are within triple quotes, and that they are in the opposite order from the one in the model path (for example, it would be [192,256] for rtmpose-m_simcc-body7_pt-body7-halpe26_700e-256x192-4d3e73dd_20230605.zip). 
+# Make sure the input_sizes are within square brackets, and that they are in the opposite order from the one in the model path (for example, it would be [192,256] for rtmpose-m_simcc-body7_pt-body7-halpe26_700e-256x192-4d3e73dd_20230605.zip). 
 # If your pose_model is not provided in skeletons.py, you may have to create your own one (see example at the end of the file).
 # Example, equivalent to mode='balanced':
 # mode = """{'det_class':'YOLOX',
@@ -68,17 +68,20 @@ mode = 'balanced' # 'lightweight', 'balanced', 'performance', or """{dictionary}
 #          '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]}"""
 
-det_frequency = 1       # Run person detection only every N frames, and inbetween track previously detected bounding boxes (keypoint detection is still run on all frames). 
+det_frequency = 4       # Run person detection only every N frames, and inbetween track previously detected bounding boxes (keypoint detection is still run on all frames). 
                         # Equal to or greater than 1, can be as high as you want in simple uncrowded cases. Much faster, but might be less accurate. 
 device = 'auto' # 'auto', 'CPU', 'CUDA', 'MPS', 'ROCM'
 backend = 'auto' # 'auto', 'openvino', 'onnxruntime', 'opencv'
-tracking_mode = 'sports2d' # 'rtmlib' or 'sports2d'. 'sports2d' is generally much more accurate and comparable in speed
+tracking_mode = 'sports2d' # 'sports2d' or 'deepsort'. 'deepsort' is slower but more robust in difficult configurations
+deepsort_params = """{'max_age':30, 'n_init':3, 'nms_max_overlap':0.8, 'max_cosine_distance':0.3, 'nn_budget':200, 'max_iou_distance':0.8, 'embedder_gpu': True}""" # """{dictionary between 3 double quotes}"""
+                  # More robust in crowded scenes but Can be tricky to parametrize. More information there: https://github.com/levan92/deep_sort_realtime/blob/master/deep_sort_realtime/deepsort_tracker.py#L51
+                  # Note: For even more robust tracking, use 'embedder':'torchreid', which runs osnet_ain_x1_0 by default. Install additional dependencies with: `pip install torchreid gdown tensorboard`
 
 
 # Processing parameters
 keypoint_likelihood_threshold = 0.3 # Keypoints whose likelihood is lower will not be taken into account
 average_likelihood_threshold = 0.5  # Person will be ignored if average likelihood of good keypoints is lower than this value
-keypoint_number_threshold = 0.3     # Person will be ignored if the number of good keypoints is less than this fraction
+keypoint_number_threshold = 0.3     # Person will be ignored if the number of good keypoints (above keypoint_likelihood_threshold) is less than this fraction
 
 
 [px_to_meters_conversion]

{sports2d-0.6.1 → sports2d-0.6.2}/Sports2D/Sports2D.py

@@ -146,6 +146,7 @@ DEFAULT_CONFIG =   {'project': {'video_input': ['demo.mp4'],
                                 'device': 'auto',
                                 'backend': 'auto',
                                 'tracking_mode': 'sports2d',
+                                'deepsort_params': """{'max_age':30, 'n_init':3, 'nms_max_overlap':0.8, 'max_cosine_distance':0.3, 'nn_budget':200, 'max_iou_distance':0.8, 'embedder_gpu': True}""",
                                 'keypoint_likelihood_threshold': 0.3,
                                 'average_likelihood_threshold': 0.5,
                                 'keypoint_number_threshold': 0.3
@@ -248,7 +249,10 @@ CONFIG_HELP =   {'config': ["C", "path to a toml configuration file"],
                 'osim_setup_path': ["", "path to OpenSim setup. '../OpenSim_setup' if not specified"],
                 'person_orientation': ["", "front, back, left, right, auto, or none. 'front none left' if not specified. If 'auto', will be either left or right depending on the direction of the motion."],
                 'close_to_zero_speed_m': ["","Sum for all keypoints: about 50 px/frame or 0.2 m/frame"], 
-                'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],                'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+                'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],
+                'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+                'deepsort_params': ["", 'Deepsort tracking parameters: """{dictionary between 3 double quotes}""". \n\
+                                    More information there: https://github.com/levan92/deep_sort_realtime/blob/master/deep_sort_realtime/deepsort_tracker.py#L51'],                 # 
                 'input_size': ["", "width, height. 1280, 720 if not specified. Lower resolution will be faster but less precise"],
                 'keypoint_likelihood_threshold': ["", "detected keypoints are not retained if likelihood is below this threshold. 0.3 if not specified"],
                 'average_likelihood_threshold': ["", "detected persons are not retained if average keypoint likelihood is below this threshold. 0.5 if not specified"],

{sports2d-0.6.1 → sports2d-0.6.2}/Sports2D/Utilities/common.py

@@ -20,7 +20,9 @@ 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
@@ -28,6 +30,7 @@ from scipy import interpolate
 import imageio_ffmpeg as ffmpeg
 import cv2
 
+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
@@ -466,7 +469,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 +477,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 +485,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 +552,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 +693,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-0.6.1 → sports2d-0.6.2}/Sports2D/process.py

@@ -60,7 +60,7 @@ from functools import partial
 from datetime import datetime
 import itertools as it
 from tqdm import tqdm
-from anytree import RenderTree, PreOrderIter
+from anytree import RenderTree
 
 import numpy as np
 import pandas as pd
@@ -68,6 +68,7 @@ import cv2
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from rtmlib import PoseTracker, BodyWithFeet, Wholebody, Body, Custom
+from deep_sort_realtime.deepsort_tracker import DeepSort
 
 from Sports2D.Utilities import filter
 from Sports2D.Utilities.common import *
@@ -337,161 +338,6 @@ def compute_angle(ang_name, person_X_flipped, person_Y, angle_dict, keypoints_id
     return ang
 
 
-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 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 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 draw_dotted_line(img, start, direction, length, color=(0, 255, 0), gap=7, dot_length=3, thickness=thickness):
     '''
     Draw a dotted line with on a cv2 image
@@ -516,109 +362,6 @@ def draw_dotted_line(img, start, direction, length, color=(0, 255, 0), gap=7, do
         cv2.line(img, tuple(line_start.astype(int)), tuple(line_end.astype(int)), color, thickness)
 
 
-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+1, color, 2, cv2.LINE_AA) 
-    
-    return img
-
-
-def draw_skel(img, X, Y, model, colors=[(255, 0, 0), (0, 255, 0), (0, 0, 255)]):
-    '''
-    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
-    node_pairs = []
-    for data_i in PreOrderIter(model.root, filter_=lambda node: node.is_leaf):
-        node_branches = [node_i.id for node_i in data_i.path]
-        node_pairs += [[node_branches[i],node_branches[i+1]] for i in range(len(node_branches)-1)]
-    node_pairs = [list(x) for x in set(tuple(x) for x in node_pairs)]
-    
-    # Draw lines
-    color_cycle = it.cycle(colors)
-    for (x,y) in zip(X,Y):
-        c = next(color_cycle)
-        if not np.isnan(x).all():
-            [cv2.line(img,
-                (int(x[n[0]]), int(y[n[0]])), (int(x[n[1]]), int(y[n[1]])), c, thickness)
-                for n in node_pairs
-                if not None in n and not (np.isnan(x[n[0]]) or np.isnan(y[n[0]]) or np.isnan(x[n[1]]) or np.isnan(y[n[1]]))] # IF NOT NONE
-
-    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
-
-
 def draw_angles(img, valid_X, valid_Y, valid_angles, valid_X_flipped, keypoints_ids, keypoints_names, angle_names, display_angle_values_on= ['body', 'list'], colors=[(255, 0, 0), (0, 255, 0), (0, 0, 255)], fontSize=0.3, thickness=1):
     '''
     Draw angles on the image.
@@ -1184,6 +927,16 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
     mode = config_dict.get('pose').get('mode')
     det_frequency = config_dict.get('pose').get('det_frequency')
     tracking_mode = config_dict.get('pose').get('tracking_mode')
+    if tracking_mode == 'deepsort':
+        deepsort_params = config_dict.get('pose').get('deepsort_params')
+        try:
+            deepsort_params = ast.literal_eval(deepsort_params)
+        except: # if within single quotes instead of double quotes when run with sports2d --mode """{dictionary}"""
+            deepsort_params = deepsort_params.strip("'").replace('\n', '').replace(" ", "").replace(",", '", "').replace(":", '":"').replace("{", '{"').replace("}", '"}').replace('":"/',':/').replace('":"\\',':\\')
+            deepsort_params = re.sub(r'"\[([^"]+)",\s?"([^"]+)\]"', r'[\1,\2]', deepsort_params) # changes "[640", "640]" to [640,640]
+            deepsort_params = json.loads(deepsort_params)
+        deepsort_tracker = DeepSort(**deepsort_params)
+        deepsort_tracker.tracker.tracks.clear()
     backend = config_dict.get('pose').get('backend')
     device = config_dict.get('pose').get('device')
     
@@ -1321,8 +1074,8 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
             logging.warning("\nInvalid mode. Must be 'lightweight', 'balanced', 'performance', or '''{dictionary}''' of parameters within triple quotes. Make sure input_sizes are within square brackets.")
             logging.warning('Using the default "balanced" mode.')
             mode = 'balanced'
-
     
+
     # Skip pose estimation or set it up:
     if load_trc:
         if not '_px' in str(load_trc): 
@@ -1341,12 +1094,21 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
         keypoints_ids = [node.id for _, _, node in RenderTree(pose_model) if node.id!=None]
         keypoints_names = [node.name for _, _, node in RenderTree(pose_model) if node.id!=None]
 
-        tracking_rtmlib = True if (tracking_mode == 'rtmlib' and multiperson) else False
-        pose_tracker = setup_pose_tracker(ModelClass, det_frequency, mode, tracking_rtmlib, backend, device)
+        # Set up pose tracker
+        try:
+            pose_tracker = setup_pose_tracker(ModelClass, det_frequency, mode, False, backend, device)
+        except:
+            logging.error('Error: Pose estimation failed. Check in Config.toml that pose_model and mode are valid.')
+            raise ValueError('Error: Pose estimation failed. Check in Config.toml that pose_model and mode are valid.')
+        
+        if tracking_mode not in ['deepsort', 'sports2d']:
+            logging.warning(f"Tracking mode {tracking_mode} not recognized. Using sports2d method.")
+            tracking_mode = 'sports2d'
         logging.info(f'\nPose tracking set up for "{pose_model_name}" model.')
         logging.info(f'Mode: {mode}.\n')
-        logging.info(f'Persons are detected every {det_frequency} frames and tracked inbetween. Multi-person is {"" if multiperson else "not "}selected.')
-        logging.info(f"Parameters: {keypoint_likelihood_threshold=}, {average_likelihood_threshold=}, {keypoint_number_threshold=}")
+        logging.info(f'Persons are detected every {det_frequency} frames and tracked inbetween. Multi-person is {"" if multiperson else "not "}selected. Tracking is done with {tracking_mode}.')
+        if tracking_mode == 'deepsort': logging.info(f'Deepsort parameters: {deepsort_params}.')
+        logging.info(f"{keypoint_likelihood_threshold=}, {average_likelihood_threshold=}, {keypoint_number_threshold=}")
 
     if flip_left_right:
         try:
@@ -1383,22 +1145,22 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
         for frame_nb in frame_iterator:
             start_time = datetime.now()
             success, frame = cap.read()
+            frame_count += 1
 
             # If frame not grabbed
             if not success:
-                logging.warning(f"Failed to grab frame {frame_count}.")
+                logging.warning(f"Failed to grab frame {frame_count-1}.")
                 if save_pose:
                     all_frames_X.append([])
                     all_frames_Y.append([])
                     all_frames_scores.append([])
                 if save_angles:
                     all_frames_angles.append([])
-                frame_count += 1
                 continue
             else:
                 cv2.putText(frame, f"Press 'q' to quit", (cam_width-int(400*fontSize), cam_height-20), cv2.FONT_HERSHEY_SIMPLEX, fontSize+0.2, (255,255,255), thickness+1, cv2.LINE_AA)
                 cv2.putText(frame, f"Press 'q' to quit", (cam_width-int(400*fontSize), cam_height-20), cv2.FONT_HERSHEY_SIMPLEX, fontSize+0.2, (0,0,255), thickness, cv2.LINE_AA)
-                frame_count += 1
+
 
             # Retrieve pose or Estimate pose and track people
             if load_trc: 
@@ -1409,13 +1171,14 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
             else: 
                 # Detect poses
                 keypoints, scores = pose_tracker(frame)
-                # Track persons
-                if tracking_rtmlib:
-                    keypoints, scores = sort_people_rtmlib(pose_tracker, keypoints, scores)
-                else:
+
+                # Track poses across frames
+                if tracking_mode == 'deepsort':
+                    keypoints, scores = sort_people_deepsort(keypoints, scores, deepsort_tracker, frame, frame_count)
+                if tracking_mode == 'sports2d': 
                     if 'prev_keypoints' not in locals(): prev_keypoints = keypoints
                     prev_keypoints, keypoints, scores = sort_people_sports2d(prev_keypoints, keypoints, scores=scores)
-                
+                        
             
             # Process coordinates and compute angles
             valid_X, valid_Y, valid_scores = [], [], []
@@ -1478,7 +1241,7 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                 img = frame.copy()
                 img = draw_bounding_box(img, valid_X, valid_Y, colors=colors, fontSize=fontSize, thickness=thickness)
                 img = draw_keypts(img, valid_X, valid_Y, valid_scores, cmap_str='RdYlGn')
-                img = draw_skel(img, valid_X, valid_Y, pose_model, colors=colors)
+                img = draw_skel(img, valid_X, valid_Y, pose_model)
                 if calculate_angles:
                     img = draw_angles(img, valid_X, valid_Y, valid_angles, valid_X_flipped, new_keypoints_ids, new_keypoints_names, angle_names, display_angle_values_on=display_angle_values_on, colors=colors, fontSize=fontSize, thickness=thickness)
 

{sports2d-0.6.1 → sports2d-0.6.2}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = sports2d
-version = 0.6.1
+version = 0.6.2
 author = David Pagnon
 author_email = contact@david-pagnon.com
 description = Detect pose and compute 2D joint angles from a video.
@@ -41,10 +41,11 @@ install_requires =
 	matplotlib
 	PyQt5
 	statsmodels
-	rtmlib_pose2sim
+	rtmlib
 	openvino
 	tqdm
 	imageio_ffmpeg
+	deep-sort-realtime
 packages = find:
 
 [options.entry_points]

{sports2d-0.6.1 → sports2d-0.6.2}/sports2d.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sports2d
-Version: 0.6.1
+Version: 0.6.2
 Summary: Detect pose and compute 2D joint angles from a video.
 Home-page: https://github.com/davidpagnon/Sports2D
 Author: David Pagnon
@@ -33,10 +33,11 @@ Requires-Dist: opencv-python
 Requires-Dist: matplotlib
 Requires-Dist: PyQt5
 Requires-Dist: statsmodels
-Requires-Dist: rtmlib_pose2sim
+Requires-Dist: rtmlib
 Requires-Dist: openvino
 Requires-Dist: tqdm
 Requires-Dist: imageio_ffmpeg
+Requires-Dist: deep-sort-realtime
 
 
 [![Continuous integration](https://github.com/davidpagnon/sports2d/actions/workflows/continuous-integration.yml/badge.svg?branch=main)](https://github.com/davidpagnon/sports2d/actions/workflows/continuous-integration.yml)
@@ -212,6 +213,9 @@ Note that it does not take distortions into account, and that it will be less ac
   ``` cmd
   sports2d --multiperson false --pose_model Body --mode lightweight --det_frequency 50 
   ```
+  ``` cmd
+  sports2d --tracking_mode deepsort --deepsort_params """{'max_age':30, 'n_init':3, 'nms_max_overlap':0.8, 'max_cosine_distance':0.3, 'nn_budget':200, 'max_iou_distance':0.8, 'embedder_gpu': True}"""
+  ```
 <br>
 
 #### Run with a toml configuration file: 
@@ -249,6 +253,7 @@ Note that any detection and pose models can be used (first [deploy them with MMP
 - Use `--det_frequency 50`: Will detect poses only every 50 frames, and track keypoints in between, which is faster.
 - Use `--multiperson false`: Can be used if one single person is present in the video. Otherwise, persons' IDs may be mixed up.
 - Use `--load_trc <path_to_file_px.trc>`: Will use pose estimation results from a file. Useful if you want to use different parameters for pixel to meter conversion or angle calculation without running detection and pose estimation all over.
+- Use `--tracking_mode sports2d`: Will use the default Sports2D tracker. Unlike DeepSort, it is faster, does not require any parametrization, and is as good in non-crowded scenes. 
 
 <br> 
 
@@ -369,7 +374,7 @@ sports2d --time_range 1.2 2.7 --ik true --person_orientation front none left
 
 ### All the parameters
 
-Have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type for a full list of the available parameters:
+For a full list of the available parameters, have a look at the [Config_Demo.toml](https://github.com/davidpagnon/Sports2D/blob/main/Sports2D/Demo/Config_demo.toml) file or type:
 
 ``` cmd
 sports2d --help
@@ -414,7 +419,10 @@ sports2d --help
 'osim_setup_path': ["", "path to OpenSim setup. '../OpenSim_setup' if not specified"],
 'person_orientation': ["", "front, back, left, right, auto, or none. 'front none left' if not specified. If 'auto', will be either left or right depending on the direction of the motion."],
 'close_to_zero_speed_m': ["","Sum for all keypoints: about 50 px/frame or 0.2 m/frame"], 
-'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],                'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'multiperson': ["", "multiperson involves tracking: will be faster if set to false. true if not specified"],
+'tracking_mode': ["", "sports2d or rtmlib. sports2d is generally much more accurate and comparable in speed. sports2d if not specified"],
+'deepsort_params': ["", 'Deepsort tracking parameters: """{dictionary between 3 double quotes}""". \n\
+                    More information there: https://github.com/levan92/deep_sort_realtime/blob/master/deep_sort_realtime/deepsort_tracker.py#L51'],     
 'input_size': ["", "width, height. 1280, 720 if not specified. Lower resolution will be faster but less precise"],
 'keypoint_likelihood_threshold': ["", "detected keypoints are not retained if likelihood is below this threshold. 0.3 if not specified"],
 'average_likelihood_threshold': ["", "detected persons are not retained if average keypoint likelihood is below this threshold. 0.5 if not specified"],
@@ -459,7 +467,7 @@ Sports2D:
 
 2. **Sets up pose estimation with RTMLib.** It can be run in lightweight, balanced, or performance mode, and for faster inference, keypoints can be tracked instead of detected for a certain number of frames. Any RTMPose model can be used. 
 
-3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. This carefully crafted `sports2d` tracker runs at a comparable speed as the RTMlib one but is much more robust. The user can still choose the RTMLib method if they need it by specifying it in the Config.toml file.
+3. **Tracks people** so that their IDs are consistent across frames. A person is associated to another in the next frame when they are at a small distance. IDs remain consistent even if the person disappears from a few frames. We crafted a 'sports2D' tracker which gives good results and runs in real time, but it is also possible to use `deepsort` in particularly challenging situations. 
 
 4. **Chooses the right persons to keep.** In single-person mode, only keeps the person with the highest average scores over the sequence. In multi-person mode, only retrieves the keypoints with high enough confidence, and only keeps the persons with high enough average confidence over each frame.
 

{sports2d-0.6.1 → sports2d-0.6.2}/sports2d.egg-info/requires.txt

@@ -8,7 +8,8 @@ opencv-python
 matplotlib
 PyQt5
 statsmodels
-rtmlib_pose2sim
+rtmlib
 openvino
 tqdm
 imageio_ffmpeg
+deep-sort-realtime