src/sata/t2m_wrapper.py

from torch_geometric.data import Batch
from .conversions.graph_to_motion import hatD_recon_motion
from fairmotion.data import bvh
from fairmotion.ops import motion as motion_ops
from .conversions.motion_to_graph import skel_2_graph, motion_2_graph
from .utils.motion_utils import motion_normalize_h2s
from .utils.bvh2joint import motion_to_joint_positions, JointTextProcessor, TextEncoder
from .utils.joint2humanml import JointToHumanML3D
import numpy as np
from os.path import join as pjoin

import torch
# TODO: fill path to reference_skeleton
def get_t2m_eval_wrapper(enc_model, enc_cfg, batch_size, ms_dict, reference_skeleton='./t2m_reference.bvh', evaluator_meta_dir=''):
    try:
        motion = bvh.load(reference_skeleton, ignore_root_skel=True, ee_as_joint=True)
    except Exception as e:
        print(f"Error loading {reference_skeleton}: {e}")
    
    motion, _ = motion_normalize_h2s(motion, False) 
    skel = motion.skel

    mean = np.load(pjoin(evaluator_meta_dir, 'mean.npy'))
    std = np.load(pjoin(evaluator_meta_dir, 'std.npy'))

    text_features = []
    text_processor = JointTextProcessor()
    text_encoder = TextEncoder()

    # Extract joint names
    joint_names = [joint.name for joint in motion.skel.joints]
    text_features = []
    for name in joint_names:
        name = text_processor.get_text(name, mode='descriptive')
        text_features.append(text_encoder.encode(name, reduce_mean=True))
    
    text_features = torch.cat(text_features, dim=0)
    skel_graph = skel_2_graph(skel, text_features)

    return T2MEvalWrapper(enc_model, skel_graph, batch_size, enc_cfg, ms_dict, mean, std)


class T2MEvalWrapper:
    def __init__(self, enc_model, ref_skel_graph, batch_size, enc_cfg, ms_dict, mean, std):
        self.enc_model = enc_model
        self.ref_skel_graph = ref_skel_graph.to(enc_model.device)
        self.batch_size = batch_size
        self.output_repr = enc_cfg["representation"]["out"]
        self.ms_dict = ms_dict
        self.mean = mean
        self.std = std

        self.enc_model.eval()
        self.converter = JointToHumanML3D(example_id="000021", data_dir='./HumanML3D/joints')
        # TODO: fill data_dir


    def decode(self, latent, valid_frames):
        B, T, C = latent.shape
        # build batched skel graph
        sliced_out_motion = []
        for i in range(B):
            skel_graph_batch = Batch.from_data_list([self.ref_skel_graph] * T)
            # Phase 1. Decode latent to fairmotion Motion
            hatD = self.enc_model.decode(latent[i], skel_graph_batch)
        
            # note that fairmotion will move motion to CPU
            out_motion_list, out_contact_list = hatD_recon_motion(
                hatD, skel_graph_batch, self.output_repr, self.ms_dict, T
            )
            out_motion = out_motion_list[0]  # Hardcode for T2M
            out_motion.set_fps(20)
            sliced_out_motion.append(out_motion)

        # Phase 2. Convert to joint positions
        joints_humanml3d_list = []
        for i in range(B):
            joints_position = motion_to_joint_positions(sliced_out_motion[i])
            # Phase 3. Convert to HumanML3D format
            joints_humanml3d = self.converter.convert(joints_position, valid_frames[i])
            joints_humanml3d_list.append(joints_humanml3d)

        joints_humanml3d = np.stack(joints_humanml3d_list, axis=0)  # (B, T, 3)
        
        # from IPython import embed; embed()
        # Phase 4. Normalize using HumanML3D mean and std
        joints_humanml3d = (joints_humanml3d - self.mean) / self.std

        # Phase 5. From numpy to torch tensor
        joints_humanml3d = torch.tensor(joints_humanml3d, device=self.enc_model.device).float()

        return joints_humanml3d

    def __getattr__(self, name):
        # Forward other attributes/methods to the encoder model.
        return getattr(self.enc_model, name)