fulyvv/pose2sim
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
3e31eaf50f
pose2sim/Pose2Sim/kinematics.py

565 lines
22 KiB
Python
Raw Blame History

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
###########################################################################
## KINEMATICS PROCESSING ##
###########################################################################
Process kinematic data using OpenSim tools.
This script performs scaling, inverse kinematics, and related processing
on 3D motion capture data (TRC files). The scaling process adjusts the
generic model to match the subject's physical dimensions, while inverse
kinematics computes the joint angles based on the motion data.
Set your parameters in Config.toml.
INPUTS:
- a directory containing TRC files
- kinematic processing parameters in Config.toml
OUTPUT:
- scaled OpenSim model files (.osim)
- joint angle data files (.mot)
'''
## 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 Pose2Sim.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, 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:
logging.error(f"Error reading TRC file at {trc_path}: {e}")
raise
def get_opensim_setup_dir():
'''
Locate the OpenSim setup directory within the Pose2Sim package.
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 == 'HALPE_26':
pose_model_file = 'Model_Pose2Sim_Halpe26.osim'
elif model_name == 'HALPE_68' or model_name == 'HALPE_136':
pose_model_file = 'Model_Pose2Sim_Halpe68_136.osim'
elif model_name == 'COCO_133':
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 == 'COCO_17':
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 OpenSim scaling setup file path.
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 == 'HALPE_26':
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Halpe26.xml'
elif model_name == 'HALPE_68' or model_name == 'HALPE_136':
scaling_setup_file = 'Scaling_Setup_Pose2Sim_Halpe68_136.xml'
elif model_name == 'COCO_133':
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 == 'COCO_17':
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 OpenSim inverse kinematics setup file path.
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 == 'HALPE_26':
ik_setup_file = 'IK_Setup_Pose2Sim_Halpe26.xml'
elif model_name == 'HALPE_68' or model_name == 'HALPE_136':
ik_setup_file = 'IK_Setup_Pose2Sim_Halpe68_136.xml'
elif model_name == 'COCO_133':
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 == 'COCO_17':
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_output_dir(config_dir, person_id):
'''
Determines the correct output directory based on the configuration and the person identifier.
INPUTS:
- config_dir (Path): The root directory where the configuration file is located.
- person_id (str): Identifier for the person (e.g., 'SinglePerson', 'P1').
OUTPUTS:
- Path: The path where the output files should be stored.
'''
output_dir = config_dir / 'kinematics' # Assuming 'opensim' as the default output subdirectory
# Append the person_id to the output directory if it's a multi-person setup
if person_id != "SinglePerson":
output_dir = output_dir / person_id
logging.info(f"Output directory determined as: {output_dir}")
# Create the directory if it does not exist
if not output_dir.exists():
output_dir.mkdir(parents=True, exist_ok=True)
return output_dir
def get_kpt_pairs_from_tree(root_node):
'''
Get name pairs for all parent-child relationships in the tree.
# Excludes the root node.
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 name pairs for all marker pairs in the scaling setup file.
'''
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):
'''
'''
measurement_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')
if right_left_symmetry:
measurement_dict[body_name] = marker_pairs
else:
if body_name.endswith('_r'):
marker_pairs_r = [pair for pair in marker_pairs if any([pair[0].startswith('R'), pair[1].startswith('R')])]
measurement_dict[body_name] = marker_pairs_r
elif body_name.endswith('_l'):
marker_pairs_l = [pair for pair in marker_pairs if any([pair[0].startswith('L'), pair[1].startswith('L')])]
measurement_dict[body_name] = marker_pairs_l
return measurement_dict
def dict_segment_ratio(scaling_root, unscaled_model, Q_coords_scaling, markers, right_left_symmetry=True):
'''
'''
# segment_pairs = get_kpt_pairs_from_tree(eval(model_name))
segment_pairs = get_kpt_pairs_from_scaling(scaling_root)
# Get model segment lengths
model_markers_locs = [unscaled_model.getMarkerSet().get(marker).getLocationInGround(unscaled_model.getWorkingState()).to_numpy() for marker in markers]
model_segment_lengths = np.array([euclidean_distance(model_markers_locs[markers.index(pt1)],
model_markers_locs[markers.index(pt2)])
for (pt1,pt2) in segment_pairs])
# 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_percent=0.5) for arr in trc_segment_lengths])
# 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 = segment_markers_dict.copy()
segment_ratio_dict.update({key: np.mean([segment_ratios[segment_pairs.index(k)]
for k in segment_markers_dict[key]])
for key in segment_markers_dict.keys()})
return segment_ratio_dict
def deactivate_measurements(scaling_root):
'''
'''
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):
'''
'''
# 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, scale in segment_ratio_dict.items():
new_scale = etree.Element('Scale')
# scales
scales_elem = etree.SubElement(new_scale, 'scales')
scales_elem.text = ' '.join([str(scale)]*3)
# 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:
- Retrieve the 80% slowest frames, excluding frames where the person is out of frame.
- From these frames, measure median segment lengths.
- Calculate ratio between model and measured segment lengths -> OpenSim manual scaling.
INPUTS:
- config_dict (dict): The configuration dictionary.
- person_id (str): The person identifier (e.g., 'SinglePerson', 'P1').
- trc_files (list): List of TRC files to be processed.
- output_dir (Path): The directory where the output files should be saved.
'''
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 / Path(scaling_path).name)
# 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) * 0.8)).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, 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 on the TRC files according to the OpenSim configuration.
INPUTS:
- config_dict (dict): The configuration dictionary.
- person_id (str): The person identifier (e.g., 'SinglePerson', 'P1').
- trc_files (list): List of TRC files to be processed.
- output_dir (Path): The directory where the output files should be saved.
'''
try:
# Retrieve data
ik_path = get_IK_Setup(model_name, osim_setup_dir)
ik_path_temp = str(kinematics_dir / Path(ik_path).name)
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 on the trc files of triangulated coordinates.
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.
'''
# 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.")
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("\nInverse 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")
Reference in New Issue View Git Blame Copy Permalink