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

Sports2D/process.py

@@ -54,6 +54,9 @@
 from pathlib import Path
 import sys
 import logging
+import json
+import ast
+from functools import partial
 from datetime import datetime
 import itertools as it
 from tqdm import tqdm
@@ -64,57 +67,13 @@ import pandas as pd
 import cv2
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-from rtmlib import PoseTracker, BodyWithFeet
+from rtmlib import PoseTracker, BodyWithFeet, Wholebody, Body, Custom
 
 from Sports2D.Utilities import filter
 from Sports2D.Utilities.common import *
 from Sports2D.Utilities.skeletons import *
 
 
-## CONSTANTS
-angle_dict = { # lowercase!
-    # joint angles
-    'right ankle': [['RKnee', 'RAnkle', 'RBigToe', 'RHeel'], 'dorsiflexion', 90, 1],
-    'left ankle': [['LKnee', 'LAnkle', 'LBigToe', 'LHeel'], 'dorsiflexion', 90, 1],
-    'right knee': [['RAnkle', 'RKnee', 'RHip'], 'flexion', -180, 1],
-    'left knee': [['LAnkle', 'LKnee', 'LHip'], 'flexion', -180, 1],
-    'right hip': [['RKnee', 'RHip', 'Hip', 'Neck'], 'flexion', 0, -1],
-    'left hip': [['LKnee', 'LHip', 'Hip', 'Neck'], 'flexion', 0, -1],
-    # 'lumbar': [['Neck', 'Hip', 'RHip', 'LHip'], 'flexion', -180, -1],
-    # 'neck': [['Head', 'Neck', 'RShoulder', 'LShoulder'], 'flexion', -180, -1],
-    'right shoulder': [['RElbow', 'RShoulder', 'Hip', 'Neck'], 'flexion', 0, -1],
-    'left shoulder': [['LElbow', 'LShoulder', 'Hip', 'Neck'], 'flexion', 0, -1],
-    'right elbow': [['RWrist', 'RElbow', 'RShoulder'], 'flexion', 180, -1],
-    'left elbow': [['LWrist', 'LElbow', 'LShoulder'], 'flexion', 180, -1],
-    'right wrist': [['RElbow', 'RWrist', 'RIndex'], 'flexion', -180, 1],
-    'left wrist': [['LElbow', 'LIndex', 'LWrist'], 'flexion', -180, 1],
-
-    # segment angles
-    'right foot': [['RBigToe', 'RHeel'], 'horizontal', 0, -1],
-    'left foot': [['LBigToe', 'LHeel'], 'horizontal', 0, -1],
-    'right shank': [['RAnkle', 'RKnee'], 'horizontal', 0, -1],
-    'left shank': [['LAnkle', 'LKnee'], 'horizontal', 0, -1],
-    'right thigh': [['RKnee', 'RHip'], 'horizontal', 0, -1],
-    'left thigh': [['LKnee', 'LHip'], 'horizontal', 0, -1],
-    'pelvis': [['LHip', 'RHip'], 'horizontal', 0, -1],
-    'trunk': [['Neck', 'Hip'], 'horizontal', 0, -1],
-    'shoulders': [['LShoulder', 'RShoulder'], 'horizontal', 0, -1],
-    'head': [['Head', 'Neck'], 'horizontal', 0, -1],
-    'right arm': [['RElbow', 'RShoulder'], 'horizontal', 0, -1],
-    'left arm': [['LElbow', 'LShoulder'], 'horizontal', 0, -1],
-    'right forearm': [['RWrist', 'RElbow'], 'horizontal', 0, -1],
-    'left forearm': [['LWrist', 'LElbow'], 'horizontal', 0, -1],
-    'right hand': [['RIndex', 'RWrist'], 'horizontal', 0, -1],
-    'left hand': [['LIndex', 'LWrist'], 'horizontal', 0, -1]
-    }
-
-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),
-            (255, 0, 125), (255, 125, 0), (0, 125, 255), (0, 255, 125), (125, 0, 255), (125, 255, 0), (0, 255, 0)]
-thickness = 1
-
-
 ## AUTHORSHIP INFORMATION
 __author__ = "David Pagnon, HunMin Kim"
 __copyright__ = "Copyright 2023, Sports2D"
@@ -153,7 +112,7 @@ def setup_webcam(webcam_id, save_vid, vid_output_path, input_size):
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, input_size[1])
     cam_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     cam_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
+    fps = round(cap.get(cv2.CAP_PROP_FPS))
     if fps == 0: fps = 30
     
     if cam_width != input_size[0] or cam_height != input_size[1]:
@@ -205,7 +164,7 @@ def setup_video(video_file_path, save_vid, vid_output_path):
     cam_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
     out_vid = None
-    fps = cap.get(cv2.CAP_PROP_FPS)
+    fps = round(cap.get(cv2.CAP_PROP_FPS))
     if fps == 0: fps = 30
     if save_vid:
         # try:
@@ -221,67 +180,78 @@ def setup_video(video_file_path, save_vid, vid_output_path):
     return cap, out_vid, cam_width, cam_height, fps
 
 
-def setup_backend_device():
+def setup_backend_device(backend='auto', device='auto'):
     '''
     Set up the backend and device for the pose tracker based on the availability of hardware acceleration.
     TensorRT is not supported by RTMLib yet: https://github.com/Tau-J/rtmlib/issues/12
 
-    Selects the best option in the following order of priority:
+    If device and backend are not specified, they are automatically set up in the following order of priority:
     1. GPU with CUDA and ONNXRuntime backend (if CUDAExecutionProvider is available)
     2. GPU with ROCm and ONNXRuntime backend (if ROCMExecutionProvider is available, for AMD GPUs)
     3. GPU with MPS or CoreML and ONNXRuntime backend (for macOS systems)
     4. CPU with OpenVINO backend (default fallback)
     '''
 
-    try:
-        import torch
-        import onnxruntime as ort
-        if torch.cuda.is_available() == True and 'CUDAExecutionProvider' in ort.get_available_providers():
-            device = 'cuda'
-            backend = 'onnxruntime'
-            logging.info(f"\nValid CUDA installation found: using ONNXRuntime backend with GPU.")
-        elif torch.cuda.is_available() == True and 'ROCMExecutionProvider' in ort.get_available_providers():
-            device = 'rocm'
-            backend = 'onnxruntime'
-            logging.info(f"\nValid ROCM installation found: using ONNXRuntime backend with GPU.")
-        else:
-            raise 
-    except:
+    if device!='auto' and backend!='auto':
+        device = device.lower()
+        backend = backend.lower()
+
+    if device=='auto' or backend=='auto':
+        if device=='auto' and backend!='auto' or device!='auto' and backend=='auto':
+            logging.warning(f"If you set device or backend to 'auto', you must set the other to 'auto' as well. Both device and backend will be determined automatically.")
+
         try:
+            import torch
             import onnxruntime as ort
-            if 'MPSExecutionProvider' in ort.get_available_providers() or 'CoreMLExecutionProvider' in ort.get_available_providers():
-                device = 'mps'
+            if torch.cuda.is_available() == True and 'CUDAExecutionProvider' in ort.get_available_providers():
+                device = 'cuda'
                 backend = 'onnxruntime'
-                logging.info(f"\nValid MPS installation found: using ONNXRuntime backend with GPU.")
+                logging.info(f"\nValid CUDA installation found: using ONNXRuntime backend with GPU.")
+            elif torch.cuda.is_available() == True and 'ROCMExecutionProvider' in ort.get_available_providers():
+                device = 'rocm'
+                backend = 'onnxruntime'
+                logging.info(f"\nValid ROCM installation found: using ONNXRuntime backend with GPU.")
             else:
-                raise
+                raise 
         except:
-            device = 'cpu'
-            backend = 'openvino'
-            logging.info(f"\nNo valid CUDA installation found: using OpenVINO backend with CPU.")
-    
+            try:
+                import onnxruntime as ort
+                if 'MPSExecutionProvider' in ort.get_available_providers() or 'CoreMLExecutionProvider' in ort.get_available_providers():
+                    device = 'mps'
+                    backend = 'onnxruntime'
+                    logging.info(f"\nValid MPS installation found: using ONNXRuntime backend with GPU.")
+                else:
+                    raise
+            except:
+                device = 'cpu'
+                backend = 'openvino'
+                logging.info(f"\nNo valid CUDA installation found: using OpenVINO backend with CPU.")
+        
     return backend, device
 
 
-def setup_pose_tracker(det_frequency, mode, tracking):
+def setup_pose_tracker(ModelClass, det_frequency, mode, tracking, backend, device):
     '''
     Set up the RTMLib pose tracker with the appropriate model and backend.
     If CUDA is available, use it with ONNXRuntime backend; else use CPU with openvino
 
     INPUTS:
+    - ModelClass: class. The RTMlib model class to use for pose detection (Body, BodyWithFeet, Wholebody)
     - det_frequency: int. The frequency of pose detection (every N frames)
     - mode: str. The mode of the pose tracker ('lightweight', 'balanced', 'performance')
     - tracking: bool. Whether to track persons across frames with RTMlib tracker
+    - backend: str. The backend to use for pose detection (onnxruntime, openvino, opencv)
+    - device: str. The device to use for pose detection (cpu, cuda, rocm, mps)
 
     OUTPUTS:
     - pose_tracker: PoseTracker. The initialized pose tracker object    
     '''
 
-    backend, device = setup_backend_device()
+    backend, device = setup_backend_device(backend=backend, device=device)
 
     # Initialize the pose tracker with Halpe26 model
     pose_tracker = PoseTracker(
-        BodyWithFeet,
+        ModelClass,
         det_frequency=det_frequency,
         mode=mode,
         backend=backend,
@@ -353,22 +323,17 @@ def compute_angle(ang_name, person_X_flipped, person_Y, angle_dict, keypoints_id
 
     ang_params = angle_dict.get(ang_name)
     if ang_params is not None:
-        if ang_name in ['pelvis', 'trunk', 'shoulders']:
-            angle_coords = [[np.abs(person_X_flipped[keypoints_ids[keypoints_names.index(kpt)]]), person_Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0] if kpt in keypoints_names]
-        else:
-            angle_coords = [[person_X_flipped[keypoints_ids[keypoints_names.index(kpt)]], person_Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0] if kpt in keypoints_names]
-        ang = points_to_angles(angle_coords)
-        ang += ang_params[2]
-        ang *= ang_params[3]
-        if ang_name in ['pelvis', 'shoulders']:
-            ang = ang-180 if ang>90 else ang
-            ang = ang+180 if ang<-90 else ang
-        else:
-            ang = ang-360 if ang>180 else ang
-            ang = ang+360 if ang<-180 else ang
+        try:
+            if ang_name in ['pelvis', 'trunk', 'shoulders']:
+                angle_coords = [[np.abs(person_X_flipped[keypoints_ids[keypoints_names.index(kpt)]]), person_Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0]]
+            else:
+                angle_coords = [[person_X_flipped[keypoints_ids[keypoints_names.index(kpt)]], person_Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0]]
+            ang = fixed_angles(angle_coords, ang_name)
+        except:
+            ang = np.nan    
     else:
         ang = np.nan
-
+    
     return ang
 
 
@@ -618,7 +583,7 @@ def draw_skel(img, X, Y, model, colors=[(255, 0, 0), (0, 255, 0), (0, 0, 255)]):
             [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 (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 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
 
@@ -692,31 +657,34 @@ def draw_angles(img, valid_X, valid_Y, valid_angles, valid_X_flipped, keypoints_
                     ang_name = angle_names[k]
                     ang_params = angle_dict.get(ang_name)
                     if ang_params is not None:
-                        ang_coords = np.array([[X[keypoints_ids[keypoints_names.index(kpt)]], Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0] if kpt in keypoints_names])
-                        X_flipped_coords = [X_flipped[keypoints_ids[keypoints_names.index(kpt)]] for kpt in ang_params[0] if kpt in keypoints_names]
-                        flip = -1 if any(x_flipped < 0 for x_flipped in X_flipped_coords) else 1
-                        flip = 1 if ang_name in ['pelvis', 'trunk', 'shoulders'] else flip
-                        right_angle = True if ang_params[2]==90 else False
-                        
-                        # Draw angle
-                        if len(ang_coords) == 2: # segment angle
-                            app_point, vec = draw_segment_angle(img, ang_coords, flip)
-                        else: # joint angle
-                            app_point, vec1, vec2 = draw_joint_angle(img, ang_coords, flip, right_angle)
-    
-                        # Write angle on body
-                        if 'body' in display_angle_values_on:
+                        kpts = ang_params[0]
+                        if not any(item not in keypoints_names+['Neck', 'Hip'] for item in kpts):
+                            ang_coords = np.array([[X[keypoints_ids[keypoints_names.index(kpt)]], Y[keypoints_ids[keypoints_names.index(kpt)]]] for kpt in ang_params[0] if kpt in keypoints_names])
+                            X_flipped = np.append(X_flipped, X[len(X_flipped):])
+                            X_flipped_coords = [X_flipped[keypoints_ids[keypoints_names.index(kpt)]] for kpt in ang_params[0] if kpt in keypoints_names]
+                            flip = -1 if any(x_flipped < 0 for x_flipped in X_flipped_coords) else 1
+                            flip = 1 if ang_name in ['pelvis', 'trunk', 'shoulders'] else flip
+                            right_angle = True if ang_params[2]==90 else False
+                            
+                            # Draw angle
                             if len(ang_coords) == 2: # segment angle
-                                write_angle_on_body(img, ang, app_point, vec, np.array([1,0]), dist=20, color=(255,255,255), fontSize=fontSize, thickness=thickness)
+                                app_point, vec = draw_segment_angle(img, ang_coords, flip)
                             else: # joint angle
-                                write_angle_on_body(img, ang, app_point, vec1, vec2, dist=40, color=(0,255,0), fontSize=fontSize, thickness=thickness)
-
-                        # Write angle as a list on image with progress bar
-                        if 'list' in display_angle_values_on:
-                            if len(ang_coords) == 2: # segment angle
-                                ang_label_line = write_angle_as_list(img, ang, ang_name, person_label_position, ang_label_line, color = (255,255,255), fontSize=fontSize, thickness=thickness)
-                            else:
-                                ang_label_line = write_angle_as_list(img, ang, ang_name, person_label_position, ang_label_line, color = (0,255,0), fontSize=fontSize, thickness=thickness)
+                                app_point, vec1, vec2 = draw_joint_angle(img, ang_coords, flip, right_angle)
+        
+                            # Write angle on body
+                            if 'body' in display_angle_values_on:
+                                if len(ang_coords) == 2: # segment angle
+                                    write_angle_on_body(img, ang, app_point, vec, np.array([1,0]), dist=20, color=(255,255,255), fontSize=fontSize, thickness=thickness)
+                                else: # joint angle
+                                    write_angle_on_body(img, ang, app_point, vec1, vec2, dist=40, color=(0,255,0), fontSize=fontSize, thickness=thickness)
+
+                            # Write angle as a list on image with progress bar
+                            if 'list' in display_angle_values_on:
+                                if len(ang_coords) == 2: # segment angle
+                                    ang_label_line = write_angle_as_list(img, ang, ang_name, person_label_position, ang_label_line, color = (255,255,255), fontSize=fontSize, thickness=thickness)
+                                else:
+                                    ang_label_line = write_angle_as_list(img, ang, ang_name, person_label_position, ang_label_line, color = (0,255,0), fontSize=fontSize, thickness=thickness)
 
     return img
 
@@ -869,32 +837,6 @@ def write_angle_as_list(img, ang, ang_name, person_label_position, ang_label_lin
     return ang_label_line
 
 
-def read_trc(trc_path):
-    '''
-    Read trc file
-
-    INPUTS:
-    - trc_path: path to the trc file
-
-    OUTPUTS:
-    - Q_coords: dataframe of coordinates
-    - frames_col: series of frames
-    - time_col: series of time
-    - markers: list of marker names
-    - header: list of header lines
-    '''
-    
-    with open(trc_path, 'r') as trc_file:
-        header = [next(trc_file) for line in range(5)]
-    markers = header[3].split('\t')[2::3]
-    
-    trc_df = pd.read_csv(trc_path, sep="\t", skiprows=4)
-    frames_col, time_col = pd.Series(trc_df.iloc[:,0], name='frames'), pd.Series(trc_df.iloc[:,1], name='time')
-    Q_coords = trc_df.drop(trc_df.columns[[0, 1]], axis=1)
-
-    return Q_coords, frames_col, time_col, markers, header
-
-
 def load_pose_file(Q_coords):
     '''
     Load 2D keypoints from a dataframe of XYZ coordinates
@@ -1075,6 +1017,7 @@ def angle_plots(angle_data_unfiltered, angle_data, person_id):
         ax = plt.subplot(111)
         plt.plot(angle_data_unfiltered.iloc[:,0], angle_data_unfiltered.iloc[:,id+1], label='unfiltered')
         plt.plot(angle_data.iloc[:,0], angle_data.iloc[:,id+1], label='filtered')
+
         ax.set_xlabel('Time (seconds)')
         ax.set_ylabel(angle+' (°)')
         plt.legend()
@@ -1144,138 +1087,6 @@ def compute_floor_line(trc_data, keypoint_names = ['LBigToe', 'RBigToe'], toe_sp
     return angle, xy_origin
 
 
-def mean_angles(Q_coords, markers, ang_to_consider = ['right knee', 'left knee', 'right hip', 'left hip']):
-    '''
-    Compute the mean angle time series from 3D points for a given list of angles.
-
-    INPUTS:
-    - Q_coords (DataFrame): The triangulated coordinates of the markers.
-    - markers (list): The list of marker names.
-    - ang_to_consider (list): The list of angles to consider (requires angle_dict).
-
-    OUTPUTS:
-    - ang_mean: The mean angle time series.
-    '''
-
-    ang_to_consider = ['right knee', 'left knee', 'right hip', 'left hip']
-
-    angs = []
-    for ang_name in ang_to_consider:
-        ang_params = angle_dict[ang_name]
-        ang_mk = ang_params[0]
-        
-        pts_for_angles = []
-        for pt in ang_mk:
-            pts_for_angles.append(Q_coords.iloc[:,markers.index(pt)*3:markers.index(pt)*3+3])
-        ang = points_to_angles(pts_for_angles)
-
-        ang += ang_params[2]
-        ang *= ang_params[3]
-        ang = np.abs(ang)
-
-        angs.append(ang)
-
-    ang_mean = np.mean(angs, axis=0)
-
-    return ang_mean
-
-
-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)
-    - frames when speed is close to zero (person is out of frame): 0.2 m/frame, or 50 px/frame
-    - frames when hip and knee angle below 45° (imprecise coordinates when person is crouching)
-    
-    INPUTS:
-    - 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)
-    - large_hip_knee_angles: int
-    - trimmed_extrema_percent
-
-    OUTPUT:
-    - Q_coords_low_speeds_low_angles: pd.DataFrame. The best coordinates for measurements
-    '''
-
-    # Add Hip column if not present
-    n_markers_init = len(keypoints_names)
-    if 'Hip' not in keypoints_names:
-        RHip_df = Q_coords.iloc[:,keypoints_names.index('RHip')*3:keypoints_names.index('RHip')*3+3]
-        LHip_df = Q_coords.iloc[:,keypoints_names.index('LHip')*3:keypoints_names.index('RHip')*3+3]
-        Hip_df = RHip_df.add(LHip_df, fill_value=0) /2
-        Hip_df.columns = [col+ str(int(Q_coords.columns[-1][1:])+1) for col in ['X','Y','Z']]
-        keypoints_names += ['Hip']
-        Q_coords = pd.concat([Q_coords, Hip_df], axis=1)
-    n_markers = len(keypoints_names)
-
-    # Using 80% slowest frames
-    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)
-    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)
-    ang_mean = mean_angles(Q_coords_low_speeds, keypoints_names, 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]
-
-    if n_markers_init < n_markers:
-        Q_coords_low_speeds_low_angles = Q_coords_low_speeds_low_angles.iloc[:,:-3]
-
-    return Q_coords_low_speeds_low_angles
-
-
-def compute_height(Q_coords, keypoints_names, fastest_frames_to_remove_percent=0.1, close_to_zero_speed=50, large_hip_knee_angles=45, trimmed_extrema_percent=0.5):
-    '''
-    Compute the height of the person from the trc data.
-
-    INPUTS:
-    - Q_coords: pd.DataFrame. The XYZ coordinates of each marker
-    - keypoints_names: list. The list of marker names
-    - fastest_frames_to_remove_percent: float. Frames with high speed are considered as outliers
-    - close_to_zero_speed: float. Sum for all keypoints: about 50 px/frame or 0.2 m/frame
-    - large_hip_knee_angles5: float. Hip and knee angles below this value are considered as imprecise
-    - trimmed_extrema_percent: float. Proportion of the most extreme segment values to remove before calculating their mean)
-    
-    OUTPUT:
-    - height: float. The estimated height of the person
-    '''
-    
-    # Retrieve most reliable coordinates
-    Q_coords_low_speeds_low_angles = best_coords_for_measurements(Q_coords, keypoints_names, 
-                                                                  fastest_frames_to_remove_percent=fastest_frames_to_remove_percent, close_to_zero_speed=close_to_zero_speed, large_hip_knee_angles=large_hip_knee_angles)
-    Q_coords_low_speeds_low_angles.columns = np.array([[m]*3 for m in keypoints_names]).flatten()
-
-    # Add MidShoulder column
-    df_MidShoulder = pd.DataFrame((Q_coords_low_speeds_low_angles['RShoulder'].values + Q_coords_low_speeds_low_angles['LShoulder'].values) /2)
-    df_MidShoulder.columns = ['MidShoulder']*3
-    Q_coords_low_speeds_low_angles = pd.concat((Q_coords_low_speeds_low_angles.reset_index(drop=True), df_MidShoulder), axis=1)
-
-    # Automatically compute the height of the person
-    pairs_up_to_shoulders = [['RHeel', 'RAnkle'], ['RAnkle', 'RKnee'], ['RKnee', 'RHip'], ['RHip', 'RShoulder'],
-                            ['LHeel', 'LAnkle'], ['LAnkle', 'LKnee'], ['LKnee', 'LHip'], ['LHip', 'LShoulder']]
-    try:
-        rfoot, rshank, rfemur, rback, lfoot, lshank, lfemur, lback = [euclidean_distance(Q_coords_low_speeds_low_angles[pair[0]],Q_coords_low_speeds_low_angles[pair[1]]) for pair in pairs_up_to_shoulders]
-    except:
-        raise ValueError('At least one of the following markers is missing for computing the height of the person:\
-                         RHeel, RAnkle, RKnee, RHip, RShoulder, LHeel, LAnkle, LKnee, LHip, LShoulder.\
-                         Make sure that the person is entirely visible, or use a calibration file instead, or set "to_meters=false".')
-    if 'Head' in keypoints_names:
-        head = euclidean_distance(Q_coords_low_speeds_low_angles['MidShoulder'], Q_coords_low_speeds_low_angles['Head'])
-    else:
-        head = euclidean_distance(Q_coords_low_speeds_low_angles['MidShoulder'], Q_coords_low_speeds_low_angles['Nose'])*1.33
-    heights = (rfoot + lfoot)/2 + (rshank + lshank)/2 + (rfemur + lfemur)/2 + (rback + lback)/2 + head
-    
-    # Remove the 20% most extreme values
-    height = trimmed_mean(heights, trimmed_extrema_percent=trimmed_extrema_percent)
-
-    return height
-
-
 def convert_px_to_meters(Q_coords_kpt, person_height_m, height_px, cx, cy, floor_angle):
     '''
     Convert pixel coordinates to meters.
@@ -1373,6 +1184,8 @@ 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')
+    backend = config_dict.get('pose').get('backend')
+    device = config_dict.get('pose').get('device')
     
     # Pixel to meters conversion
     to_meters = config_dict.get('px_to_meters_conversion').get('to_meters')
@@ -1406,6 +1219,7 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
     fontSize = config_dict.get('angles').get('fontSize')
     thickness = 1 if fontSize < 0.8 else 2
     flip_left_right = config_dict.get('angles').get('flip_left_right')
+    correct_segment_angles_with_floor_angle = config_dict.get('angles').get('correct_segment_angles_with_floor_angle')
 
     # Post-processing settings
     interpolate = config_dict.get('post-processing').get('interpolate')    
@@ -1464,9 +1278,52 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
         cv2.namedWindow(f'{video_file} Sports2D', cv2.WINDOW_NORMAL + cv2.WINDOW_KEEPRATIO)
         cv2.setWindowProperty(f'{video_file} Sports2D', cv2.WND_PROP_ASPECT_RATIO, cv2.WINDOW_FULLSCREEN)
 
+    # Select the appropriate model based on the model_type
+    if pose_model.upper() in ('HALPE_26', 'BODY_WITH_FEET'):
+        model_name = 'HALPE_26'
+        ModelClass = BodyWithFeet # 26 keypoints(halpe26)
+        logging.info(f"Using HALPE_26 model (body and feet) for pose estimation.")
+    elif pose_model.upper() in ('COCO_133', 'WHOLE_BODY', 'WHOLE_BODY_WRIST'):
+        model_name = 'COCO_133'
+        ModelClass = Wholebody
+        logging.info(f"Using COCO_133 model (body, feet, hands, and face) for pose estimation.")
+    elif pose_model.upper() in ('COCO_17', 'BODY'):
+        model_name = 'COCO_17'
+        ModelClass = Body
+        logging.info(f"Using COCO_17 model (body) for pose estimation.")
+    else:
+        raise ValueError(f"Invalid model_type: {model_name}. Must be 'HALPE_26', 'COCO_133', or 'COCO_17'. Use another network (MMPose, DeepLabCut, OpenPose, AlphaPose, BlazePose...) and convert the output files if you need another model. See documentation.")
+    pose_model_name = pose_model
+    pose_model = eval(model_name)
+
+    # Manually select the models if mode is a dictionary rather than 'lightweight', 'balanced', or 'performance'
+    if not mode in ['lightweight', 'balanced', 'performance']:
+        try:
+            try:
+                mode = ast.literal_eval(mode)
+            except: # if within single quotes instead of double quotes when run with sports2d --mode """{dictionary}"""
+                mode = mode.strip("'").replace('\n', '').replace(" ", "").replace(",", '", "').replace(":", '":"').replace("{", '{"').replace("}", '"}').replace('":"/',':/').replace('":"\\',':\\')
+                mode = re.sub(r'"\[([^"]+)",\s?"([^"]+)\]"', r'[\1,\2]', mode) # changes "[640", "640]" to [640,640]
+                mode = json.loads(mode)
+            det_class = mode.get('det_class')
+            det = mode.get('det_model')
+            det_input_size = mode.get('det_input_size')
+            pose_class = mode.get('pose_class')
+            pose = mode.get('pose_model')
+            pose_input_size = mode.get('pose_input_size')
+
+            ModelClass = partial(Custom,
+                        det_class=det_class, det=det, det_input_size=det_input_size,
+                        pose_class=pose_class, pose=pose, pose_input_size=pose_input_size,
+                        backend=backend, device=device)
+            
+        except (json.JSONDecodeError, TypeError):
+            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:
-    model = eval(pose_model)
     if load_trc:
         if not '_px' in str(load_trc): 
             logging.error(f'\n{load_trc} file needs to be in px, not in meters.')
@@ -1475,26 +1332,39 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
         Q_coords, _, _, keypoints_names, _ = read_trc(load_trc)
         keypoints_ids = [i for i in range(len(keypoints_names))]
         keypoints_all, scores_all = load_pose_file(Q_coords)
-        for pre, _, node in RenderTree(model):
+        for pre, _, node in RenderTree(model_name):
             if node.name in keypoints_names:
                 node.id = keypoints_names.index(node.name)
     
     else:
         # Retrieve keypoint names from model
-        keypoints_ids = [node.id for _, _, node in RenderTree(model) if node.id!=None]
-        keypoints_names = [node.name for _, _, node in RenderTree(model) if node.id!=None]
+        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(det_frequency, mode, tracking_rtmlib)
-        logging.info(f'\nPose tracking set up for BodyWithFeet model in {mode} mode.')
+        pose_tracker = setup_pose_tracker(ModelClass, det_frequency, mode, tracking_rtmlib, backend, device)
+        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=}")
 
-    Ltoe_idx = keypoints_ids[keypoints_names.index('LBigToe')]
-    LHeel_idx = keypoints_ids[keypoints_names.index('LHeel')]
-    Rtoe_idx = keypoints_ids[keypoints_names.index('RBigToe')]
-    RHeel_idx = keypoints_ids[keypoints_names.index('RHeel')]
-    L_R_direction_idx = [Ltoe_idx, LHeel_idx, Rtoe_idx, RHeel_idx]
+    if flip_left_right:
+        try:
+            Ltoe_idx = keypoints_ids[keypoints_names.index('LBigToe')]
+            LHeel_idx = keypoints_ids[keypoints_names.index('LHeel')]
+            Rtoe_idx = keypoints_ids[keypoints_names.index('RBigToe')]
+            RHeel_idx = keypoints_ids[keypoints_names.index('RHeel')]
+            L_R_direction_idx = [Ltoe_idx, LHeel_idx, Rtoe_idx, RHeel_idx]
+        except ValueError:
+            logging.warning(f"Missing 'LBigToe', 'LHeel', 'RBigToe', 'RHeel' keypoints. flip_left_right will be set to False")
+            flip_left_right = False
+
+    if calculate_angles:
+        for ang_name in angle_names:
+            ang_params = angle_dict.get(ang_name)
+            kpts = ang_params[0]
+            if any(item not in keypoints_names+['Neck', 'Hip'] for item in kpts):
+                logging.warning(f"Skipping {ang_name} angle computation because at least one of the following keypoints is not provided by the model: {ang_params[0]}.")
 
 
     # Process video or webcam feed
@@ -1568,9 +1438,6 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                         person_X = np.full_like(person_X, np.nan)
                         person_Y = np.full_like(person_Y, np.nan)
                         person_scores = np.full_like(person_scores, np.nan)
-                valid_X.append(person_X)
-                valid_Y.append(person_Y)
-                valid_scores.append(person_scores)
 
 
                 # Compute angles
@@ -1580,24 +1447,40 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                         person_X_flipped = flip_left_right_direction(person_X, L_R_direction_idx, keypoints_names, keypoints_ids)
                     else:
                         person_X_flipped = person_X.copy()
-                    valid_X_flipped.append(person_X_flipped)
                         
                     # Compute angles
                     person_angles = []
+                    # Add Neck and Hip if not provided
+                    new_keypoints_names, new_keypoints_ids = keypoints_names.copy(), keypoints_ids.copy()
+                    for kpt in ['Neck', 'Hip']:
+                        if kpt not in new_keypoints_names:
+                            person_X_flipped, person_Y, person_scores = add_neck_hip_coords(kpt, person_X_flipped, person_Y, person_scores, new_keypoints_ids, new_keypoints_names)
+                            person_X, _, _ = add_neck_hip_coords(kpt, person_X, person_Y, person_scores, new_keypoints_ids, new_keypoints_names)
+                            new_keypoints_names.append(kpt)
+                            new_keypoints_ids.append(len(person_X_flipped)-1)
+
                     for ang_name in angle_names:
-                        ang = compute_angle(ang_name, person_X_flipped, person_Y, angle_dict, keypoints_ids, keypoints_names)
+                        ang_params = angle_dict.get(ang_name)
+                        kpts = ang_params[0]
+                        if not any(item not in new_keypoints_names for item in kpts):
+                            ang = compute_angle(ang_name, person_X_flipped, person_Y, angle_dict, new_keypoints_ids, new_keypoints_names)
+                        else:
+                            ang = np.nan
                         person_angles.append(ang)
                     valid_angles.append(person_angles)
-
+                    valid_X_flipped.append(person_X_flipped)
+                valid_X.append(person_X)
+                valid_Y.append(person_Y)
+                valid_scores.append(person_scores)
 
             # Draw keypoints and skeleton
             if show_realtime_results or save_vid or save_img:
                 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, scores, cmap_str='RdYlGn')
-                img = draw_skel(img, valid_X, valid_Y, model, colors=colors)
+                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)
                 if calculate_angles:
-                    img = draw_angles(img, valid_X, valid_Y, valid_angles, valid_X_flipped, keypoints_ids, keypoints_names, angle_names, display_angle_values_on=display_angle_values_on, colors=colors, fontSize=fontSize, thickness=thickness)
+                    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)
 
                 if show_realtime_results:
                     cv2.imshow(f'{video_file} Sports2D', img)
@@ -1618,6 +1501,8 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                 elapsed_time = (datetime.now() - start_time).total_seconds()
                 frame_processing_times.append(elapsed_time)
 
+
+        # End of the video is reached
         cap.release()
         logging.info(f"Video processing completed.")
         if save_vid:
@@ -1730,9 +1615,13 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                 if show_plots and not to_meters:
                     pose_plots(trc_data_unfiltered_i, trc_data_i, i)
             
+
         # Convert px to meters
         if to_meters:
             logging.info('\nConverting pose to meters:')
+            if calib_on_person_id>=len(trc_data):
+                logging.warning(f'Person #{calib_on_person_id} not detected in the video. Calibrating on person #0 instead.')
+                calib_on_person_id = 0
             if calib_file:
                 logging.info(f'Using calibration file to convert coordinates in meters: {calib_file}.')
                 calib_params_dict = retrieve_calib_params(calib_file)
@@ -1748,10 +1637,20 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
 
                 if floor_angle == 'auto' or xy_origin == 'auto':
                     # estimated from the line formed by the toes when they are on the ground (where speed = 0)
-                    toe_speed_below = 1 # m/s (below which the foot is considered to be stationary)
-                    px_per_m = height_px/person_height_m
-                    toe_speed_below_px_frame = toe_speed_below * px_per_m / fps
-                    floor_angle_estim, xy_origin_estim = compute_floor_line(trc_data[calib_on_person_id], keypoint_names=['LBigToe', 'RBigToe'], toe_speed_below=toe_speed_below_px_frame)
+                    try:
+                        toe_speed_below = 1 # m/s (below which the foot is considered to be stationary)
+                        px_per_m = height_px/person_height_m
+                        toe_speed_below_px_frame = toe_speed_below * px_per_m / fps
+                        try:
+                            floor_angle_estim, xy_origin_estim = compute_floor_line(trc_data[calib_on_person_id], keypoint_names=['LBigToe', 'RBigToe'], toe_speed_below=toe_speed_below_px_frame)
+                        except: # no feet points
+                            floor_angle_estim, xy_origin_estim = compute_floor_line(trc_data[calib_on_person_id], keypoint_names=['LAnkle', 'RAnkle'], toe_speed_below=toe_speed_below_px_frame)
+                            xy_origin_estim[0] = xy_origin_estim[0]-0.13
+                            logging.warning(f'The RBigToe and LBigToe are missing from your model. Using ankles - 13 cm to compute the floor line.')
+                    except:
+                        floor_angle_estim = 0
+                        xy_origin_estim = cam_width/2, cam_height/2
+                        logging.warning(f'Could not estimate the floor angle and xy_origin. Make sure that the full body is visible. Using floor angle = 0° and xy_origin = [{cam_width/2}, {cam_height/2}].')
                 if not floor_angle == 'auto':
                     floor_angle_estim = floor_angle
                 if xy_origin == 'auto':
@@ -1776,7 +1675,7 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                     # Write to trc file
                     idx_path = selected_person_id if not multiperson and not calib_file else i
                     pose_path_person_m_i = (pose_output_path.parent / (pose_output_path_m.stem + f'_person{idx_path:02d}.trc'))
-                    make_trc_with_trc_data(trc_data_m_i, pose_path_person_m_i)
+                    make_trc_with_trc_data(trc_data_m_i, pose_path_person_m_i, fps=fps)
                     logging.info(f'Person {idx_path}: Pose in meters saved to {pose_path_person_m_i.resolve()}.')
                     
                 
@@ -1887,6 +1786,17 @@ def process_fun(config_dict, video_file, time_range, frame_rate, result_dir):
                     logging.info(f'Filtering with {args}')
                     all_frames_angles_person_filt = all_frames_angles_person_interp.apply(filter.filter1d, axis=0, args=filter_options)
 
+                # Remove columns with all nan values
+                all_frames_angles_person_filt.dropna(axis=1, how='all', inplace=True)
+                all_frames_angles_person = all_frames_angles_person[all_frames_angles_person_filt.columns]
+
+                # Add floor_angle_estim to segment angles
+                if correct_segment_angles_with_floor_angle and to_meters: 
+                    logging.info(f'Correcting segment angles by removing the {round(np.degrees(floor_angle_estim),2)}° floor angle.')
+                    for ang_name in all_frames_angles_person_filt.columns:
+                        if 'horizontal' in angle_dict[ang_name][1]:
+                            all_frames_angles_person_filt[ang_name] -= np.degrees(floor_angle_estim)
+
                 # Build mot file
                 angle_data = make_mot_with_angles(all_frames_angles_person_filt, all_frames_time, str(angles_path_person))
                 logging.info(f'Angles saved to {angles_path_person.resolve()}.')