pose2sim/Pose2Sim/kinematics.py
Laizo 3d6dbfc6d5
Progress on the integration of sports2D & Real-time (#139)
* 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
2024-11-15 17:30:43 +01:00

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")