3d6dbfc6d5
* Renamed an internal variable move save_video from pose to project * Rename of variables * Start importing from sports2d * setup_capture_directories import * rename of display_detectiob * Add deprecated warning message for display_detection * Loop most import from sports2d * Filed reorganized * Fix variable initiation * Move function to sports2d * Fixed imports TODO: fix file organisation * update for webcam usage * begin of parralelisation * Advancement on parallel process * Skeletons from sports2d * Creation of the new process * Combined display * Forgot in commit * Advancement on video connexion stabilisation * Code simplified * code simplififcation * fixed multiple issues * Progress on webcam connexion * Update for thread managment * Fix codec * Progress on webcam connection * fix display issues * Optimisation attempt * fix pose_tracker initiation * blocking process while searching for webcam * Common process * Improve code stability * try to fix video * Code simplification and working on debug * code simplifications * fix return fonction issue * Still try to fix issue of frames skipped * Progress on new process * Fix frame ixd number * frame range fix * move frame range
635 lines
27 KiB
Python
635 lines
27 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
|
|
'''
|
|
###########################################################################
|
|
## KINEMATICS PROCESSING ##
|
|
###########################################################################
|
|
|
|
Runs OpenSim scaling and inverse kinematics
|
|
|
|
Scaling:
|
|
- No need for a static trial: scaling is done on the triangulated coordinates (trc file)
|
|
- Remove 10% fastest frames (potential outliers)
|
|
- Remove frames where coordinate speed is null (person probably out of frame)
|
|
- Remove 40% most extreme calculated segment values (potential outliers)
|
|
- For each segment, scale on the mean of the remaining segment values
|
|
|
|
Inverse Kinematics:
|
|
- Run on the scaled model with the same trc file
|
|
- Model markers follow the triangulated markers while respecting the model kinematic constraints
|
|
- Joint angles are computed
|
|
|
|
INPUTS:
|
|
- config_dict (dict): Generated from a .toml calibration file
|
|
|
|
OUTPUTS:
|
|
- A scaled .osim model for each person
|
|
- Joint angle data files (.mot) for each person
|
|
- Optionally, OpenSim scaling and IK setup files saved to the kinematics directory
|
|
- Pose2Sim and OpenSim logs saved to files
|
|
'''
|
|
|
|
|
|
## INIT
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import pandas as pd
|
|
from lxml import etree
|
|
import logging
|
|
from anytree import PreOrderIter
|
|
|
|
import opensim
|
|
|
|
from Pose2Sim.common import natural_sort_key, euclidean_distance, trimmed_mean
|
|
from Sports2D.Utilities.skeletons import *
|
|
|
|
|
|
|
|
## AUTHORSHIP INFORMATION
|
|
__author__ = "Ivan Sun, David Pagnon"
|
|
__copyright__ = "Copyright 2021, Pose2Sim"
|
|
__credits__ = ["Ivan Sun", "David Pagnon"]
|
|
__license__ = "BSD 3-Clause License"
|
|
__version__ = "0.10.0"
|
|
__maintainer__ = "David Pagnon"
|
|
__email__ = "contact@david-pagnon.com"
|
|
__status__ = "Development"
|
|
|
|
|
|
## FUNCTIONS
|
|
def read_trc(trc_path):
|
|
'''
|
|
Read a TRC file and extract its contents.
|
|
|
|
INPUTS:
|
|
- trc_path (str): The path to the TRC file.
|
|
|
|
OUTPUTS:
|
|
- tuple: A tuple containing the Q coordinates, frames column, time column, marker names, and header.
|
|
'''
|
|
|
|
try:
|
|
with open(trc_path, 'r') as trc_file:
|
|
header = [next(trc_file) for _ in range(5)]
|
|
markers = header[3].split('\t')[2::3][:-1]
|
|
|
|
trc_df = pd.read_csv(trc_path, sep="\t", skiprows=4, encoding='utf-8')
|
|
frames_col, time_col = trc_df.iloc[:, 0], trc_df.iloc[:, 1]
|
|
Q_coords = trc_df.drop(trc_df.columns[[0, 1, -1]], axis=1)
|
|
|
|
return Q_coords, frames_col, time_col, markers, header
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error reading TRC file at {trc_path}: {e}")
|
|
|
|
|
|
def get_opensim_setup_dir():
|
|
'''
|
|
Locate the OpenSim setup directory within the Pose2Sim package.
|
|
|
|
INPUTS:
|
|
- None
|
|
|
|
OUTPUTS:
|
|
- Path: The path to the OpenSim setup directory.
|
|
'''
|
|
|
|
pose2sim_path = Path(sys.modules['Pose2Sim'].__file__).resolve().parent
|
|
setup_dir = pose2sim_path / 'OpenSim_Setup'
|
|
return setup_dir
|
|
|
|
|
|
def get_model_path(model_name, osim_setup_dir):
|
|
'''
|
|
Retrieve the path of the OpenSim model file.
|
|
|
|
INPUTS:
|
|
- model_name (str): Name of the model
|
|
- osim_setup_dir (Path): Path to the OpenSim setup directory.
|
|
|
|
OUTPUTS:
|
|
- pose_model_path: (Path) Path to the OpenSim model file.
|
|
'''
|
|
|
|
if model_name == 'BODY_25B':
|
|
pose_model_file = 'Model_Setup_Pose2Sim_Body25b.osim'
|
|
elif model_name == 'BODY_25':
|
|
pose_model_file = 'Model_Pose2Sim_Body25.osim'
|
|
elif model_name == 'BODY_135':
|
|
pose_model_file = 'Model_Pose2Sim_Body135.osim'
|
|
elif model_name == 'BLAZEPOSE':
|
|
pose_model_file = 'Model_Pose2Sim_Blazepose.osim'
|
|
elif model_name in ('HALPE_26', 'BODY_WITH_FEET'):
|
|
pose_model_file = 'Model_Pose2Sim_Halpe26.osim'
|
|
elif model_name in ('HALPE_68', 'HALPE_136'):
|
|
pose_model_file = 'Model_Pose2Sim_Halpe68_136.osim'
|
|
elif model_name in ('COCO_133', 'WHOLE_BODY'):
|
|
pose_model_file = 'Model_Pose2Sim_Coco133.osim'
|
|
# elif model_name == 'COCO' or model_name == 'MPII':
|
|
# pose_model_file = 'Model_Pose2Sim_Coco.osim'
|
|
elif model_name in ('COCO_17', 'BODY'):
|
|
pose_model_file = 'Model_Pose2Sim_Coco17.osim'
|
|
elif model_name == 'LSTM':
|
|
pose_model_file = 'Model_Pose2Sim_LSTM.osim'
|
|
else:
|
|
raise ValueError(f"Pose model '{model_name}' not found.")
|
|
|
|
unscaled_model_path = osim_setup_dir / pose_model_file
|
|
|
|
return unscaled_model_path
|
|
|
|
|
|
def get_scaling_setup(model_name, osim_setup_dir):
|
|
'''
|
|
Retrieve the path of the OpenSim scaling setup file.
|
|
|
|
INPUTS:
|
|
- model_name (str): Name of the model
|
|
- osim_setup_dir (Path): Path to the OpenSim setup directory.
|
|
|
|
OUTPUTS:
|
|
- scaling_setup_path: (Path) Path to the OpenSim scaling setup file.
|
|
'''
|
|
|
|
if model_name == 'BODY_25B':
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Body25b.xml'
|
|
elif model_name == 'BODY_25':
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Body25.xml'
|
|
elif model_name == 'BODY_135':
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Body135.xml'
|
|
elif model_name == 'BLAZEPOSE':
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Blazepose.xml'
|
|
elif model_name in ('HALPE_26', 'BODY_WITH_FEET'):
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Halpe26.xml'
|
|
elif model_name in ('HALPE_68', 'HALPE_136'):
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Halpe68_136.xml'
|
|
elif model_name in ('COCO_133', 'WHOLE_BODY'):
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Coco133.xml'
|
|
# elif model_name == 'COCO' or model_name == 'MPII':
|
|
# scaling_setup_file = 'Scaling_Setup_Pose2Sim_Coco.xml'
|
|
elif model_name in ('COCO_17', 'BODY'):
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Coco17.xml'
|
|
elif model_name == 'LSTM':
|
|
scaling_setup_file = 'Scaling_Setup_Pose2Sim_LSTM.xml'
|
|
else:
|
|
raise ValueError(f"Pose model '{model_name}' not found.")
|
|
|
|
scaling_setup_path = osim_setup_dir / scaling_setup_file
|
|
|
|
return scaling_setup_path
|
|
|
|
|
|
def get_IK_Setup(model_name, osim_setup_dir):
|
|
'''
|
|
Retrieve the path of the OpenSim inverse kinematics setup file.
|
|
|
|
INPUTS:
|
|
- model_name (str): Name of the model
|
|
- osim_setup_dir (Path): Path to the OpenSim setup directory.
|
|
|
|
OUTPUTS:
|
|
- ik_setup_path: (Path) Path to the OpenSim IK setup file.
|
|
'''
|
|
|
|
if model_name == 'BODY_25B':
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Body25b.xml'
|
|
elif model_name == 'BODY_25':
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Body25.xml'
|
|
elif model_name == 'BODY_135':
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Body135.xml'
|
|
elif model_name == 'BLAZEPOSE':
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Blazepose.xml'
|
|
elif model_name in ('HALPE_26', 'BODY_WITH_FEET'):
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Halpe26.xml'
|
|
elif model_name in ('HALPE_68', 'HALPE_136'):
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Halpe68_136.xml'
|
|
elif model_name in ('COCO_133', 'WHOLE_BODY'):
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Coco133.xml'
|
|
# elif model_name == 'COCO' or model_name == 'MPII':
|
|
# ik_setup_file = 'IK_Setup_Pose2Sim_Coco.xml'
|
|
elif model_name in ('COCO_17', 'BODY'):
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_Coco17.xml'
|
|
elif model_name == 'LSTM':
|
|
ik_setup_file = 'IK_Setup_Pose2Sim_LSTM.xml'
|
|
else:
|
|
raise ValueError(f"Pose model '{model_name}' not found.")
|
|
|
|
ik_setup_path = osim_setup_dir / ik_setup_file
|
|
return ik_setup_path
|
|
|
|
|
|
def get_kpt_pairs_from_tree(root_node):
|
|
'''
|
|
Get marker pairs for all parent-child relationships in the tree.
|
|
# Excludes the root node.
|
|
# Not used in the current version.
|
|
|
|
INPUTS:
|
|
- root_node (Node): The root node of the tree.
|
|
|
|
OUTPUTS:
|
|
- list: A list of name pairs for all parent-child relationships in the tree.
|
|
'''
|
|
|
|
pairs = []
|
|
for node in PreOrderIter(root_node):
|
|
# if node.is_root:
|
|
# continue
|
|
for child in node.children:
|
|
pairs.append([node.name, child.name])
|
|
|
|
return pairs
|
|
|
|
|
|
def get_kpt_pairs_from_scaling(scaling_root):
|
|
'''
|
|
Get all marker pairs from the scaling setup file.
|
|
|
|
INPUTS:
|
|
- scaling_root (Element): The root element of the scaling setup file.
|
|
|
|
OUTPUTS:
|
|
- pairs: A list of marker pairs.
|
|
'''
|
|
|
|
pairs = [pair.find('markers').text.strip().split(' ')
|
|
for pair in scaling_root[0].findall(".//MarkerPair")]
|
|
|
|
return pairs
|
|
|
|
|
|
def dict_segment_marker_pairs(scaling_root, right_left_symmetry=True):
|
|
'''
|
|
Get a dictionary of segment names and their corresponding marker pairs.
|
|
|
|
INPUTS:
|
|
- scaling_root (Element): The root element of the scaling setup file.
|
|
- right_left_symmetry (bool): Whether to consider right and left side of equal size.
|
|
|
|
OUTPUTS:
|
|
- segment_markers_dict: A dictionary of segment names and their corresponding marker pairs.
|
|
'''
|
|
|
|
segment_markers_dict = {}
|
|
for measurement in scaling_root.findall(".//Measurement"):
|
|
# Collect all marker pairs for this measurement
|
|
marker_pairs = [pair.find('markers').text.strip().split() for pair in measurement.findall(".//MarkerPair")]
|
|
|
|
# Collect all body scales for this measurement
|
|
for body_scale in measurement.findall(".//BodyScale"):
|
|
body_name = body_scale.get('name')
|
|
axes = body_scale.find('axes').text.strip().split()
|
|
for axis in axes:
|
|
body_name_axis = f"{body_name}_{axis}"
|
|
if right_left_symmetry:
|
|
segment_markers_dict.setdefault(body_name_axis, []).extend(marker_pairs)
|
|
else:
|
|
if body_name.endswith('_r'):
|
|
marker_pairs_r = [pair for pair in marker_pairs if any([pair[0].upper().startswith('R'), pair[1].upper().startswith('R')])]
|
|
segment_markers_dict.setdefault(body_name_axis, []).extend(marker_pairs_r)
|
|
elif body_name.endswith('_l'):
|
|
marker_pairs_l = [pair for pair in marker_pairs if any([pair[0].upper().startswith('L'), pair[1].upper().startswith('L')])]
|
|
segment_markers_dict.setdefault(body_name_axis, []).extend(marker_pairs_l)
|
|
else:
|
|
segment_markers_dict.setdefault(body_name_axis, []).extend(marker_pairs)
|
|
|
|
return segment_markers_dict
|
|
|
|
|
|
def dict_segment_ratio(scaling_root, unscaled_model, Q_coords_scaling, markers, trimmed_extrema_percent=0.5, right_left_symmetry=True):
|
|
'''
|
|
Calculate the ratios between the size of the actual segment and the size of the model segment.
|
|
X, Y, and Z ratios are calculated separately if the original scaling setup file asks for it.
|
|
|
|
INPUTS:
|
|
- scaling_root (Element): The root element of the scaling setup file.
|
|
- unscaled_model (Model): The original OpenSim model before scaling.
|
|
- Q_coords_scaling (DataFrame): The triangulated coordinates of the markers.
|
|
- markers (list): The list of marker names.
|
|
- trimmed_extrema_percent (float): The proportion of the most extreme segment values to remove before calculating their mean.
|
|
- right_left_symmetry (bool): Whether to consider right and left side of equal size.
|
|
|
|
OUTPUTS:
|
|
- segment_ratio_dict: A dictionary of segment names and their corresponding X, Y, and Z ratios.
|
|
'''
|
|
|
|
# segment_pairs = get_kpt_pairs_from_tree(eval(model_name))
|
|
segment_pairs = get_kpt_pairs_from_scaling(scaling_root)
|
|
|
|
# Get median segment lengths from Q_coords_scaling. Trimmed mean works better than mean or median
|
|
trc_segment_lengths = np.array([euclidean_distance(Q_coords_scaling.iloc[:,markers.index(pt1)*3:markers.index(pt1)*3+3],
|
|
Q_coords_scaling.iloc[:,markers.index(pt2)*3:markers.index(pt2)*3+3])
|
|
for (pt1,pt2) in segment_pairs])
|
|
# trc_segment_lengths = np.median(trc_segment_lengths, axis=1)
|
|
# trc_segment_lengths = np.mean(trc_segment_lengths, axis=1)
|
|
trc_segment_lengths = np.array([trimmed_mean(arr, trimmed_extrema_percent=trimmed_extrema_percent) for arr in trc_segment_lengths])
|
|
|
|
# Get model segment lengths
|
|
model_markers = [marker for marker in markers if marker in [m.getName() for m in unscaled_model.getMarkerSet()]]
|
|
model_markers_locs = [unscaled_model.getMarkerSet().get(marker).getLocationInGround(unscaled_model.getWorkingState()).to_numpy() for marker in model_markers]
|
|
model_segment_lengths = np.array([euclidean_distance(model_markers_locs[model_markers.index(pt1)],
|
|
model_markers_locs[model_markers.index(pt2)])
|
|
for (pt1,pt2) in segment_pairs])
|
|
|
|
# Calculate ratio for each segment
|
|
segment_ratios = trc_segment_lengths / model_segment_lengths
|
|
segment_markers_dict = dict_segment_marker_pairs(scaling_root, right_left_symmetry=right_left_symmetry)
|
|
segment_ratio_dict_temp = segment_markers_dict.copy()
|
|
segment_ratio_dict_temp.update({key: np.mean([segment_ratios[segment_pairs.index(k)]
|
|
for k in segment_markers_dict[key]])
|
|
for key in segment_markers_dict.keys()})
|
|
# Merge X, Y, Z ratios into single key
|
|
segment_ratio_dict={}
|
|
xyz_keys = list(set([key[:-2] for key in segment_ratio_dict_temp.keys()]))
|
|
for key in xyz_keys:
|
|
segment_ratio_dict[key] = [segment_ratio_dict_temp[key+'_X'], segment_ratio_dict_temp[key+'_Y'], segment_ratio_dict_temp[key+'_Z']]
|
|
|
|
return segment_ratio_dict
|
|
|
|
|
|
def deactivate_measurements(scaling_root):
|
|
'''
|
|
Deactivate all scalings based on marker positions (called 'measurements' in OpenSim) in the scaling setup file.
|
|
(will use scaling based on segment sizes instead (called 'manual' in OpenSim))
|
|
|
|
INPUTS:
|
|
- scaling_root (Element): The root element of the scaling setup file.
|
|
|
|
OUTPUTS:
|
|
- scaling_root with deactivated measurements.
|
|
'''
|
|
|
|
measurement_set = scaling_root.find(".//MeasurementSet/objects")
|
|
for measurement in measurement_set.findall('Measurement'):
|
|
apply_elem = measurement.find('apply')
|
|
apply_elem.text = 'false'
|
|
|
|
|
|
def update_scale_values(scaling_root, segment_ratio_dict):
|
|
'''
|
|
Remove previous scaling values ('manual') and
|
|
add new scaling values based on calculated segment ratios.
|
|
|
|
INPUTS:
|
|
- scaling_root (Element): The root element of the scaling setup file.
|
|
- segment_ratio_dict (dict): A dictionary of segment names and their corresponding X, Y, and Z ratios.
|
|
|
|
OUTPUTS:
|
|
- scaling_root with updated scaling values.
|
|
'''
|
|
|
|
# Get the ScaleSet/objects element
|
|
scale_set = scaling_root.find(".//ScaleSet/objects")
|
|
|
|
# Remove all existing Scale elements
|
|
for scale in scale_set.findall('Scale'):
|
|
scale_set.remove(scale)
|
|
|
|
# Add new Scale elements based on scale_dict
|
|
for segment, scales in segment_ratio_dict.items():
|
|
new_scale = etree.Element('Scale')
|
|
# scales
|
|
scales_elem = etree.SubElement(new_scale, 'scales')
|
|
scales_elem.text = ' '.join(map(str, scales))
|
|
# segment name
|
|
segment_elem = etree.SubElement(new_scale, 'segment')
|
|
segment_elem.text = segment
|
|
# apply True
|
|
apply_elem = etree.SubElement(new_scale, 'apply')
|
|
apply_elem.text = 'true'
|
|
|
|
scale_set.append(new_scale)
|
|
|
|
|
|
def perform_scaling(trc_file, kinematics_dir, osim_setup_dir, model_name, right_left_symmetry=True, subject_height=1.75, subject_mass=70, remove_scaling_setup=True):
|
|
'''
|
|
Perform model scaling based on the (not necessarily static) TRC file:
|
|
- Remove 10% fastest frames (potential outliers)
|
|
- Remove frames where coordinate speed is null (person probably out of frame)
|
|
- Remove 40% most extreme calculated segment values (potential outliers)
|
|
- For each segment, scale on the mean of the remaining segment values
|
|
|
|
INPUTS:
|
|
- trc_file (Path): The path to the TRC file.
|
|
- kinematics_dir (Path): The directory where the kinematics files are saved.
|
|
- osim_setup_dir (Path): The directory where the OpenSim setup and model files are stored.
|
|
- model_name (str): The name of the model.
|
|
- right_left_symmetry (bool): Whether to consider right and left side of equal size.
|
|
- subject_height (float): The height of the subject.
|
|
- subject_mass (float): The mass of the subject.
|
|
- remove_scaling_setup (bool): Whether to remove the scaling setup file after scaling.
|
|
|
|
OUTPUTS:
|
|
- A scaled OpenSim model file.
|
|
'''
|
|
|
|
fastest_frames_to_remove_percent = 0.1
|
|
trimmed_extrema_percent = 0.4 # proportion of the most extreme segment values to remove before calculating their mean
|
|
|
|
try:
|
|
# Load model
|
|
opensim.ModelVisualizer.addDirToGeometrySearchPaths(str(osim_setup_dir / 'Geometry'))
|
|
unscaled_model_path = get_model_path(model_name, osim_setup_dir)
|
|
if not unscaled_model_path:
|
|
raise ValueError(f"Unscaled OpenSim model not found at: {unscaled_model_path}")
|
|
unscaled_model = opensim.Model(str(unscaled_model_path))
|
|
unscaled_model.initSystem()
|
|
scaled_model_path = (kinematics_dir / (trc_file.stem + '.osim')).resolve()
|
|
|
|
# Load scaling setup
|
|
scaling_path = get_scaling_setup(model_name, osim_setup_dir)
|
|
scaling_tree = etree.parse(scaling_path)
|
|
scaling_root = scaling_tree.getroot()
|
|
scaling_path_temp = str(kinematics_dir / (trc_file.stem + '_scaling_setup.xml'))
|
|
|
|
# Read trc file
|
|
Q_coords, _, _, markers, _ = read_trc(trc_file)
|
|
|
|
# Using 80% slowest frames for scaling, removing frames when person is out of frame
|
|
Q_diff = Q_coords.diff(axis=0).sum(axis=1)
|
|
Q_diff = Q_diff[Q_diff != 0] # remove when speed is 0 (person out of frame)
|
|
min_speed_indices = Q_diff.abs().nsmallest(int(len(Q_diff) * (1-fastest_frames_to_remove_percent))).index
|
|
Q_coords_scaling = Q_coords.iloc[min_speed_indices].reset_index(drop=True)
|
|
|
|
# Get manual scale values (scale on trimmed mean of measured segments rather than on raw keypoints)
|
|
segment_ratio_dict = dict_segment_ratio(scaling_root, unscaled_model, Q_coords_scaling, markers,
|
|
trimmed_extrema_percent=trimmed_extrema_percent, right_left_symmetry=right_left_symmetry)
|
|
|
|
# Update scaling setup file
|
|
scaling_root[0].find('mass').text = str(subject_mass)
|
|
scaling_root[0].find('height').text = str(subject_height)
|
|
scaling_root[0].find('GenericModelMaker').find('model_file').text = str(unscaled_model_path)
|
|
scaling_root[0].find(".//scaling_order").text = ' manualScale measurements'
|
|
deactivate_measurements(scaling_root)
|
|
update_scale_values(scaling_root, segment_ratio_dict)
|
|
for mk_f in scaling_root[0].findall(".//marker_file"): mk_f.text = "Unassigned"
|
|
scaling_root[0].find('ModelScaler').find('output_model_file').text = str(scaled_model_path)
|
|
|
|
etree.indent(scaling_tree, space='\t', level=0)
|
|
scaling_tree.write(scaling_path_temp, pretty_print=True, xml_declaration=True, encoding='utf-8')
|
|
|
|
# Run scaling
|
|
opensim.ScaleTool(scaling_path_temp).run()
|
|
|
|
# Remove scaling setup
|
|
if remove_scaling_setup:
|
|
Path(scaling_path_temp).unlink()
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error during scaling for {trc_file}: {e}")
|
|
raise
|
|
|
|
|
|
def perform_IK(trc_file, kinematics_dir, osim_setup_dir, model_name, remove_IK_setup=True):
|
|
'''
|
|
Perform inverse kinematics based on a TRC file and a scaled OpenSim model:
|
|
- Model markers follow the triangulated markers while respecting the model kinematic constraints
|
|
- Joint angles are computed
|
|
|
|
INPUTS:
|
|
- trc_file (Path): The path to the TRC file.
|
|
- kinematics_dir (Path): The directory where the kinematics files are saved.
|
|
- osim_setup_dir (Path): The directory where the OpenSim setup and model files are stored.
|
|
- model_name (str): The name of the model.
|
|
- remove_IK_setup (bool): Whether to remove the IK setup file after running IK.
|
|
|
|
OUTPUTS:
|
|
- A joint angle data file (.mot).
|
|
'''
|
|
|
|
try:
|
|
# Retrieve data
|
|
ik_path = get_IK_Setup(model_name, osim_setup_dir)
|
|
ik_path_temp = str(kinematics_dir / (trc_file.stem + '_ik_setup.xml'))
|
|
scaled_model_path = (kinematics_dir / (trc_file.stem + '.osim')).resolve()
|
|
output_motion_file = Path(kinematics_dir, trc_file.stem + '.mot').resolve()
|
|
if not trc_file.exists():
|
|
raise FileNotFoundError(f"TRC file does not exist: {trc_file}")
|
|
_, _, time_col, _, _ = read_trc(trc_file)
|
|
start_time, end_time = time_col.iloc[0], time_col.iloc[-1]
|
|
|
|
# Update IK setup file
|
|
ik_tree = etree.parse(ik_path)
|
|
ik_root = ik_tree.getroot()
|
|
ik_root.find('.//model_file').text = str(scaled_model_path)
|
|
ik_root.find('.//time_range').text = f'{start_time} {end_time}'
|
|
ik_root.find('.//output_motion_file').text = str(output_motion_file)
|
|
ik_root.find('.//marker_file').text = str(trc_file.resolve())
|
|
ik_tree.write(ik_path_temp)
|
|
|
|
# Run IK
|
|
opensim.InverseKinematicsTool(str(ik_path_temp)).run()
|
|
|
|
# Remove IK setup
|
|
if remove_IK_setup:
|
|
Path(ik_path_temp).unlink()
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error during IK for {trc_file}: {e}")
|
|
raise
|
|
|
|
|
|
def kinematics(config_dict):
|
|
'''
|
|
Runs OpenSim scaling and inverse kinematics
|
|
|
|
Scaling:
|
|
- No need for a static trial: scaling is done on the triangulated coordinates (trc file)
|
|
- Remove 10% fastest frames (potential outliers)
|
|
- Remove frames where coordinate speed is null (person probably out of frame)
|
|
- Remove 40% most extreme calculated segment values (potential outliers)
|
|
- For each segment, scale on the mean of the remaining segment values
|
|
|
|
Inverse Kinematics:
|
|
- Run on the scaled model with the same trc file
|
|
- Model markers follow the triangulated markers while respecting the model kinematic constraints
|
|
- Joint angles are computed
|
|
|
|
INPUTS:
|
|
- config_dict (dict): Generated from a .toml calibration file
|
|
|
|
OUTPUTS:
|
|
- A scaled .osim model for each person
|
|
- Joint angle data files (.mot) for each person
|
|
- Optionally, OpenSim scaling and IK setup files saved to the kinematics directory
|
|
- Pose2Sim and OpenSim logs saved to files
|
|
'''
|
|
|
|
# Read config_dict
|
|
project_dir = config_dict.get('project').get('project_dir')
|
|
# if batch
|
|
session_dir = Path(project_dir) / '..'
|
|
# if single trial
|
|
session_dir = session_dir if 'Config.toml' in os.listdir(session_dir) else os.getcwd()
|
|
use_augmentation = config_dict.get('kinematics').get('use_augmentation')
|
|
if use_augmentation: model_name = 'LSTM'
|
|
else: model_name = config_dict.get('pose').get('pose_model').upper()
|
|
right_left_symmetry = config_dict.get('kinematics').get('right_left_symmetry')
|
|
remove_scaling_setup = config_dict.get('kinematics').get('remove_individual_scaling_setup')
|
|
remove_IK_setup = config_dict.get('kinematics').get('remove_individual_IK_setup')
|
|
subject_height = config_dict.get('project').get('participant_height')
|
|
subject_mass = config_dict.get('project').get('participant_mass')
|
|
|
|
pose3d_dir = Path(project_dir) / 'pose-3d'
|
|
kinematics_dir = Path(project_dir) / 'kinematics'
|
|
kinematics_dir.mkdir(parents=True, exist_ok=True)
|
|
osim_setup_dir = get_opensim_setup_dir()
|
|
|
|
# OpenSim logs saved to a different file
|
|
opensim_logs_file = kinematics_dir / 'opensim_logs.txt'
|
|
opensim.Logger.setLevelString('Info')
|
|
opensim.Logger.removeFileSink()
|
|
opensim.Logger.addFileSink(str(opensim_logs_file))
|
|
|
|
# Find all trc files
|
|
trc_files = []
|
|
if use_augmentation:
|
|
trc_files = [f for f in pose3d_dir.glob('*.trc') if '_LSTM' in f.name]
|
|
if len(trc_files) == 0:
|
|
model_name = config_dict.get('pose').get('pose_model').upper()
|
|
logging.warning("No LSTM trc files found. Using non augmented trc files instead.")
|
|
if len(trc_files) == 0: # filtered files by default
|
|
trc_files = [f for f in pose3d_dir.glob('*.trc') if '_LSTM' not in f.name and '_filt' in f.name and '_scaling' not in f.name]
|
|
if len(trc_files) == 0:
|
|
trc_files = [f for f in pose3d_dir.glob('*.trc') if '_LSTM' not in f.name and '_scaling' not in f.name]
|
|
if len(trc_files) == 0:
|
|
raise ValueError(f'No trc files found in {pose3d_dir}.')
|
|
sorted(trc_files, key=natural_sort_key)
|
|
|
|
# Get subject heights and masses
|
|
if subject_height is None or subject_height == 0:
|
|
subject_height = [1.75] * len(trc_files)
|
|
logging.warning("No subject height found in Config.toml. Using default height of 1.75m.")
|
|
elif not type(subject_height) == list: # int or float
|
|
subject_height = [subject_height]
|
|
elif len(subject_height) < len(trc_files):
|
|
logging.warning("Number of subject heights does not match number of TRC files. Missing heights are set to 1.75m.")
|
|
subject_height += [1.75] * (len(trc_files) - len(subject_height))
|
|
|
|
if subject_mass is None or subject_mass == 0:
|
|
subject_mass = [70] * len(trc_files)
|
|
logging.warning("No subject mass found in Config.toml. Using default mass of 70kg.")
|
|
elif not type(subject_mass) == list:
|
|
subject_mass = [subject_mass]
|
|
elif len(subject_mass) < len(trc_files):
|
|
logging.warning("Number of subject masses does not match number of TRC files. Missing masses are set to 70kg.\n")
|
|
subject_mass += [70] * (len(trc_files) - len(subject_mass))
|
|
|
|
# Perform scaling and IK for each trc file
|
|
for p, trc_file in enumerate(trc_files):
|
|
logging.info(f"Processing TRC file: {trc_file.resolve()}")
|
|
|
|
logging.info("Scaling...")
|
|
perform_scaling(trc_file, kinematics_dir, osim_setup_dir, model_name, right_left_symmetry=right_left_symmetry, subject_height=subject_height[p], subject_mass=subject_mass[p], remove_scaling_setup=remove_scaling_setup)
|
|
logging.info(f"\tDone. OpenSim logs saved to {opensim_logs_file.resolve()}.")
|
|
logging.info(f"\tScaled model saved to {(kinematics_dir / (trc_file.stem + '_scaled.osim')).resolve()}")
|
|
|
|
logging.info("Inverse Kinematics...")
|
|
perform_IK(trc_file, kinematics_dir, osim_setup_dir, model_name, remove_IK_setup=remove_IK_setup)
|
|
logging.info(f"\tDone. OpenSim logs saved to {opensim_logs_file.resolve()}.")
|
|
logging.info(f"\tJoint angle data saved to {(kinematics_dir / (trc_file.stem + '.mot')).resolve()}\n")
|