pose2sim/Pose2Sim/augmenter.py

import os
import numpy as np
from . import utilsDataman
import copy
import tensorflow as tf
from .utils import TRC2numpy
import json
import os
import glob


# subject_height must be in meters
def get_midhip_data(trc_file):
    try:
        # Find MidHip data
        midhip_data = trc_file.marker("CHip")
        if midhip_data is None or len(midhip_data) == 0:
            raise ValueError("MidHip data is empty")
    except (KeyError, ValueError):
        # If MidHip data is not found, calculate it from RHip and LHip
        rhip_data = trc_file.marker("RHip")
        lhip_data = trc_file.marker("LHip")
        midhip_data = (rhip_data + lhip_data) / 2

    return midhip_data


def augmentTRC(config_dict):

    # get parameters from Config.toml
    project_dir = config_dict.get('project').get('project_dir')
    session_dir = os.path.realpath(os.path.join(project_dir, '..', '..'))
    pathInputTRCFile = os.path.realpath(os.path.join(project_dir, 'pose-3d'))
    pathOutputTRCFile = os.path.realpath(os.path.join(project_dir, 'pose-3d'))
    subject_height = config_dict.get('project').get('participant_height')
    if subject_height is None or subject_height == 0:
        raise ValueError("Subject height is not set or invalid in the config file.")
    subject_mass = config_dict.get('project').get('participant_mass')
    augmenterDir = os.path.join(session_dir, '..', '..', 'MarkerAugmenter')
    augmenterModelName = 'LSTM'
    augmenter_model = 'v0.3'
    offset = True

    # Apply all trc files
    trc_files = [f for f in glob.glob(os.path.join(pathInputTRCFile, '*.trc')) if '_LSTM' not in f]
    for pathInputTRCFile in trc_files:
        pathOutputTRCFile = os.path.splitext(pathInputTRCFile)[0] + "_LSTM.trc"
    
    # This is by default - might need to be adjusted in the future.
    featureHeight = True
    featureWeight = True
    
    # Augmenter types
    if augmenter_model == 'v0.0':
        from .utils import getOpenPoseMarkers_fullBody
        feature_markers_full, response_markers_full = getOpenPoseMarkers_fullBody()         
        augmenterModelType_all = [augmenter_model]
        feature_markers_all = [feature_markers_full]
        response_markers_all = [response_markers_full]            
    elif augmenter_model == 'v0.1' or augmenter_model == 'v0.2':
        # Lower body           
        augmenterModelType_lower = '{}_lower'.format(augmenter_model)
        from .utils import getOpenPoseMarkers_lowerExtremity
        feature_markers_lower, response_markers_lower = getOpenPoseMarkers_lowerExtremity()
        # Upper body
        augmenterModelType_upper = '{}_upper'.format(augmenter_model)
        from .utils import getMarkers_upperExtremity_noPelvis
        feature_markers_upper, response_markers_upper = getMarkers_upperExtremity_noPelvis()        
        augmenterModelType_all = [augmenterModelType_lower, augmenterModelType_upper]
        feature_markers_all = [feature_markers_lower, feature_markers_upper]
        response_markers_all = [response_markers_lower, response_markers_upper]
    else:
        # Lower body           
        augmenterModelType_lower = '{}_lower'.format(augmenter_model)
        from .utils import getOpenPoseMarkers_lowerExtremity2
        feature_markers_lower, response_markers_lower = getOpenPoseMarkers_lowerExtremity2()
        # Upper body
        augmenterModelType_upper = '{}_upper'.format(augmenter_model)
        from .utils import getMarkers_upperExtremity_noPelvis2
        feature_markers_upper, response_markers_upper = getMarkers_upperExtremity_noPelvis2()        
        augmenterModelType_all = [augmenterModelType_lower, augmenterModelType_upper]
        feature_markers_all = [feature_markers_lower, feature_markers_upper]
        response_markers_all = [response_markers_lower, response_markers_upper]
    print('Using augmenter model: {}'.format(augmenter_model))
    
    # %% Process data.
    # Import TRC file
    trc_file = utilsDataman.TRCFile(pathInputTRCFile)
    
    # Loop over augmenter types to handle separate augmenters for lower and
    # upper bodies.
    outputs_all = {}
    n_response_markers_all = 0
    for idx_augm, augmenterModelType in enumerate(augmenterModelType_all):
        outputs_all[idx_augm] = {}
        feature_markers = feature_markers_all[idx_augm]
        response_markers = response_markers_all[idx_augm]
        
        augmenterModelDir = os.path.join(augmenterDir, augmenterModelName, 
                                         augmenterModelType)
        
        # %% Pre-process inputs.
        # Step 1: import .trc file with OpenPose marker trajectories.  
        trc_data = TRC2numpy(pathInputTRCFile, feature_markers)

        # # Add these lines to get RHip and LHip data and calculate midHip
        # rhip_data = trc_file.marker("RHip")  # Replace "RHip" with the actual name in your trc file
        # lhip_data = trc_file.marker("LHip")  # Replace "LHip" with the actual name in your trc file

        # Calculate the midHip marker as the average of RHip and LHip
        midhip_data = get_midhip_data(trc_file)

        trc_data_data = trc_data[:,1:]

        # Step 2: Normalize with reference marker position.
        with open(os.path.join(augmenterModelDir, "metadata.json"), 'r') as f:
            metadata = json.load(f)

        # Use midhip_data as the reference marker data
        referenceMarker_data = midhip_data  # instead of trc_file.marker(referenceMarker)

        norm_trc_data_data = np.zeros((trc_data_data.shape[0],
                                    trc_data_data.shape[1]))
        for i in range(0,trc_data_data.shape[1],3):
            norm_trc_data_data[:,i:i+3] = (trc_data_data[:,i:i+3] - 
                                        referenceMarker_data)
            
            
        # Step 3: Normalize with subject's height.
        norm2_trc_data_data = copy.deepcopy(norm_trc_data_data)
        norm2_trc_data_data = norm2_trc_data_data / subject_height
        
        # Step 4: Add remaining features.
        inputs = copy.deepcopy(norm2_trc_data_data)
        if featureHeight:    
            inputs = np.concatenate(
                (inputs, subject_height*np.ones((inputs.shape[0],1))), axis=1)
        if featureWeight:    
            inputs = np.concatenate(
                (inputs, subject_mass*np.ones((inputs.shape[0],1))), axis=1)
            
        # Step 5: Pre-process data
        pathMean = os.path.join(augmenterModelDir, "mean.npy")
        pathSTD = os.path.join(augmenterModelDir, "std.npy")
        if os.path.isfile(pathMean):
            trainFeatures_mean = np.load(pathMean, allow_pickle=True)
            inputs -= trainFeatures_mean
        if os.path.isfile(pathSTD):
            trainFeatures_std = np.load(pathSTD, allow_pickle=True)
            inputs /= trainFeatures_std 
            
        # Step 6: Reshape inputs if necessary (eg, LSTM)
        if augmenterModelName == "LSTM":
            inputs = np.reshape(inputs, (1, inputs.shape[0], inputs.shape[1]))
            
        # %% Load model and weights, and predict outputs.
        json_file = open(os.path.join(augmenterModelDir, "model.json"), 'r')
        pretrainedModel_json = json_file.read()
        json_file.close()
        model = tf.keras.models.model_from_json(pretrainedModel_json)
        model.load_weights(os.path.join(augmenterModelDir, "weights.h5"))  
        outputs = model.predict(inputs)
        
        # %% Post-process outputs.
        # Step 1: Reshape if necessary (eg, LSTM)
        if augmenterModelName == "LSTM":
            outputs = np.reshape(outputs, (outputs.shape[1], outputs.shape[2]))
            
        # Step 2: Un-normalize with subject's height.
        unnorm_outputs = outputs * subject_height
        
        # Step 2: Un-normalize with reference marker position.
        unnorm2_outputs = np.zeros((unnorm_outputs.shape[0],
                                    unnorm_outputs.shape[1]))
        for i in range(0,unnorm_outputs.shape[1],3):
            unnorm2_outputs[:,i:i+3] = (unnorm_outputs[:,i:i+3] + 
                                        referenceMarker_data)
            
        # %% Add markers to .trc file.
        for c, marker in enumerate(response_markers):
            x = unnorm2_outputs[:,c*3]
            y = unnorm2_outputs[:,c*3+1]
            z = unnorm2_outputs[:,c*3+2]
            trc_file.add_marker(marker, x, y, z)
            
        # %% Gather data for computing minimum y-position.
        outputs_all[idx_augm]['response_markers'] = response_markers   
        outputs_all[idx_augm]['response_data'] = unnorm2_outputs
        n_response_markers_all += len(response_markers)
        
    # %% Extract minimum y-position across response markers. This is used
    # to align feet and floor when visualizing.
    responses_all_conc = np.zeros((unnorm2_outputs.shape[0],
                                   n_response_markers_all*3))
    idx_acc_res = 0
    for idx_augm in outputs_all:
        idx_acc_res_end = (idx_acc_res + 
                           (len(outputs_all[idx_augm]['response_markers']))*3)
        responses_all_conc[:,idx_acc_res:idx_acc_res_end] = (
            outputs_all[idx_augm]['response_data'])
        idx_acc_res = idx_acc_res_end
    # Minimum y-position across response markers.
    min_y_pos = np.min(responses_all_conc[:,1::3])
        
    # %% If offset
    if offset:
        trc_file.offset('y', -(min_y_pos-0.01))
        
    # %% Return augmented .trc file   
    trc_file.write(pathOutputTRCFile)
    
    return min_y_pos