include data folder

.gitignore

@@ -105,7 +105,7 @@ venv.bak/
 .mypy_cache/
 
 # custom
-data
+# data
 # data for pytest moved to http server
 # !tests/data
 .vscode

mGPT/data/HumanML3D.py

@@ -0,0 +1,119 @@
+import numpy as np
+import torch
+from os.path import join as pjoin
+import os
+from .humanml.utils.word_vectorizer import WordVectorizer
+from .humanml.scripts.motion_process import (process_file, recover_from_ric)
+from . import BASEDataModule
+from .humanml import Text2MotionDatasetEval, Text2MotionDataset, Text2MotionDatasetCB, MotionDataset, MotionDatasetVQ, Text2MotionDatasetToken, Text2MotionDatasetM2T
+from .utils import humanml3d_collate
+script_directory = os.path.dirname(os.path.abspath(__file__))
+class HumanML3DDataModule(BASEDataModule):
+    def __init__(self, cfg, **kwargs):
+
+        super().__init__(collate_fn=humanml3d_collate)
+        self.cfg = cfg
+        self.save_hyperparameters(logger=False)
+        
+        # Basic info of the dataset
+        cfg.DATASET.JOINT_TYPE = 'humanml3d'
+        self.name = "humanml3d"
+        self.njoints = 22
+        
+        # Path to the dataset
+        data_root = cfg.DATASET.HUMANML3D.ROOT
+        print(data_root)
+        assets_root = cfg.DATASET.HUMANML3D.ASSETS_ROOT
+        self.hparams.data_root = data_root
+        self.hparams.text_dir = pjoin(data_root, "texts")
+        self.hparams.motion_dir = pjoin(data_root, 'new_joint_vecs')
+        
+        # Mean and std of the dataset
+        self.hparams.mean = np.load(pjoin(script_directory, "mean.npy"))
+        self.hparams.std = np.load(pjoin(script_directory, "std.npy"))
+        
+        # Mean and std for fair evaluation
+        self.hparams.mean_eval = np.load(pjoin(script_directory, "mean_eval.npy"))
+        self.hparams.std_eval = np.load(pjoin(script_directory, "std_eval.npy"))
+        
+        # Length of the dataset
+        self.hparams.max_motion_length = cfg.DATASET.HUMANML3D.MAX_MOTION_LEN
+        self.hparams.min_motion_length = cfg.DATASET.HUMANML3D.MIN_MOTION_LEN
+        self.hparams.max_text_len = cfg.DATASET.HUMANML3D.MAX_TEXT_LEN
+        self.hparams.unit_length = cfg.DATASET.HUMANML3D.UNIT_LEN
+
+        # Additional parameters
+        self.hparams.debug = cfg.DEBUG
+        self.hparams.stage = cfg.TRAIN.STAGE
+
+        # Dataset switch
+        self.DatasetEval = Text2MotionDatasetEval
+
+        if cfg.TRAIN.STAGE == "vae":
+            if cfg.model.params.motion_vae.target.split('.')[-1].lower() == "vqvae":
+                self.hparams.win_size = 64
+                self.Dataset = MotionDatasetVQ
+            else:
+                self.Dataset = MotionDataset
+        elif 'lm' in cfg.TRAIN.STAGE:
+            self.hparams.code_path = cfg.DATASET.CODE_PATH
+            self.hparams.task_path = cfg.DATASET.TASK_PATH
+            self.hparams.std_text = cfg.DATASET.HUMANML3D.STD_TEXT
+            self.Dataset = Text2MotionDatasetCB
+        elif cfg.TRAIN.STAGE == "token":
+            self.Dataset = Text2MotionDatasetToken
+            self.DatasetEval = Text2MotionDatasetToken
+        elif cfg.TRAIN.STAGE == "m2t":
+            self.Dataset = Text2MotionDatasetM2T
+            self.DatasetEval = Text2MotionDatasetM2T
+        else:
+            self.Dataset = Text2MotionDataset
+
+        # Get additional info of the dataset
+        self.nfeats = 263
+        cfg.DATASET.NFEATS = self.nfeats
+        
+
+    def feats2joints(self, features):
+        mean = torch.tensor(self.hparams.mean).to(features)
+        std = torch.tensor(self.hparams.std).to(features)
+        features = features * std + mean
+        return recover_from_ric(features, self.njoints)
+
+    def joints2feats(self, features):
+        features = process_file(features, self.njoints)[0]
+        return features
+
+    def normalize(self, features):
+        mean = torch.tensor(self.hparams.mean).to(features)
+        std = torch.tensor(self.hparams.std).to(features)
+        features = (features - mean) / std
+        return features
+
+    def denormalize(self, features):
+        mean = torch.tensor(self.hparams.mean).to(features)
+        std = torch.tensor(self.hparams.std).to(features)
+        features = features * std + mean
+        return features
+
+    def renorm4t2m(self, features):
+        # renorm to t2m norms for using t2m evaluators
+        ori_mean = torch.tensor(self.hparams.mean).to(features)
+        ori_std = torch.tensor(self.hparams.std).to(features)
+        eval_mean = torch.tensor(self.hparams.mean_eval).to(features)
+        eval_std = torch.tensor(self.hparams.std_eval).to(features)
+        features = features * ori_std + ori_mean
+        features = (features - eval_mean) / eval_std
+        return features
+
+    def mm_mode(self, mm_on=True):
+        if mm_on:
+            self.is_mm = True
+            self.name_list = self.test_dataset.name_list
+            self.mm_list = np.random.choice(self.name_list,
+                                            self.cfg.METRIC.MM_NUM_SAMPLES,
+                                            replace=False)
+            self.test_dataset.name_list = self.mm_list
+        else:
+            self.is_mm = False
+            self.test_dataset.name_list = self.name_list

mGPT/data/Kit.py

@@ -0,0 +1,88 @@
+import numpy as np
+import torch
+from os.path import join as pjoin
+from .humanml.utils.word_vectorizer import WordVectorizer
+from .humanml.scripts.motion_process import (process_file, recover_from_ric)
+from .HumanML3D import HumanML3DDataModule
+from .humanml import Text2MotionDatasetEval, Text2MotionDataset, Text2MotionDatasetCB, MotionDataset, MotionDatasetVQ, Text2MotionDatasetToken
+
+
+class KitDataModule(HumanML3DDataModule):
+    def __init__(self, cfg, **kwargs):
+
+        super().__init__(cfg, **kwargs)
+
+        # Basic info of the dataset
+        self.name = "kit"
+        self.njoints = 21
+
+        # Path to the dataset
+        data_root = cfg.DATASET.KIT.ROOT
+        self.hparams.data_root = data_root
+        self.hparams.text_dir = pjoin(data_root, "texts")
+        self.hparams.motion_dir = pjoin(data_root, 'new_joint_vecs')
+
+        # Mean and std of the dataset
+        dis_data_root = pjoin(cfg.DATASET.KIT.MEAN_STD_PATH, 'kit',
+                              "VQVAEV3_CB1024_CMT_H1024_NRES3", "meta")
+        self.hparams.mean = np.load(pjoin(dis_data_root, "mean.npy"))
+        self.hparams.std = np.load(pjoin(dis_data_root, "std.npy"))
+
+        # Mean and std for fair evaluation
+        dis_data_root_eval = pjoin(cfg.DATASET.KIT.MEAN_STD_PATH, 't2m',
+                                   "Comp_v6_KLD005", "meta")
+        self.hparams.mean_eval = np.load(pjoin(dis_data_root_eval, "mean.npy"))
+        self.hparams.std_eval = np.load(pjoin(dis_data_root_eval, "std.npy"))
+
+        # Length of the dataset
+        self.hparams.max_motion_length = cfg.DATASET.KIT.MAX_MOTION_LEN
+        self.hparams.min_motion_length = cfg.DATASET.KIT.MIN_MOTION_LEN
+        self.hparams.max_text_len = cfg.DATASET.KIT.MAX_TEXT_LEN
+        self.hparams.unit_length = cfg.DATASET.KIT.UNIT_LEN
+        
+        # Get additional info of the dataset
+        self._sample_set = self.get_sample_set(overrides={"split": "test", "tiny": True})
+        self.nfeats = self._sample_set.nfeats
+        cfg.DATASET.NFEATS = self.nfeats
+
+    def feats2joints(self, features):
+        mean = torch.tensor(self.hparams.mean).to(features)
+        std = torch.tensor(self.hparams.std).to(features)
+        features = features * std + mean
+        return recover_from_ric(features, self.njoints)
+
+    def joints2feats(self, features):
+        features = process_file(features, self.njoints)[0]
+        # mean = torch.tensor(self.hparams.mean).to(features)
+        # std = torch.tensor(self.hparams.std).to(features)
+        # features = (features - mean) / std
+        return features
+
+    def normalize(self, features):
+        mean = torch.tensor(self.hparams.mean).to(features)
+        std = torch.tensor(self.hparams.std).to(features)
+        features = (features - mean) / std
+        return features
+
+    def renorm4t2m(self, features):
+        # renorm to t2m norms for using t2m evaluators
+        ori_mean = torch.tensor(self.hparams.mean).to(features)
+        ori_std = torch.tensor(self.hparams.std).to(features)
+        eval_mean = torch.tensor(self.hparams.mean_eval).to(features)
+        eval_std = torch.tensor(self.hparams.std_eval).to(features)
+        features = features * ori_std + ori_mean
+        features = (features - eval_mean) / eval_std
+        return features
+
+    def mm_mode(self, mm_on=True):
+        # random select samples for mm
+        if mm_on:
+            self.is_mm = True
+            self.name_list = self.test_dataset.name_list
+            self.mm_list = np.random.choice(self.name_list,
+                                            self.cfg.METRIC.MM_NUM_SAMPLES,
+                                            replace=False)
+            self.test_dataset.name_list = self.mm_list
+        else:
+            self.is_mm = False
+            self.test_dataset.name_list = self.name_list

mGPT/data/__init__.py

@@ -0,0 +1,103 @@
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+
+
+class BASEDataModule(pl.LightningDataModule):
+    def __init__(self, collate_fn):
+        super().__init__()
+
+        self.dataloader_options = {"collate_fn": collate_fn}
+        self.persistent_workers = True
+        self.is_mm = False
+
+        self._train_dataset = None
+        self._val_dataset = None
+        self._test_dataset = None
+
+    def get_sample_set(self, overrides={}):
+        sample_params = self.hparams.copy()
+        sample_params.update(overrides)
+        return self.DatasetEval(**sample_params)
+
+    @property
+    def train_dataset(self):
+        if self._train_dataset is None:
+            self._train_dataset = self.Dataset(split=self.cfg.TRAIN.SPLIT,
+                                               **self.hparams)
+        return self._train_dataset
+
+    @property
+    def val_dataset(self):
+        if self._val_dataset is None:
+            params = self.hparams.copy()
+            params['code_path'] = None
+            params['split'] = self.cfg.EVAL.SPLIT
+            self._val_dataset = self.DatasetEval(**params)
+        return self._val_dataset
+
+    @property
+    def test_dataset(self):
+        if self._test_dataset is None:
+            # self._test_dataset = self.DatasetEval(split=self.cfg.TEST.SPLIT,
+            #                                       **self.hparams)
+            params = self.hparams.copy()
+            params['code_path'] = None
+            params['split'] = self.cfg.TEST.SPLIT
+            self._test_dataset = self.DatasetEval( **params)
+        return self._test_dataset
+
+    def setup(self, stage=None):
+        # Use the getter the first time to load the data
+        if stage in (None, "fit"):
+            _ = self.train_dataset
+            _ = self.val_dataset
+        if stage in (None, "test"):
+            _ = self.test_dataset
+
+    def train_dataloader(self):
+        dataloader_options = self.dataloader_options.copy()
+        dataloader_options["batch_size"] = self.cfg.TRAIN.BATCH_SIZE
+        dataloader_options["num_workers"] = self.cfg.TRAIN.NUM_WORKERS
+        return DataLoader(
+            self.train_dataset,
+            shuffle=False,
+            persistent_workers=True,
+            **dataloader_options,
+        )
+
+    def predict_dataloader(self):
+        dataloader_options = self.dataloader_options.copy()
+        dataloader_options[
+            "batch_size"] = 1 if self.is_mm else self.cfg.TEST.BATCH_SIZE
+        dataloader_options["num_workers"] = self.cfg.TEST.NUM_WORKERS
+        dataloader_options["shuffle"] = False
+        return DataLoader(
+            self.test_dataset,
+            persistent_workers=True,
+            **dataloader_options,
+        )
+
+    def val_dataloader(self):
+        # overrides batch_size and num_workers
+        dataloader_options = self.dataloader_options.copy()
+        dataloader_options["batch_size"] = self.cfg.EVAL.BATCH_SIZE
+        dataloader_options["num_workers"] = self.cfg.EVAL.NUM_WORKERS
+        dataloader_options["shuffle"] = False
+        return DataLoader(
+            self.val_dataset,
+            persistent_workers=True,
+            **dataloader_options,
+        )
+
+    def test_dataloader(self):
+        # overrides batch_size and num_workers
+        dataloader_options = self.dataloader_options.copy()
+        dataloader_options[
+            "batch_size"] = 1 if self.is_mm else self.cfg.TEST.BATCH_SIZE
+        dataloader_options["num_workers"] = self.cfg.TEST.NUM_WORKERS
+        dataloader_options["shuffle"] = False
+        return DataLoader(
+            self.test_dataset,
+            persistent_workers=True,
+            **dataloader_options,
+        )

mGPT/data/build_data.py

@@ -0,0 +1,15 @@
+from omegaconf import OmegaConf
+from os.path import join as pjoin
+from ..config import instantiate_from_config
+
+
+def build_data(cfg, phase="train"):
+    data_config = OmegaConf.to_container(cfg.DATASET, resolve=True)
+    data_config['params'] = {'cfg': cfg, 'phase': phase}
+    if isinstance(data_config['target'], str):
+        return instantiate_from_config(data_config)
+    elif isinstance(data_config['target'], list):
+        data_config_tmp = data_config.copy()
+        data_config_tmp['params']['dataModules'] = data_config['target']
+        data_config_tmp['target'] = 'mGPT.data.Concat.ConcatDataModule'
+        return instantiate_from_config(data_config)

mGPT/data/humanml/README.md

@@ -0,0 +1 @@
+This code is based on https://github.com/EricGuo5513/text-to-motion.git

mGPT/data/humanml/__init__.py

@@ -0,0 +1,7 @@
+from .dataset_t2m import Text2MotionDataset
+from .dataset_t2m_eval import Text2MotionDatasetEval
+from .dataset_t2m_cb import Text2MotionDatasetCB
+from .dataset_t2m_token import Text2MotionDatasetToken
+from .dataset_t2m_m2t import Text2MotionDatasetM2T
+from .dataset_m import MotionDataset
+from .dataset_m_vq import MotionDatasetVQ

mGPT/data/humanml/common/quaternion.py

@@ -0,0 +1,423 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import torch
+import numpy as np
+
+_EPS4 = np.finfo(float).eps * 4.0
+
+_FLOAT_EPS = np.finfo(np.float64).eps
+
+# PyTorch-backed implementations
+def qinv(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+
+def qinv_np(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return qinv(torch.from_numpy(q).float()).numpy()
+
+
+def qnormalize(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return q / torch.norm(q, dim=-1, keepdim=True)
+
+
+def qmul(q, r):
+    """
+    Multiply quaternion(s) q with quaternion(s) r.
+    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
+    Returns q*r as a tensor of shape (*, 4).
+    """
+    assert q.shape[-1] == 4
+    assert r.shape[-1] == 4
+
+    original_shape = q.shape
+
+    # Compute outer product
+    terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
+
+    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
+    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
+    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
+    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
+    return torch.stack((w, x, y, z), dim=1).view(original_shape)
+
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    # print(q.shape)
+    q = q.contiguous().view(-1, 4)
+    v = v.contiguous().view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+
+def qeuler(q, order, epsilon=0, deg=True):
+    """
+    Convert quaternion(s) q to Euler angles.
+    Expects a tensor of shape (*, 4), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+
+    original_shape = list(q.shape)
+    original_shape[-1] = 3
+    q = q.view(-1, 4)
+
+    q0 = q[:, 0]
+    q1 = q[:, 1]
+    q2 = q[:, 2]
+    q3 = q[:, 3]
+
+    if order == 'xyz':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    elif order == 'yzx':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon))
+    elif order == 'zxy':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'xzy':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon))
+    elif order == 'yxz':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'zyx':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    else:
+        raise
+
+    if deg:
+        return torch.stack((x, y, z), dim=1).view(original_shape) * 180 / np.pi
+    else:
+        return torch.stack((x, y, z), dim=1).view(original_shape)
+
+
+# Numpy-backed implementations
+
+def qmul_np(q, r):
+    q = torch.from_numpy(q).contiguous().float()
+    r = torch.from_numpy(r).contiguous().float()
+    return qmul(q, r).numpy()
+
+
+def qrot_np(q, v):
+    q = torch.from_numpy(q).contiguous().float()
+    v = torch.from_numpy(v).contiguous().float()
+    return qrot(q, v).numpy()
+
+
+def qeuler_np(q, order, epsilon=0, use_gpu=False):
+    if use_gpu:
+        q = torch.from_numpy(q).cuda().float()
+        return qeuler(q, order, epsilon).cpu().numpy()
+    else:
+        q = torch.from_numpy(q).contiguous().float()
+        return qeuler(q, order, epsilon).numpy()
+
+
+def qfix(q):
+    """
+    Enforce quaternion continuity across the time dimension by selecting
+    the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)
+    between two consecutive frames.
+
+    Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.
+    Returns a tensor of the same shape.
+    """
+    assert len(q.shape) == 3
+    assert q.shape[-1] == 4
+
+    result = q.copy()
+    dot_products = np.sum(q[1:] * q[:-1], axis=2)
+    mask = dot_products < 0
+    mask = (np.cumsum(mask, axis=0) % 2).astype(bool)
+    result[1:][mask] *= -1
+    return result
+
+
+def euler2quat(e, order, deg=True):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.view(-1, 3)
+
+    ## if euler angles in degrees
+    if deg:
+        e = e * np.pi / 180.
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1)
+    ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1)
+    rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.view(original_shape)
+
+
+def expmap_to_quaternion(e):
+    """
+    Convert axis-angle rotations (aka exponential maps) to quaternions.
+    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
+    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 4).
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+    e = e.reshape(-1, 3)
+
+    theta = np.linalg.norm(e, axis=1).reshape(-1, 1)
+    w = np.cos(0.5 * theta).reshape(-1, 1)
+    xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e
+    return np.concatenate((w, xyz), axis=1).reshape(original_shape)
+
+
+def euler_to_quaternion(e, order):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.reshape(-1, 3)
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1)
+    ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1)
+    rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul_np(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.reshape(original_shape)
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def quaternion_to_matrix_np(quaternions):
+    q = torch.from_numpy(quaternions).contiguous().float()
+    return quaternion_to_matrix(q).numpy()
+
+
+def quaternion_to_cont6d_np(quaternions):
+    rotation_mat = quaternion_to_matrix_np(quaternions)
+    cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1)
+    return cont_6d
+
+
+def quaternion_to_cont6d(quaternions):
+    rotation_mat = quaternion_to_matrix(quaternions)
+    cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1)
+    return cont_6d
+
+
+def cont6d_to_matrix(cont6d):
+    assert cont6d.shape[-1] == 6, "The last dimension must be 6"
+    x_raw = cont6d[..., 0:3]
+    y_raw = cont6d[..., 3:6]
+
+    x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)
+    z = torch.cross(x, y_raw, dim=-1)
+    z = z / torch.norm(z, dim=-1, keepdim=True)
+
+    y = torch.cross(z, x, dim=-1)
+
+    x = x[..., None]
+    y = y[..., None]
+    z = z[..., None]
+
+    mat = torch.cat([x, y, z], dim=-1)
+    return mat
+
+
+def cont6d_to_matrix_np(cont6d):
+    q = torch.from_numpy(cont6d).contiguous().float()
+    return cont6d_to_matrix(q).numpy()
+
+
+def qpow(q0, t, dtype=torch.float):
+    ''' q0 : tensor of quaternions
+    t: tensor of powers
+    '''
+    q0 = qnormalize(q0)
+    theta0 = torch.acos(q0[..., 0])
+
+    ## if theta0 is close to zero, add epsilon to avoid NaNs
+    mask = (theta0 <= 10e-10) * (theta0 >= -10e-10)
+    theta0 = (1 - mask) * theta0 + mask * 10e-10
+    v0 = q0[..., 1:] / torch.sin(theta0).view(-1, 1)
+
+    if isinstance(t, torch.Tensor):
+        q = torch.zeros(t.shape + q0.shape)
+        theta = t.view(-1, 1) * theta0.view(1, -1)
+    else:  ## if t is a number
+        q = torch.zeros(q0.shape)
+        theta = t * theta0
+
+    q[..., 0] = torch.cos(theta)
+    q[..., 1:] = v0 * torch.sin(theta).unsqueeze(-1)
+
+    return q.to(dtype)
+
+
+def qslerp(q0, q1, t):
+    '''
+    q0: starting quaternion
+    q1: ending quaternion
+    t: array of points along the way
+
+    Returns:
+    Tensor of Slerps: t.shape + q0.shape
+    '''
+
+    q0 = qnormalize(q0)
+    q1 = qnormalize(q1)
+    q_ = qpow(qmul(q1, qinv(q0)), t)
+
+    return qmul(q_,
+                q0.contiguous().view(torch.Size([1] * len(t.shape)) + q0.shape).expand(t.shape + q0.shape).contiguous())
+
+
+def qbetween(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v = torch.cross(v0, v1)
+    w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1,
+                                                                                                              keepdim=True)
+    return qnormalize(torch.cat([w, v], dim=-1))
+
+
+def qbetween_np(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v0 = torch.from_numpy(v0).float()
+    v1 = torch.from_numpy(v1).float()
+    return qbetween(v0, v1).numpy()
+
+
+def lerp(p0, p1, t):
+    if not isinstance(t, torch.Tensor):
+        t = torch.Tensor([t])
+
+    new_shape = t.shape + p0.shape
+    new_view_t = t.shape + torch.Size([1] * len(p0.shape))
+    new_view_p = torch.Size([1] * len(t.shape)) + p0.shape
+    p0 = p0.view(new_view_p).expand(new_shape)
+    p1 = p1.view(new_view_p).expand(new_shape)
+    t = t.view(new_view_t).expand(new_shape)
+
+    return p0 + t * (p1 - p0)

mGPT/data/humanml/common/skeleton.py

@@ -0,0 +1,199 @@
+from .quaternion import *
+import scipy.ndimage.filters as filters
+
+class Skeleton(object):
+    def __init__(self, offset, kinematic_tree, device):
+        self.device = device
+        self._raw_offset_np = offset.numpy()
+        self._raw_offset = offset.clone().detach().to(device).float()
+        self._kinematic_tree = kinematic_tree
+        self._offset = None
+        self._parents = [0] * len(self._raw_offset)
+        self._parents[0] = -1
+        for chain in self._kinematic_tree:
+            for j in range(1, len(chain)):
+                self._parents[chain[j]] = chain[j-1]
+
+    def njoints(self):
+        return len(self._raw_offset)
+
+    def offset(self):
+        return self._offset
+
+    def set_offset(self, offsets):
+        self._offset = offsets.clone().detach().to(self.device).float()
+
+    def kinematic_tree(self):
+        return self._kinematic_tree
+
+    def parents(self):
+        return self._parents
+
+    # joints (batch_size, joints_num, 3)
+    def get_offsets_joints_batch(self, joints):
+        assert len(joints.shape) == 3
+        _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # joints (joints_num, 3)
+    def get_offsets_joints(self, joints):
+        assert len(joints.shape) == 2
+        _offsets = self._raw_offset.clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            # print(joints.shape)
+            _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder
+    # joints (batch_size, joints_num, 3)
+    def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False):
+        assert len(face_joint_idx) == 4
+        '''Get Forward Direction'''
+        l_hip, r_hip, sdr_r, sdr_l = face_joint_idx
+        across1 = joints[:, r_hip] - joints[:, l_hip]
+        across2 = joints[:, sdr_r] - joints[:, sdr_l]
+        across = across1 + across2
+        across = across / np.sqrt((across**2).sum(axis=-1))[:, np.newaxis]
+        # print(across1.shape, across2.shape)
+
+        # forward (batch_size, 3)
+        forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+        if smooth_forward:
+            forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest')
+            # forward (batch_size, 3)
+        forward = forward / np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis]
+
+        '''Get Root Rotation'''
+        target = np.array([[0,0,1]]).repeat(len(forward), axis=0)
+        root_quat = qbetween_np(forward, target)
+
+        '''Inverse Kinematics'''
+        # quat_params (batch_size, joints_num, 4)
+        # print(joints.shape[:-1])
+        quat_params = np.zeros(joints.shape[:-1] + (4,))
+        # print(quat_params.shape)
+        root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        quat_params[:, 0] = root_quat
+        # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        for chain in self._kinematic_tree:
+            R = root_quat
+            for j in range(len(chain) - 1):
+                # (batch, 3)
+                u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0)
+                # print(u.shape)
+                # (batch, 3)
+                v = joints[:, chain[j+1]] - joints[:, chain[j]]
+                v = v / np.sqrt((v**2).sum(axis=-1))[:, np.newaxis]
+                # print(u.shape, v.shape)
+                rot_u_v = qbetween_np(u, v)
+
+                R_loc = qmul_np(qinv_np(R), rot_u_v)
+
+                quat_params[:,chain[j + 1], :] = R_loc
+                R = qmul_np(R, R_loc)
+
+        return quat_params
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        joints = torch.zeros(quat_params.shape[:-1] + (3,)).to(self.device)
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(len(quat_params), -1).detach().to(self.device)
+            for i in range(1, len(chain)):
+                R = qmul(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot(R, offset_vec) + joints[:, chain[i-1]]
+        return joints
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics_np(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(quat_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(len(quat_params), axis=0)
+            for i in range(1, len(chain)):
+                R = qmul_np(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot_np(R, offset_vec) + joints[:, chain[i - 1]]
+        return joints
+
+    def forward_kinematics_cont6d_np(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(cont6d_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix_np(cont6d_params[:, 0])
+            else:
+                matR = np.eye(3)[np.newaxis, :].repeat(len(cont6d_params), axis=0)
+            for i in range(1, len(chain)):
+                matR = np.matmul(matR, cont6d_to_matrix_np(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]][..., np.newaxis]
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = np.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints
+
+    def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            # skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device)
+        joints[..., 0, :] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix(cont6d_params[:, 0])
+            else:
+                matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device)
+            for i in range(1, len(chain)):
+                matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]].unsqueeze(-1)
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints
+
+
+
+
+

mGPT/data/humanml/dataset_m.py

@@ -0,0 +1,156 @@
+import os
+import rich
+import random
+import pickle
+import codecs as cs
+import numpy as np
+from torch.utils import data
+from rich.progress import track
+from os.path import join as pjoin
+
+
+class MotionDataset(data.Dataset):
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length=196,
+        min_motion_length=20,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+        
+        # restrian the length of motion and text
+        self.max_motion_length = max_motion_length
+        self.min_motion_length = min_motion_length
+        self.unit_length = unit_length
+        
+        # Data mean and std
+        self.mean = mean
+        self.std = std
+        
+        # Data path
+        split_file = pjoin(data_root, split + '.txt')
+        motion_dir = pjoin(data_root, 'new_joint_vecs')
+        text_dir = pjoin(data_root, 'texts')
+
+        # Data id list
+        self.id_list = []     
+        with cs.open(split_file, "r") as f:
+            for line in f.readlines():
+                self.id_list.append(line.strip())
+                
+        # Debug mode
+        if tiny or debug:
+            enumerator = enumerate(
+                track(
+                    self.id_list,
+                    f"Loading HumanML3D {split}",
+                ))
+            maxdata = 100
+            subset = '_tiny'
+        else:
+            enumerator = enumerate(self.id_list)
+            maxdata = 1e10
+            subset = ''
+
+        new_name_list = []
+        motion_dict = {}
+
+        # Fast loading
+        if os.path.exists(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl')):
+            with rich.progress.open(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl'),
+                    'rb', description=f"Loading HumanML3D {split}") as file:
+                motion_dict = pickle.load(file)
+            with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'), 'rb') as file:
+                new_name_list = pickle.load(file)
+        else:
+            for idx, name in enumerator:
+                if len(new_name_list) > maxdata:
+                    break
+                try:
+                    motion = [np.load(pjoin(motion_dir, name + ".npy"))]
+
+                    # Read text
+                    with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                        text_data = []
+                        flag = False
+                        lines = f.readlines()
+
+                        for line in lines:
+                            try:
+                                line_split = line.strip().split('#')
+                                f_tag = float(line_split[2])
+                                to_tag = float(line_split[3])
+                                f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                                to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                                if f_tag == 0.0 and to_tag == 0.0:
+                                    flag = True
+                                else:
+                                    motion_new = [tokens[int(f_tag*fps/unit_length) : int(to_tag*fps/unit_length)] for tokens in motion if int(f_tag*fps/unit_length) < int(to_tag*fps/unit_length)]
+
+                                    if len(motion_new) == 0:
+                                        continue
+                                    new_name = '%s_%f_%f'%(name, f_tag, to_tag)
+
+                                    motion_dict[new_name] = {
+                                        'motion': motion_new, 
+                                        "length": [len(m[0]) for m in motion_new]}
+                                    new_name_list.append(new_name)
+                            except:
+                                pass
+
+                    if flag:
+                        motion_dict[name] = {
+                            'motion': motion, 
+                            "length": [len(m[0]) for m in motion]}
+                        new_name_list.append(name)
+                except:
+                    pass
+                
+            if tmpFile:
+                os.makedirs(pjoin(data_root, 'tmp'), exist_ok=True)
+                
+                with open(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl'),'wb') as file:
+                    pickle.dump(motion_dict, file)
+                with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'), 'wb') as file:
+                    pickle.dump(new_name_list, file)
+
+        self.motion_dict = motion_dict
+        self.name_list = new_name_list
+        self.nfeats = motion_dict[new_name_list[0]]['motion'][0].shape[1]
+
+    def __len__(self):
+        return len(self.name_list)
+
+    def __getitem__(self, item):
+        data = self.motion_dict[self.name_list[item]]
+        motion_list, m_length = data["motion"], data["length"]
+
+        # Randomly select a motion
+        motion = random.choice(motion_list)
+
+        # Crop the motions in to times of 4, and introduce small variations
+        if self.unit_length < 10:
+            coin2 = np.random.choice(["single", "single", "double"])
+        else:
+            coin2 = "single"
+
+        if coin2 == "double":
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == "single":
+            m_length = (m_length // self.unit_length) * self.unit_length
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx + m_length]
+        
+        # Z Normalization
+        motion = (motion - self.mean) / self.std
+
+        return None, motion, m_length, None, None, None, None,

mGPT/data/humanml/dataset_m_vq.py

@@ -0,0 +1,54 @@
+import random
+import codecs as cs
+import numpy as np
+from torch.utils import data
+from rich.progress import track
+from os.path import join as pjoin
+from .dataset_m import MotionDataset
+from .dataset_t2m import Text2MotionDataset
+
+
+class MotionDatasetVQ(Text2MotionDataset):
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length,
+        min_motion_length,
+        win_size,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+        super().__init__(data_root, split, mean, std, max_motion_length,
+                         min_motion_length, unit_length, fps, tmpFile, tiny,
+                         debug, **kwargs)
+
+        # Filter out the motions that are too short
+        self.window_size = win_size
+        name_list = list(self.name_list)
+        for name in self.name_list:
+            motion = self.data_dict[name]["motion"]
+            if motion.shape[0] < self.window_size:
+                name_list.remove(name)
+                self.data_dict.pop(name)
+        self.name_list = name_list
+
+    def __len__(self):
+        return len(self.name_list)
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, length = data["motion"], data["length"]
+
+        idx = random.randint(0, motion.shape[0] - self.window_size)
+        motion = motion[idx:idx + self.window_size]
+        motion = (motion - self.mean) / self.std
+
+        return None, motion, length, None, None, None, None,

mGPT/data/humanml/dataset_t2m.py

@@ -0,0 +1,211 @@
+import os
+import rich
+import random
+import pickle
+import codecs as cs
+import numpy as np
+from torch.utils import data
+from rich.progress import track
+from os.path import join as pjoin
+
+
+class Text2MotionDataset(data.Dataset):
+
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length=196,
+        min_motion_length=40,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+
+        # restrian the length of motion and text
+        self.max_length = 20
+        self.max_motion_length = max_motion_length
+        self.min_motion_length = min_motion_length
+        self.unit_length = unit_length
+
+        # Data mean and std
+        self.mean = mean
+        self.std = std
+
+        # Data path
+        split_file = pjoin(data_root, split + '.txt')
+        motion_dir = pjoin(data_root, 'new_joint_vecs')
+        text_dir = pjoin(data_root, 'texts')
+
+        # Data id list
+        self.id_list = []
+        with cs.open(split_file, "r") as f:
+            for line in f.readlines():
+                self.id_list.append(line.strip())
+
+        # Debug mode
+        if tiny or debug:
+            enumerator = enumerate(self.id_list)
+            maxdata = 100
+            subset = '_tiny'
+        else:
+            enumerator = enumerate(
+                track(
+                    self.id_list,
+                    f"Loading HumanML3D {split}",
+                ))
+            maxdata = 1e10
+            subset = ''
+
+        new_name_list = []
+        length_list = []
+        data_dict = {}
+
+        # Fast loading
+        if os.path.exists(pjoin(data_root, f'tmp/{split}{subset}_data.pkl')):
+            if tiny or debug:
+                with open(pjoin(data_root, f'tmp/{split}{subset}_data.pkl'),
+                          'rb') as file:
+                    data_dict = pickle.load(file)
+            else:
+                with rich.progress.open(
+                        pjoin(data_root, f'tmp/{split}{subset}_data.pkl'),
+                        'rb',
+                        description=f"Loading HumanML3D {split}") as file:
+                    data_dict = pickle.load(file)
+            with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'),
+                      'rb') as file:
+                name_list = pickle.load(file)
+            for name in new_name_list:
+                length_list.append(data_dict[name]['length'])
+
+        else:
+            for idx, name in enumerator:
+                if len(new_name_list) > maxdata:
+                    break
+                try:
+                    motion = np.load(pjoin(motion_dir, name + ".npy"))
+                    if (len(motion)) < self.min_motion_length or (len(motion)
+                                                                  >= 200):
+                        continue
+
+                    # Read text
+                    text_data = []
+                    flag = False
+                    with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                        lines = f.readlines()
+                        for line in lines:
+                            text_dict = {}
+                            line_split = line.strip().split('#')
+                            caption = line_split[0]
+                            t_tokens = line_split[1].split(' ')
+                            f_tag = float(line_split[2])
+                            to_tag = float(line_split[3])
+                            f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                            to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                            text_dict['caption'] = caption
+                            text_dict['tokens'] = t_tokens
+                            if f_tag == 0.0 and to_tag == 0.0:
+                                flag = True
+                                text_data.append(text_dict)
+                            else:
+                                motion_new = motion[int(f_tag *
+                                                        fps):int(to_tag * fps)]
+                                if (len(motion_new)
+                                    ) < self.min_motion_length or (
+                                        len(motion_new) >= 200):
+                                    continue
+                                new_name = random.choice(
+                                    'ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                while new_name in new_name_list:
+                                    new_name = random.choice(
+                                        'ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                name_count = 1
+                                while new_name in data_dict:
+                                    new_name += '_' + name_count
+                                    name_count += 1
+                                data_dict[new_name] = {
+                                    'motion': motion_new,
+                                    "length": len(motion_new),
+                                    'text': [text_dict]
+                                }
+                                new_name_list.append(new_name)
+                                length_list.append(len(motion_new))
+
+                    if flag:
+                        data_dict[name] = {
+                            'motion': motion,
+                            "length": len(motion),
+                            'text': text_data
+                        }
+                        new_name_list.append(name)
+                        length_list.append(len(motion))
+                except:
+                    pass
+
+            name_list, length_list = zip(
+                *sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+
+            if tmpFile:
+                os.makedirs(pjoin(data_root, 'tmp'), exist_ok=True)
+                with open(pjoin(data_root, f'tmp/{split}{subset}_data.pkl'),
+                          'wb') as file:
+                    pickle.dump(data_dict, file)
+                with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'),
+                          'wb') as file:
+                    pickle.dump(name_list, file)
+
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+        self.nfeats = data_dict[name_list[0]]['motion'].shape[1]
+        self.reset_max_len(self.max_length)
+
+    def reset_max_len(self, length):
+        assert length <= self.max_motion_length
+        self.pointer = np.searchsorted(self.length_arr, length)
+        print("Pointer Pointing at %d" % self.pointer)
+        self.max_length = length
+
+    def __len__(self):
+        return len(self.name_list) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data["motion"], data["length"], data[
+            "text"]
+
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption = text_data["caption"]
+
+        all_captions = [
+            ' '.join([token.split('/')[0] for token in text_dic['tokens']])
+            for text_dic in text_list
+        ]
+
+        # Crop the motions in to times of 4, and introduce small variations
+        if self.unit_length < 10:
+            coin2 = np.random.choice(["single", "single", "double"])
+        else:
+            coin2 = "single"
+
+        if coin2 == "double":
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == "single":
+            m_length = (m_length // self.unit_length) * self.unit_length
+
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx + m_length]
+
+        # Z Normalization
+        motion = (motion - self.mean) / self.std
+
+        return caption, motion, m_length, None, None, None, None, all_captions

mGPT/data/humanml/dataset_t2m_cb.py

@@ -0,0 +1,211 @@
+import rich
+import random
+import pickle
+import os
+import numpy as np
+import codecs as cs
+from torch.utils import data
+from os.path import join as pjoin
+from rich.progress import track
+import json
+import spacy
+
+class Text2MotionDatasetCB(data.Dataset):
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length=196,
+        min_motion_length=20,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        stage='lm_pretrain',
+        code_path='VQVAE',
+        task_path=None,
+        std_text=False,
+        **kwargs,
+    ):
+        self.tiny = tiny
+        self.unit_length = unit_length
+
+        # Data mean and std
+        self.mean = mean
+        self.std = std
+
+        # Data path
+        split = 'train'
+        split_file = pjoin(data_root, split + '.txt')
+        motion_dir = pjoin(data_root, code_path)
+        text_dir = pjoin(data_root, 'texts')
+        
+        if task_path:
+            instructions = task_path
+        elif stage == 'lm_pretrain':
+            instructions = pjoin(data_root, 'template_pretrain.json')
+        elif stage in ['lm_instruct', "lm_rl"]:
+            instructions = pjoin(data_root, 'template_instructions.json')
+        else:
+            raise NotImplementedError(f"stage {stage} not implemented")
+
+        # Data id list
+        self.id_list = []
+        with cs.open(split_file, "r") as f:
+            for line in f.readlines():
+                self.id_list.append(line.strip())
+
+        # Debug mode
+        if tiny or debug:
+            enumerator = enumerate(self.id_list)
+            maxdata = 100
+            subset = '_tiny'
+        else:
+            enumerator = enumerate(
+                track(
+                    self.id_list,
+                    f"Loading HumanML3D {split}",
+                ))
+            maxdata = 1e10
+            subset = ''
+
+        new_name_list = []
+        data_dict = {}
+
+        # Fast loading
+        for i, name in enumerator:
+            if len(new_name_list) > maxdata:
+                break
+            try:
+                # Load motion tokens
+                m_token_list = np.load(pjoin(motion_dir, f'{name}.npy'))
+                # Read text
+                with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                    text_data = []
+                    flag = False
+                    lines = f.readlines()
+
+                    for line in lines:
+                        try:
+                            text_dict = {}
+                            line_split = line.strip().split('#')
+                            caption = line_split[0]
+                            t_tokens = line_split[1].split(' ')
+                            f_tag = float(line_split[2])
+                            to_tag = float(line_split[3])
+                            f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                            to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                            text_dict['caption'] = caption
+                            text_dict['tokens'] = t_tokens
+                            if f_tag == 0.0 and to_tag == 0.0:
+                                flag = True
+                                text_data.append(text_dict)
+                            else:
+                                m_token_list_new = [
+                                    tokens[int(f_tag * fps / unit_length
+                                                ):int(to_tag * fps /
+                                                        unit_length)]
+                                    for tokens in m_token_list
+                                    if int(f_tag * fps / unit_length) <
+                                    int(to_tag * fps / unit_length)
+                                ]
+
+                                if len(m_token_list_new) == 0:
+                                    continue
+                                new_name = '%s_%f_%f' % (name, f_tag,
+                                                            to_tag)
+
+                                data_dict[new_name] = {
+                                    'm_token_list': m_token_list_new,
+                                    'text': [text_dict]
+                                }
+                                new_name_list.append(new_name)
+                        except:
+                            pass
+
+                if flag:
+                    data_dict[name] = {
+                        'm_token_list': m_token_list,
+                        'text': text_data
+                    }
+                    new_name_list.append(name)
+            except:
+                pass
+
+        if tmpFile:
+            os.makedirs(pjoin(data_root, 'tmp'), exist_ok=True)
+            with open(
+                    pjoin(data_root,
+                            f'tmp/{split}{subset}_tokens_data.pkl'),
+                    'wb') as file:
+                pickle.dump(data_dict, file)
+            with open(
+                    pjoin(data_root,
+                            f'tmp/{split}{subset}_tokens_index.pkl'),
+                    'wb') as file:
+                pickle.dump(new_name_list, file)
+
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+        self.nlp = spacy.load('en_core_web_sm')
+        self.std_text = std_text
+        self.instructions = json.load(open(instructions, 'r'))
+        self.tasks = []
+        for task in self.instructions.keys():
+            for subtask in self.instructions[task].keys():
+                self.tasks.append(self.instructions[task][subtask])
+
+    def __len__(self):
+        return len(self.name_list) * len(self.tasks)
+
+    def __getitem__(self, item):
+        data_idx = item % len(self.name_list)
+        task_idx = item // len(self.name_list)
+
+        data = self.data_dict[self.name_list[data_idx]]
+        m_token_list, text_list = data['m_token_list'], data['text']
+
+        m_tokens = random.choice(m_token_list)
+        text_data = random.choice(text_list)
+        caption = text_data['caption']
+        if self.std_text:
+            doc = self.nlp(caption)
+            word_list = []
+            pos_list = []
+            for token in doc:
+                word = token.text
+                if not word.isalpha():
+                    continue
+                if (token.pos_ == 'NOUN'
+                        or token.pos_ == 'VERB') and (word != 'left'):
+                    word_list.append(token.lemma_)
+                else:
+                    word_list.append(word)
+                pos_list.append(token.pos_)
+                
+            caption = ' '.join(word_list)
+        
+        all_captions = [
+            ' '.join([token.split('/')[0] for token in text_dic['tokens']])
+            for text_dic in text_list
+        ]
+
+        coin = np.random.choice([False, False, True])
+
+        if coin:
+            # drop one token at the head or tail
+            coin2 = np.random.choice([True, False])
+            if coin2:
+                m_tokens = m_tokens[:-1]
+            else:
+                m_tokens = m_tokens[1:]
+
+        m_tokens_len = m_tokens.shape[0]
+
+        tasks = self.tasks[task_idx]
+
+        return caption, m_tokens, m_tokens_len, None, None, None, None, all_captions, tasks

mGPT/data/humanml/dataset_t2m_eval.py

@@ -0,0 +1,92 @@
+import random
+import numpy as np
+from .dataset_t2m import Text2MotionDataset
+
+
+class Text2MotionDatasetEval(Text2MotionDataset):
+
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        w_vectorizer,
+        max_motion_length=196,
+        min_motion_length=40,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+        super().__init__(data_root, split, mean, std, max_motion_length,
+                         min_motion_length, unit_length, fps, tmpFile, tiny,
+                         debug, **kwargs)
+
+        self.w_vectorizer = w_vectorizer
+
+
+    def __getitem__(self, item):
+        # Get text data
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data["motion"], data["length"], data["text"]
+
+        all_captions = [
+            ' '.join([token.split('/')[0] for token in text_dic['tokens']])
+            for text_dic in text_list
+        ]
+
+        if len(all_captions) > 3:
+            all_captions = all_captions[:3]
+        elif len(all_captions) == 2:
+            all_captions = all_captions + all_captions[0:1]
+        elif len(all_captions) == 1:
+            all_captions = all_captions * 3
+
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption, tokens = text_data["caption"], text_data["tokens"]
+
+        # Text
+        max_text_len = 20
+        if len(tokens) < max_text_len:
+            # pad with "unk"
+            tokens = ["sos/OTHER"] + tokens + ["eos/OTHER"]
+            sent_len = len(tokens)
+            tokens = tokens + ["unk/OTHER"] * (max_text_len + 2 - sent_len)
+        else:
+            # crop
+            tokens = tokens[:max_text_len]
+            tokens = ["sos/OTHER"] + tokens + ["eos/OTHER"]
+            sent_len = len(tokens)
+        pos_one_hots = []
+        word_embeddings = []
+        for token in tokens:
+            word_emb, pos_oh = self.w_vectorizer[token]
+            pos_one_hots.append(pos_oh[None, :])
+            word_embeddings.append(word_emb[None, :])
+        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+        word_embeddings = np.concatenate(word_embeddings, axis=0)
+        
+        # Random crop
+        if self.unit_length < 10:
+            coin2 = np.random.choice(["single", "single", "double"])
+        else:
+            coin2 = "single"
+
+        if coin2 == "double":
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == "single":
+            m_length = (m_length // self.unit_length) * self.unit_length
+
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx + m_length]
+        
+        # Z Normalization
+        motion = (motion - self.mean) / self.std
+
+        return caption, motion, m_length, word_embeddings, pos_one_hots, sent_len, "_".join(
+            tokens), all_captions

mGPT/data/humanml/dataset_t2m_m2t.py

@@ -0,0 +1,119 @@
+import random
+import numpy as np
+from torch.utils import data
+from .dataset_t2m import Text2MotionDataset
+import codecs as cs
+from os.path import join as pjoin
+
+
+class Text2MotionDatasetM2T(data.Dataset):
+
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length=196,
+        min_motion_length=40,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+        
+        self.max_motion_length = max_motion_length
+        self.min_motion_length = min_motion_length
+        self.unit_length = unit_length
+        
+        # Data mean and std
+        self.mean = mean
+        self.std = std
+        
+        # Data path
+        split_file = pjoin(data_root, split + '.txt')
+        motion_dir = pjoin(data_root, 'new_joint_vecs')
+        text_dir = pjoin(data_root, 'texts')
+
+        # Data id list
+        self.id_list = []
+        with cs.open(split_file, "r") as f:
+            for line in f.readlines():
+                self.id_list.append(line.strip())
+                
+        new_name_list = []
+        length_list = []
+        data_dict = {}
+        for name in self.id_list:
+            # try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if (len(motion)) <  self.min_motion_length or (len(motion) >= 200):
+                    continue
+                
+                        
+                text_data = []
+                flag = False
+                
+                with cs.open(pjoin(text_dir, name + '.txt')) as f:
+                    for line in f.readlines():
+                        text_dict = {}
+                        line_split = line.strip().split('#')
+                        caption = line_split[0]
+                        tokens = line_split[1].split(' ')
+                        f_tag = float(line_split[2])
+                        to_tag = float(line_split[3])
+                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                        text_dict['caption'] = caption
+                        text_dict['tokens'] = tokens
+                        if f_tag == 0.0 and to_tag == 0.0:
+                            flag = True
+                            text_data.append(text_dict)
+                        else:
+                            try:
+                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
+
+                                if (len(n_motion)) < min_motion_length or (len(n_motion) >= 200):
+                                    continue
+                                
+                                new_name = "%s_%f_%f"%(name, f_tag, to_tag)
+                                data_dict[new_name] = {'motion': n_motion,
+                                                       'length': len(n_motion),
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                            except:
+                                print(line_split)
+                                print(line_split[2], line_split[3], f_tag, to_tag, name)
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                    'length': len(motion),
+                                    'name': name,
+                                    'text': text_data}
+                    
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            # except:
+            #     # Some motion may not exist in KIT dataset
+            #     pass
+
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+        self.nfeats = motion.shape[-1]
+    
+    
+    def __len__(self):
+        return len(self.data_dict)  
+        
+    def __getitem__(self, item):
+        name = self.name_list[item]
+        data = self.data_dict[name]
+        motion, m_length = data['motion'], data['length']
+
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return name, motion, m_length, True, True, True, True, True, True

mGPT/data/humanml/dataset_t2m_token.py

@@ -0,0 +1,86 @@
+import random
+import numpy as np
+from torch.utils import data
+from .dataset_t2m import Text2MotionDataset
+import codecs as cs
+from os.path import join as pjoin
+
+
+class Text2MotionDatasetToken(data.Dataset):
+
+    def __init__(
+        self,
+        data_root,
+        split,
+        mean,
+        std,
+        max_motion_length=196,
+        min_motion_length=40,
+        unit_length=4,
+        fps=20,
+        tmpFile=True,
+        tiny=False,
+        debug=False,
+        **kwargs,
+    ):
+        
+        self.max_motion_length = max_motion_length
+        self.min_motion_length = min_motion_length
+        self.unit_length = unit_length
+        
+        # Data mean and std
+        self.mean = mean
+        self.std = std
+        
+        # Data path
+        split_file = pjoin(data_root, split + '.txt')
+        motion_dir = pjoin(data_root, 'new_joint_vecs')
+        text_dir = pjoin(data_root, 'texts')
+
+        # Data id list
+        self.id_list = []
+        with cs.open(split_file, "r") as f:
+            for line in f.readlines():
+                self.id_list.append(line.strip())
+                
+        new_name_list = []
+        length_list = []
+        data_dict = {}
+        for name in self.id_list:
+            try:
+                motion = np.load(pjoin(motion_dir, name + '.npy'))
+                if (len(motion)) <  self.min_motion_length or (len(motion) >= 200):
+                    continue
+
+                data_dict[name] = {'motion': motion,
+                                'length': len(motion),
+                                'name': name}
+                new_name_list.append(name)
+                length_list.append(len(motion))
+            except:
+                # Some motion may not exist in KIT dataset
+                pass
+
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+        self.nfeats = motion.shape[-1]
+    
+    
+    def __len__(self):
+        return len(self.data_dict)  
+        
+    def __getitem__(self, item):
+        name = self.name_list[item]
+        data = self.data_dict[name]
+        motion, m_length = data['motion'], data['length']
+
+        m_length = (m_length // self.unit_length) * self.unit_length
+
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx+m_length]
+
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return name, motion, m_length, True, True, True, True, True, True

mGPT/data/humanml/scripts/motion_process.py

@@ -0,0 +1,529 @@
+from os.path import join as pjoin
+
+from ..common.skeleton import Skeleton
+import numpy as np
+import os
+from ..common.quaternion import *
+from ..utils.paramUtil import *
+
+import torch
+from tqdm import tqdm
+
+# positions (batch, joint_num, 3)
+def uniform_skeleton(positions, target_offset):
+    src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))
+    src_offset = src_offset.numpy()
+    tgt_offset = target_offset.numpy()
+    # print(src_offset)
+    # print(tgt_offset)
+    '''Calculate Scale Ratio as the ratio of legs'''
+    src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max()
+    tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max()
+
+    scale_rt = tgt_leg_len / src_leg_len
+    # print(scale_rt)
+    src_root_pos = positions[:, 0]
+    tgt_root_pos = src_root_pos * scale_rt
+
+    '''Inverse Kinematics'''
+    quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx)
+    # print(quat_params.shape)
+
+    '''Forward Kinematics'''
+    src_skel.set_offset(target_offset)
+    new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)
+    return new_joints
+
+
+def extract_features(positions, feet_thre, n_raw_offsets, kinematic_chain, face_joint_indx, fid_r, fid_l):
+    global_positions = positions.copy()
+    """ Get Foot Contacts """
+
+    def foot_detect(positions, thres):
+        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])
+
+        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
+        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
+        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
+        #     feet_l_h = positions[:-1,fid_l,1]
+        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float64)
+        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float64)
+
+        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
+        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
+        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
+        #     feet_r_h = positions[:-1,fid_r,1]
+        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float64)
+        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float64)
+        return feet_l, feet_r
+
+    #
+    feet_l, feet_r = foot_detect(positions, feet_thre)
+    # feet_l, feet_r = foot_detect(positions, 0.002)
+
+    '''Quaternion and Cartesian representation'''
+    r_rot = None
+
+    def get_rifke(positions):
+        '''Local pose'''
+        positions[..., 0] -= positions[:, 0:1, 0]
+        positions[..., 2] -= positions[:, 0:1, 2]
+        '''All pose face Z+'''
+        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
+        return positions
+
+    def get_quaternion(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)
+
+        '''Fix Quaternion Discontinuity'''
+        quat_params = qfix(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        quat_params[1:, 0] = r_velocity
+        # (seq_len, joints_num, 4)
+        return quat_params, r_velocity, velocity, r_rot
+
+    def get_cont6d_params(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)
+
+        '''Quaternion to continuous 6D'''
+        cont_6d_params = quaternion_to_cont6d_np(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        # (seq_len, joints_num, 4)
+        return cont_6d_params, r_velocity, velocity, r_rot
+
+    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)
+    positions = get_rifke(positions)
+
+    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)
+    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')
+    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')
+    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    '''Root height'''
+    root_y = positions[:, 0, 1:2]
+
+    '''Root rotation and linear velocity'''
+    # (seq_len-1, 1) rotation velocity along y-axis
+    # (seq_len-1, 2) linear velovity on xz plane
+    r_velocity = np.arcsin(r_velocity[:, 2:3])
+    l_velocity = velocity[:, [0, 2]]
+    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)
+    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)
+
+    '''Get Joint Rotation Representation'''
+    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints
+    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)
+
+    '''Get Joint Rotation Invariant Position Represention'''
+    # (seq_len, (joints_num-1)*3) local joint position
+    ric_data = positions[:, 1:].reshape(len(positions), -1)
+
+    '''Get Joint Velocity Representation'''
+    # (seq_len-1, joints_num*3)
+    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
+                        global_positions[1:] - global_positions[:-1])
+    local_vel = local_vel.reshape(len(local_vel), -1)
+
+    data = root_data
+    data = np.concatenate([data, ric_data[:-1]], axis=-1)
+    data = np.concatenate([data, rot_data[:-1]], axis=-1)
+    #     print(dataset.shape, local_vel.shape)
+    data = np.concatenate([data, local_vel], axis=-1)
+    data = np.concatenate([data, feet_l, feet_r], axis=-1)
+
+    return data
+
+
+def process_file(positions, feet_thre):
+    # (seq_len, joints_num, 3)
+    #     '''Down Sample'''
+    #     positions = positions[::ds_num]
+
+    '''Uniform Skeleton'''
+    positions = uniform_skeleton(positions, tgt_offsets)
+
+    '''Put on Floor'''
+    floor_height = positions.min(axis=0).min(axis=0)[1]
+    positions[:, :, 1] -= floor_height
+    #     print(floor_height)
+
+    #     plot_3d_motion("./positions_1.mp4", kinematic_chain, positions, 'title', fps=20)
+
+    '''XZ at origin'''
+    root_pos_init = positions[0]
+    root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])
+    positions = positions - root_pose_init_xz
+
+    # '''Move the first pose to origin '''
+    # root_pos_init = positions[0]
+    # positions = positions - root_pos_init[0]
+
+    '''All initially face Z+'''
+    r_hip, l_hip, sdr_r, sdr_l = face_joint_indx
+    across1 = root_pos_init[r_hip] - root_pos_init[l_hip]
+    across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]
+    across = across1 + across2
+    across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]
+
+    # forward (3,), rotate around y-axis
+    forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+    # forward (3,)
+    forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]
+
+    #     print(forward_init)
+
+    target = np.array([[0, 0, 1]])
+    root_quat_init = qbetween_np(forward_init, target)
+    root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init
+
+    positions_b = positions.copy()
+
+    positions = qrot_np(root_quat_init, positions)
+
+    #     plot_3d_motion("./positions_2.mp4", kinematic_chain, positions, 'title', fps=20)
+
+    '''New ground truth positions'''
+    global_positions = positions.copy()
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(positions[:, 0, 0], positions[:, 0, 2], marker='o', color='r')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    """ Get Foot Contacts """
+
+    def foot_detect(positions, thres):
+        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])
+
+        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
+        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
+        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
+        #     feet_l_h = positions[:-1,fid_l,1]
+        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float64)
+        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float64)
+
+        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
+        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
+        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
+        #     feet_r_h = positions[:-1,fid_r,1]
+        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float64)
+        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float64)
+        return feet_l, feet_r
+    #
+    feet_l, feet_r = foot_detect(positions, feet_thre)
+    # feet_l, feet_r = foot_detect(positions, 0.002)
+
+    '''Quaternion and Cartesian representation'''
+    r_rot = None
+
+    def get_rifke(positions):
+        '''Local pose'''
+        positions[..., 0] -= positions[:, 0:1, 0]
+        positions[..., 2] -= positions[:, 0:1, 2]
+        '''All pose face Z+'''
+        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
+        return positions
+
+    def get_quaternion(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)
+
+        '''Fix Quaternion Discontinuity'''
+        quat_params = qfix(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        quat_params[1:, 0] = r_velocity
+        # (seq_len, joints_num, 4)
+        return quat_params, r_velocity, velocity, r_rot
+
+    def get_cont6d_params(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)
+
+        '''Quaternion to continuous 6D'''
+        cont_6d_params = quaternion_to_cont6d_np(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        # (seq_len, joints_num, 4)
+        return cont_6d_params, r_velocity, velocity, r_rot
+
+    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)
+    positions = get_rifke(positions)
+
+    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)
+    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')
+    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')
+    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    '''Root height'''
+    root_y = positions[:, 0, 1:2]
+
+    '''Root rotation and linear velocity'''
+    # (seq_len-1, 1) rotation velocity along y-axis
+    # (seq_len-1, 2) linear velovity on xz plane
+    r_velocity = np.arcsin(r_velocity[:, 2:3])
+    l_velocity = velocity[:, [0, 2]]
+    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)
+    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)
+
+    '''Get Joint Rotation Representation'''
+    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints
+    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)
+
+    '''Get Joint Rotation Invariant Position Represention'''
+    # (seq_len, (joints_num-1)*3) local joint position
+    ric_data = positions[:, 1:].reshape(len(positions), -1)
+
+    '''Get Joint Velocity Representation'''
+    # (seq_len-1, joints_num*3)
+    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
+                        global_positions[1:] - global_positions[:-1])
+    local_vel = local_vel.reshape(len(local_vel), -1)
+
+    data = root_data
+    data = np.concatenate([data, ric_data[:-1]], axis=-1)
+    data = np.concatenate([data, rot_data[:-1]], axis=-1)
+    #     print(dataset.shape, local_vel.shape)
+    data = np.concatenate([data, local_vel], axis=-1)
+    data = np.concatenate([data, feet_l, feet_r], axis=-1)
+
+    return data, global_positions, positions, l_velocity
+
+
+# Recover global angle and positions for rotation dataset
+# root_rot_velocity (B, seq_len, 1)
+# root_linear_velocity (B, seq_len, 2)
+# root_y (B, seq_len, 1)
+# ric_data (B, seq_len, (joint_num - 1)*3)
+# rot_data (B, seq_len, (joint_num - 1)*6)
+# local_velocity (B, seq_len, joint_num*3)
+# foot contact (B, seq_len, 4)
+def recover_root_rot_pos(data):
+    rot_vel = data[..., 0]
+    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
+    '''Get Y-axis rotation from rotation velocity'''
+    r_rot_ang[..., 1:] = rot_vel[..., :-1]
+    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
+
+    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
+    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
+    r_rot_quat[..., 2] = torch.sin(r_rot_ang)
+
+    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
+    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
+    '''Add Y-axis rotation to root position'''
+    r_pos = qrot(qinv(r_rot_quat), r_pos)
+
+    r_pos = torch.cumsum(r_pos, dim=-2)
+
+    r_pos[..., 1] = data[..., 3]
+    return r_rot_quat, r_pos
+
+
+def recover_from_rot(data, joints_num, skeleton):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+
+    r_rot_cont6d = quaternion_to_cont6d(r_rot_quat)
+
+    start_indx = 1 + 2 + 1 + (joints_num - 1) * 3
+    end_indx = start_indx + (joints_num - 1) * 6
+    cont6d_params = data[..., start_indx:end_indx]
+    #     print(r_rot_cont6d.shape, cont6d_params.shape, r_pos.shape)
+    cont6d_params = torch.cat([r_rot_cont6d, cont6d_params], dim=-1)
+    cont6d_params = cont6d_params.view(-1, joints_num, 6)
+
+    positions = skeleton.forward_kinematics_cont6d(cont6d_params, r_pos)
+
+    return positions
+
+def recover_rot(data):
+    # dataset [bs, seqlen, 263/251] HumanML/KIT
+    joints_num = 22 if data.shape[-1] == 263 else 21
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+    r_pos_pad = torch.cat([r_pos, torch.zeros_like(r_pos)], dim=-1).unsqueeze(-2)
+    r_rot_cont6d = quaternion_to_cont6d(r_rot_quat)
+    start_indx = 1 + 2 + 1 + (joints_num - 1) * 3
+    end_indx = start_indx + (joints_num - 1) * 6
+    cont6d_params = data[..., start_indx:end_indx]
+    cont6d_params = torch.cat([r_rot_cont6d, cont6d_params], dim=-1)
+    cont6d_params = cont6d_params.view(-1, joints_num, 6)
+    cont6d_params = torch.cat([cont6d_params, r_pos_pad], dim=-2)
+    return cont6d_params
+
+
+def recover_from_ric(data, joints_num):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+    positions = data[..., 4:(joints_num - 1) * 3 + 4]
+    positions = positions.view(positions.shape[:-1] + (-1, 3))
+
+    '''Add Y-axis rotation to local joints'''
+    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)
+
+    '''Add root XZ to joints'''
+    positions[..., 0] += r_pos[..., 0:1]
+    positions[..., 2] += r_pos[..., 2:3]
+
+    '''Concate root and joints'''
+    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
+
+    return positions
+'''
+For Text2Motion Dataset
+'''
+'''
+if __name__ == "__main__":
+    example_id = "000021"
+    # Lower legs
+    l_idx1, l_idx2 = 5, 8
+    # Right/Left foot
+    fid_r, fid_l = [8, 11], [7, 10]
+    # Face direction, r_hip, l_hip, sdr_r, sdr_l
+    face_joint_indx = [2, 1, 17, 16]
+    # l_hip, r_hip
+    r_hip, l_hip = 2, 1
+    joints_num = 22
+    # ds_num = 8
+    data_dir = '../dataset/pose_data_raw/joints/'
+    save_dir1 = '../dataset/pose_data_raw/new_joints/'
+    save_dir2 = '../dataset/pose_data_raw/new_joint_vecs/'
+
+    n_raw_offsets = torch.from_numpy(t2m_raw_offsets)
+    kinematic_chain = t2m_kinematic_chain
+
+    # Get offsets of target skeleton
+    example_data = np.load(os.path.join(data_dir, example_id + '.npy'))
+    example_data = example_data.reshape(len(example_data), -1, 3)
+    example_data = torch.from_numpy(example_data)
+    tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    # (joints_num, 3)
+    tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])
+    # print(tgt_offsets)
+
+    source_list = os.listdir(data_dir)
+    frame_num = 0
+    for source_file in tqdm(source_list):
+        source_data = np.load(os.path.join(data_dir, source_file))[:, :joints_num]
+        try:
+            dataset, ground_positions, positions, l_velocity = process_file(source_data, 0.002)
+            rec_ric_data = recover_from_ric(torch.from_numpy(dataset).unsqueeze(0).float(), joints_num)
+            np.save(pjoin(save_dir1, source_file), rec_ric_data.squeeze().numpy())
+            np.save(pjoin(save_dir2, source_file), dataset)
+            frame_num += dataset.shape[0]
+        except Exception as e:
+            print(source_file)
+            print(e)
+
+    print('Total clips: %d, Frames: %d, Duration: %fm' %
+          (len(source_list), frame_num, frame_num / 20 / 60))
+'''
+
+if __name__ == "__main__":
+    example_id = "03950_gt"
+    # Lower legs
+    l_idx1, l_idx2 = 17, 18
+    # Right/Left foot
+    fid_r, fid_l = [14, 15], [19, 20]
+    # Face direction, r_hip, l_hip, sdr_r, sdr_l
+    face_joint_indx = [11, 16, 5, 8]
+    # l_hip, r_hip
+    r_hip, l_hip = 11, 16
+    joints_num = 21
+    # ds_num = 8
+    data_dir = '../dataset/kit_mocap_dataset/joints/'
+    save_dir1 = '../dataset/kit_mocap_dataset/new_joints/'
+    save_dir2 = '../dataset/kit_mocap_dataset/new_joint_vecs/'
+
+    n_raw_offsets = torch.from_numpy(kit_raw_offsets)
+    kinematic_chain = kit_kinematic_chain
+
+    '''Get offsets of target skeleton'''
+    example_data = np.load(os.path.join(data_dir, example_id + '.npy'))
+    example_data = example_data.reshape(len(example_data), -1, 3)
+    example_data = torch.from_numpy(example_data)
+    tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    # (joints_num, 3)
+    tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])
+    # print(tgt_offsets)
+
+    source_list = os.listdir(data_dir)
+    frame_num = 0
+    '''Read source dataset'''
+    for source_file in tqdm(source_list):
+        source_data = np.load(os.path.join(data_dir, source_file))[:, :joints_num]
+        try:
+            name = ''.join(source_file[:-7].split('_')) + '.npy'
+            data, ground_positions, positions, l_velocity = process_file(source_data, 0.05)
+            rec_ric_data = recover_from_ric(torch.from_numpy(data).unsqueeze(0).float(), joints_num)
+            if np.isnan(rec_ric_data.numpy()).any():
+                print(source_file)
+                continue
+            np.save(pjoin(save_dir1, name), rec_ric_data.squeeze().numpy())
+            np.save(pjoin(save_dir2, name), data)
+            frame_num += data.shape[0]
+        except Exception as e:
+            print(source_file)
+            print(e)
+
+    print('Total clips: %d, Frames: %d, Duration: %fm' %
+          (len(source_list), frame_num, frame_num / 12.5 / 60))

mGPT/data/humanml/utils/paramUtil.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+# Define a kinematic tree for the skeletal struture
+kit_kinematic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]]
+
+kit_raw_offsets = np.array(
+    [
+        [0, 0, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1]
+    ]
+)
+
+t2m_raw_offsets = np.array([[0,0,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,0,1],
+                           [0,0,1],
+                           [0,1,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,0,1],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0]])
+
+t2m_kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]]
+t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]]
+t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]]
+
+
+kit_tgt_skel_id = '03950'
+
+t2m_tgt_skel_id = '000021'
+

mGPT/data/humanml/utils/word_vectorizer.py

@@ -0,0 +1,79 @@
+import numpy as np
+import pickle
+from os.path import join as pjoin
+
+POS_enumerator = {
+    'VERB': 0,
+    'NOUN': 1,
+    'DET': 2,
+    'ADP': 3,
+    'NUM': 4,
+    'AUX': 5,
+    'PRON': 6,
+    'ADJ': 7,
+    'ADV': 8,
+    'Loc_VIP': 9,
+    'Body_VIP': 10,
+    'Obj_VIP': 11,
+    'Act_VIP': 12,
+    'Desc_VIP': 13,
+    'OTHER': 14,
+}
+
+Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward',
+            'up', 'down', 'straight', 'curve')
+
+Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh')
+
+Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball')
+
+Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn',
+            'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll',
+            'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb')
+
+Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily', 'angrily', 'sadly')
+
+VIP_dict = {
+    'Loc_VIP': Loc_list,
+    'Body_VIP': Body_list,
+    'Obj_VIP': Obj_List,
+    'Act_VIP': Act_list,
+    'Desc_VIP': Desc_list,
+}
+
+
+class WordVectorizer(object):
+    def __init__(self, meta_root, prefix):
+        vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix))
+        words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb'))
+        word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb'))
+        self.word2vec = {w: vectors[word2idx[w]] for w in words}
+
+    def _get_pos_ohot(self, pos):
+        pos_vec = np.zeros(len(POS_enumerator))
+        if pos in POS_enumerator:
+            pos_vec[POS_enumerator[pos]] = 1
+        else:
+            pos_vec[POS_enumerator['OTHER']] = 1
+        return pos_vec
+
+    def __len__(self):
+        return len(self.word2vec)
+
+    def __getitem__(self, item):
+        word, pos = item.split('/')
+        if word in self.word2vec:
+            word_vec = self.word2vec[word]
+            vip_pos = None
+            for key, values in VIP_dict.items():
+                if word in values:
+                    vip_pos = key
+                    break
+            if vip_pos is not None:
+                pos_vec = self._get_pos_ohot(vip_pos)
+            else:
+                pos_vec = self._get_pos_ohot(pos)
+        else:
+            word_vec = self.word2vec['unk']
+            pos_vec = self._get_pos_ohot('OTHER')
+        return word_vec, pos_vec

mGPT/data/tools/__init__.py

@@ -0,0 +1,2 @@
+from .tensors import lengths_to_mask
+from .collate import collate_text_and_length, collate_pairs_and_text, collate_datastruct_and_text, collate_tensor_with_padding

mGPT/data/tools/collate.py

@@ -0,0 +1,99 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import List, Dict
+from torch import Tensor
+
+
+def collate_tensor_with_padding(batch: List[Tensor]) -> Tensor:
+    dims = batch[0].dim()
+    max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
+    size = (len(batch),) + tuple(max_size)
+    canvas = batch[0].new_zeros(size=size)
+    for i, b in enumerate(batch):
+        sub_tensor = canvas[i]
+        for d in range(dims):
+            sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
+        sub_tensor.add_(b)
+    return canvas
+
+
+def collate_datastruct_and_text(lst_elements: List) -> Dict:
+    collate_datastruct = lst_elements[0]["datastruct"].transforms.collate
+
+    batch = {
+        # Collate with padding for the datastruct
+        "datastruct": collate_datastruct([x["datastruct"] for x in lst_elements]),
+        # Collate normally for the length
+        "length": [x["length"] for x in lst_elements],
+        # Collate the text
+        "text": [x["text"] for x in lst_elements]}
+
+    # add keyid for example
+    otherkeys = [x for x in lst_elements[0].keys() if x not in batch]
+    for key in otherkeys:
+        batch[key] = [x[key] for x in lst_elements]
+
+    return batch
+
+def collate_length_and_text(lst_elements: List) -> Dict:
+
+    batch = {
+            "length_0": [x["length_0"] for x in lst_elements], 
+            "length_1": [x["length_1"] for x in lst_elements], 
+            "length_transition": [x["length_transition"] for x in lst_elements], 
+            "length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
+            "text_0": [x["text_0"] for x in lst_elements],
+            "text_1": [x["text_1"] for x in lst_elements]
+    }
+
+    return batch
+
+def collate_pairs_and_text(lst_elements: List, ) -> Dict:
+    if 'features_0' not in lst_elements[0]: # test set
+        collate_datastruct = lst_elements[0]["datastruct"].transforms.collate
+        batch = {"datastruct": collate_datastruct([x["datastruct"] for x in lst_elements]),
+                "length_0": [x["length_0"] for x in lst_elements], 
+                "length_1": [x["length_1"] for x in lst_elements], 
+                "length_transition": [x["length_transition"] for x in lst_elements], 
+                "length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
+                "text_0": [x["text_0"] for x in lst_elements],
+                "text_1": [x["text_1"] for x in lst_elements]
+        }
+
+    else:
+        batch = {"motion_feats_0": collate_tensor_with_padding([el["features_0"] for el in lst_elements]),
+                "motion_feats_1": collate_tensor_with_padding([el["features_1"] for el in lst_elements]),
+                "motion_feats_1_with_transition": collate_tensor_with_padding([el["features_1_with_transition"] for el in lst_elements]),
+                "length_0": [x["length_0"] for x in lst_elements], 
+                "length_1": [x["length_1"] for x in lst_elements], 
+                "length_transition": [x["length_transition"] for x in lst_elements], 
+                "length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
+                "text_0": [x["text_0"] for x in lst_elements],
+                "text_1": [x["text_1"] for x in lst_elements]
+                }
+    return batch
+
+
+def collate_text_and_length(lst_elements: Dict) -> Dict:
+    batch = {"length": [x["length"] for x in lst_elements],
+             "text": [x["text"] for x in lst_elements]}
+
+    # add keyid for example
+    otherkeys = [x for x in lst_elements[0].keys() if x not in batch and x != "datastruct"]
+    for key in otherkeys:
+        batch[key] = [x[key] for x in lst_elements]
+    return batch

mGPT/data/tools/easyconvert.py

@@ -0,0 +1,72 @@
+from .geometry import *
+
+def nfeats_of(rottype):
+    if rottype in ["rotvec", "axisangle"]:
+        return 3
+    elif rottype in ["rotquat", "quaternion"]:
+        return 4
+    elif rottype in ["rot6d", "6drot", "rotation6d"]:
+        return 6
+    elif rottype in ["rotmat"]:
+        return 9
+    else:
+        return TypeError("This rotation type doesn't have features.")
+
+
+def axis_angle_to(newtype, rotations):
+    if newtype in ["matrix"]:
+        rotations = axis_angle_to_matrix(rotations)
+        return rotations
+    elif newtype in ["rotmat"]:
+        rotations = axis_angle_to_matrix(rotations)
+        rotations = matrix_to("rotmat", rotations)
+        return rotations
+    elif newtype in ["rot6d", "6drot", "rotation6d"]:
+        rotations = axis_angle_to_matrix(rotations)
+        rotations = matrix_to("rot6d", rotations)
+        return rotations
+    elif newtype in ["rotquat", "quaternion"]:
+        rotations = axis_angle_to_quaternion(rotations)
+        return rotations
+    elif newtype in ["rotvec", "axisangle"]:
+        return rotations
+    else:
+        raise NotImplementedError
+
+
+def matrix_to(newtype, rotations):
+    if newtype in ["matrix"]:
+        return rotations
+    if newtype in ["rotmat"]:
+        rotations = rotations.reshape((*rotations.shape[:-2], 9))
+        return rotations
+    elif newtype in ["rot6d", "6drot", "rotation6d"]:
+        rotations = matrix_to_rotation_6d(rotations)
+        return rotations
+    elif newtype in ["rotquat", "quaternion"]:
+        rotations = matrix_to_quaternion(rotations)
+        return rotations
+    elif newtype in ["rotvec", "axisangle"]:
+        rotations = matrix_to_axis_angle(rotations)
+        return rotations
+    else:
+        raise NotImplementedError
+
+
+def to_matrix(oldtype, rotations):
+    if oldtype in ["matrix"]:
+        return rotations
+    if oldtype in ["rotmat"]:
+        rotations = rotations.reshape((*rotations.shape[:-2], 3, 3))
+        return rotations
+    elif oldtype in ["rot6d", "6drot", "rotation6d"]:
+        rotations = rotation_6d_to_matrix(rotations)
+        return rotations
+    elif oldtype in ["rotquat", "quaternion"]:
+        rotations = quaternion_to_matrix(rotations)
+        return rotations
+    elif oldtype in ["rotvec", "axisangle"]:
+        rotations = axis_angle_to_matrix(rotations)
+        return rotations
+    else:
+        raise NotImplementedError

mGPT/data/tools/geometry.py

@@ -0,0 +1,566 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+# Check PYTORCH3D_LICENCE before use
+
+import functools
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+"""
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+
+# Added
+def matrix_of_angles(cos, sin, inv=False, dim=2):
+    assert dim in [2, 3]
+    sin = -sin if inv else sin
+    if dim == 2:
+        row1 = torch.stack((cos, -sin), axis=-1)
+        row2 = torch.stack((sin, cos), axis=-1)
+        return torch.stack((row1, row2), axis=-2)
+    elif dim == 3:
+        row1 = torch.stack((cos, -sin, 0*cos), axis=-1)
+        row2 = torch.stack((sin, cos, 0*cos), axis=-1)
+        row3 = torch.stack((0*sin, 0*cos, 1+0*cos), axis=-1)
+        return torch.stack((row1, row2, row3),axis=-2)
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def _copysign(a, b):
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x):
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    return ret
+
+
+def matrix_to_quaternion(matrix):
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    m00 = matrix[..., 0, 0]
+    m11 = matrix[..., 1, 1]
+    m22 = matrix[..., 2, 2]
+    o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22)
+    x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22)
+    y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22)
+    z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22)
+    o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2])
+    o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0])
+    o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1])
+    return torch.stack((o0, o1, o2, o3), -1)
+
+
+def _axis_angle_rotation(axis: str, angle):
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    if axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    if axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles, convention: str):
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1))
+    return functools.reduce(torch.matmul, matrices)
+
+
+def _angle_from_tan(
+    axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
+):
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in data as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str):
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+
+
+def matrix_to_euler_angles(matrix, convention: str):
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(
+            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
+        )
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(
+            convention[0], convention[1], matrix[..., i2], False, tait_bryan
+        ),
+        central_angle,
+        _angle_from_tan(
+            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
+        ),
+    )
+    return torch.stack(o, -1)
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random rotations as 3x3 rotation matrices.
+
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(
+        n, dtype=dtype, device=device, requires_grad=requires_grad
+    )
+    return quaternion_to_matrix(quaternions)
+
+
+def random_rotation(
+    dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate a single random 3x3 rotation matrix.
+
+    Args:
+        dtype: Type to return
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set
+
+    Returns:
+        Rotation matrix as tensor of shape (3, 3).
+    """
+    return random_rotations(1, dtype, device, requires_grad)[0]
+
+
+def standardize_quaternion(quaternions):
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a, b):
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a, b):
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion):
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    return quaternion * quaternion.new_tensor([1, -1, -1, -1])
+
+
+def quaternion_apply(quaternion, point):
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, f{point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1,))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle):
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix):
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle):
+    """
+    Convert rotations given as axis/angle to quaternions.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = 0.5 * angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+
+def quaternion_to_axis_angle(quaternions):
+    """
+    Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalisation per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)

mGPT/data/tools/tensors.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import List, Dict
+import torch
+from torch import Tensor
+
+
+def lengths_to_mask(lengths: List[int], device: torch.device) -> Tensor:
+    lengths = torch.tensor(lengths, device=device)
+    max_len = max(lengths)
+    mask = torch.arange(max_len, device=device).expand(len(lengths), max_len) < lengths.unsqueeze(1)
+    return mask

mGPT/data/transforms/__init__.py

@@ -0,0 +1,15 @@
+from .base import Transform
+from .smpl import SMPLTransform
+from .xyz import XYZTransform
+
+# rots2rfeats
+from .rots2rfeats import Rots2Rfeats
+from .rots2rfeats import Globalvelandy
+
+# rots2joints
+from .rots2joints import Rots2Joints
+from .rots2joints import SMPLH, SMPLX
+
+# joints2jfeats
+from .joints2jfeats import Joints2Jfeats
+from .joints2jfeats import Rifke

mGPT/data/transforms/base.py

@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from dataclasses import dataclass, fields
+
+
+class Transform:
+
+    def collate(self, lst_datastruct):
+        from ..tools import collate_tensor_with_padding
+        example = lst_datastruct[0]
+
+        def collate_or_none(key):
+            if example[key] is None:
+                return None
+            key_lst = [x[key] for x in lst_datastruct]
+            return collate_tensor_with_padding(key_lst)
+
+        kwargs = {key: collate_or_none(key) for key in example.datakeys}
+
+        return self.Datastruct(**kwargs)
+
+
+# Inspired from SMPLX library
+# need to define "datakeys" and transforms
+@dataclass
+class Datastruct:
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        self.__dict__[key] = value
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+    def __iter__(self):
+        return self.keys()
+
+    def keys(self):
+        keys = [t.name for t in fields(self)]
+        return iter(keys)
+
+    def values(self):
+        values = [getattr(self, t.name) for t in fields(self)]
+        return iter(values)
+
+    def items(self):
+        data = [(t.name, getattr(self, t.name)) for t in fields(self)]
+        return iter(data)
+
+    def to(self, *args, **kwargs):
+        for key in self.datakeys:
+            if self[key] is not None:
+                self[key] = self[key].to(*args, **kwargs)
+        return self
+
+    @property
+    def device(self):
+        return self[self.datakeys[0]].device
+
+    def detach(self):
+
+        def detach_or_none(tensor):
+            if tensor is not None:
+                return tensor.detach()
+            return None
+
+        kwargs = {key: detach_or_none(self[key]) for key in self.datakeys}
+        return self.transforms.Datastruct(**kwargs)

mGPT/data/transforms/identity.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+from torch import Tensor
+
+from .base import Datastruct, dataclass, Transform
+
+
+class IdentityTransform(Transform):
+    def __init__(self, **kwargs):
+        return
+
+    def Datastruct(self, **kwargs):
+        return IdentityDatastruct(**kwargs)
+
+    def __repr__(self):
+        return "IdentityTransform()"
+
+
+@dataclass
+class IdentityDatastruct(Datastruct):
+    transforms: IdentityTransform
+
+    features: Optional[Tensor] = None
+
+    def __post_init__(self):
+        self.datakeys = ["features"]
+
+    def __len__(self):
+        return len(self.rfeats)

mGPT/data/transforms/joints2jfeats/__init__.py

@@ -0,0 +1,2 @@
+from .base import Joints2Jfeats
+from .rifke import Rifke

mGPT/data/transforms/joints2jfeats/base.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+from pathlib import Path
+import os
+
+
+class Joints2Jfeats(nn.Module):
+
+    def __init__(self,
+                 path: Optional[str] = None,
+                 normalization: bool = False,
+                 eps: float = 1e-12,
+                 **kwargs) -> None:
+        if normalization and path is None:
+            raise TypeError(
+                "You should provide a path if normalization is on.")
+
+        super().__init__()
+        self.normalization = normalization
+        self.eps = eps
+        # workaround for cluster local/sync
+        if path is not None:
+            # rel_p = path.split('/')
+            # rel_p = rel_p[rel_p.index('deps'):]
+            # rel_p = '/'.join(rel_p)
+            pass
+        if normalization:
+            mean_path = Path(path) / "jfeats_mean.pt"
+            std_path = Path(path) / "jfeats_std.pt"
+            self.register_buffer('mean', torch.load(mean_path))
+            self.register_buffer('std', torch.load(std_path))
+
+    def normalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = (features - self.mean) / (self.std + self.eps)
+        return features
+
+    def unnormalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = features * self.std + self.mean
+        return features

mGPT/data/transforms/joints2jfeats/rifke.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from .tools import get_forward_direction, get_floor, gaussian_filter1d  # noqa
+from mGPT.utils.geometry_tools import matrix_of_angles
+from .base import Joints2Jfeats
+
+
+class Rifke(Joints2Jfeats):
+
+    def __init__(self,
+                 jointstype: str = "mmm",
+                 path: Optional[str] = None,
+                 normalization: bool = False,
+                 forward_filter: bool = False,
+                 **kwargs) -> None:
+        #
+        # if jointstype != "mmm":
+        #     print("This function assume that the root is the first index")
+        #     raise NotImplementedError("This jointstype is not implemented.")
+
+        super().__init__(path=path, normalization=normalization)
+        self.jointstype = jointstype
+        self.forward_filter = forward_filter
+
+    def forward(self, joints: Tensor) -> Tensor:
+        # Joints to rotation invariant poses (Holden et. al.)
+        # Similar function than fke2rifke in Language2Pose repository
+        # Adapted to pytorch
+        # Put the origin center of the root joint instead of the ground projection
+        poses = joints.clone()
+        poses[..., 1] -= get_floor(poses, jointstype=self.jointstype)
+
+        translation = poses[..., 0, :].clone()
+        # Let the root have the Y translation --> gravity axis
+        root_y = translation[..., 1]
+
+        # Trajectory => Translation without gravity axis (Y)
+        trajectory = translation[..., [0, 2]]
+
+        # Delete the root joints of the poses
+        poses = poses[..., 1:, :]
+
+        # Remove the trajectory of the poses
+        poses[..., [0, 2]] -= trajectory[..., None, :]
+
+        # Compute the trajectory
+        vel_trajectory = torch.diff(trajectory, dim=-2)
+        # 0 for the first one => keep the dimentionality
+        vel_trajectory = torch.cat(
+            (0 * vel_trajectory[..., [0], :], vel_trajectory), dim=-2)
+
+        # Compute the forward direction
+        forward = get_forward_direction(poses, jointstype=self.jointstype)
+        if self.forward_filter:
+            # Smoothing to remove high frequencies
+            forward = gaussian_filter1d(forward, 2)
+            # normalize again to get real directions
+            forward = torch.nn.functional.normalize(forward, dim=-1)
+        # changed this also for New pytorch
+        angles = torch.atan2(*(forward.transpose(0, -1))).transpose(0, -1)
+        vel_angles = torch.diff(angles, dim=-1)
+        # 0 for the first one => keep the dimentionality
+        vel_angles = torch.cat((0 * vel_angles[..., [0]], vel_angles), dim=-1)
+
+        # Construct the inverse rotation matrix
+        sin, cos = forward[..., 0], forward[..., 1]
+        rotations_inv = matrix_of_angles(cos, sin, inv=True)
+
+        # Rotate the poses
+        poses_local = torch.einsum("...lj,...jk->...lk", poses[..., [0, 2]],
+                                   rotations_inv)
+        poses_local = torch.stack(
+            (poses_local[..., 0], poses[..., 1], poses_local[..., 1]), axis=-1)
+
+        # stack the xyz joints into feature vectors
+        poses_features = rearrange(poses_local,
+                                   "... joints xyz -> ... (joints xyz)")
+
+        # Rotate the vel_trajectory
+        vel_trajectory_local = torch.einsum("...j,...jk->...k", vel_trajectory,
+                                            rotations_inv)
+
+        # Stack things together
+        features = torch.cat((root_y[..., None], poses_features,
+                              vel_angles[..., None], vel_trajectory_local), -1)
+
+        # Normalize if needed
+        features = self.normalize(features)
+        return features
+
+    def inverse(self, features: Tensor) -> Tensor:
+        features = self.unnormalize(features)
+        root_y, poses_features, vel_angles, vel_trajectory_local = self.extract(
+            features)
+
+        # already have the good dimensionality
+        angles = torch.cumsum(vel_angles, dim=-1)
+        # First frame should be 0, but if infered it is better to ensure it
+        angles = angles - angles[..., [0]]
+
+        cos, sin = torch.cos(angles), torch.sin(angles)
+        rotations = matrix_of_angles(cos, sin, inv=False)
+
+        # Get back the poses
+        poses_local = rearrange(poses_features,
+                                "... (joints xyz) -> ... joints xyz",
+                                xyz=3)
+
+        # Rotate the poses
+        poses = torch.einsum("...lj,...jk->...lk", poses_local[..., [0, 2]],
+                             rotations)
+        poses = torch.stack(
+            (poses[..., 0], poses_local[..., 1], poses[..., 1]), axis=-1)
+
+        # Rotate the vel_trajectory
+        vel_trajectory = torch.einsum("...j,...jk->...k", vel_trajectory_local,
+                                      rotations)
+        # Integrate the trajectory
+        # Already have the good dimensionality
+        trajectory = torch.cumsum(vel_trajectory, dim=-2)
+        # First frame should be 0, but if infered it is better to ensure it
+        trajectory = trajectory - trajectory[..., [0], :]
+
+        # Add the root joints (which is still zero)
+        poses = torch.cat((0 * poses[..., [0], :], poses), -2)
+
+        # put back the root joint y
+        poses[..., 0, 1] = root_y
+
+        # Add the trajectory globally
+        poses[..., [0, 2]] += trajectory[..., None, :]
+        return poses
+
+    def extract(self, features: Tensor):
+        root_y = features[..., 0]
+        poses_features = features[..., 1:-3]
+        vel_angles = features[..., -3]
+        vel_trajectory_local = features[..., -2:]
+
+        return root_y, poses_features, vel_angles, vel_trajectory_local

mGPT/data/transforms/joints2jfeats/tools.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import torch
+import torch.nn.functional as F
+
+from mGPT.utils.joints import mmm_joints
+
+# Get the indexes of particular body part SMPLH case
+# Feet
+# LM, RM = smplh_joints.index("left_ankle"), smplh_joints.index("right_ankle")
+# LF, RF = smplh_joints.index("left_foot"), smplh_joints.index("right_foot")
+# # Shoulders
+# LS, RS = smplh_joints.index("left_shoulder"), smplh_joints.index("right_shoulder")
+# # Hips
+# LH, RH = smplh_joints.index("left_hip"), smplh_joints.index("right_hip")
+
+# Get the indexes of particular body part
+# Feet
+LM, RM = mmm_joints.index("LMrot"), mmm_joints.index("RMrot")
+LF, RF = mmm_joints.index("LF"), mmm_joints.index("RF")
+# Shoulders
+LS, RS = mmm_joints.index("LS"), mmm_joints.index("RS")
+# Hips
+LH, RH = mmm_joints.index("LH"), mmm_joints.index("RH")
+
+
+def get_forward_direction(poses, jointstype="mmm"):
+    # assert jointstype == 'mmm'
+    across = poses[..., RH, :] - poses[..., LH, :] + poses[..., RS, :] - poses[
+        ..., LS, :]
+    forward = torch.stack((-across[..., 2], across[..., 0]), axis=-1)
+    forward = torch.nn.functional.normalize(forward, dim=-1)
+    return forward
+
+
+def get_floor(poses, jointstype="mmm"):
+    # assert jointstype == 'mmm'
+    ndim = len(poses.shape)
+    foot_heights = poses[..., (LM, LF, RM, RF), 1].min(-1).values
+    floor_height = softmin(foot_heights, softness=0.5, dim=-1)
+    # changed this thing Mathis version 1.11 pytorch
+    return floor_height[(ndim - 2) * [None]].transpose(0, -1)
+
+
+def softmax(x, softness=1.0, dim=None):
+    maxi, mini = x.max(dim=dim).values, x.min(dim=dim).values
+    return maxi + torch.log(softness + torch.exp(mini - maxi))
+
+
+def softmin(x, softness=1.0, dim=0):
+    return -softmax(-x, softness=softness, dim=dim)
+
+
+def gaussian_filter1d(_inputs, sigma, truncate=4.0):
+    # Code adapted/mixed from scipy library into pytorch
+    # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L211
+    # and gaussian kernel
+    # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L179
+    # Correspond to mode="nearest" and order = 0
+    # But works batched
+    if len(_inputs.shape) == 2:
+        inputs = _inputs[None]
+    else:
+        inputs = _inputs
+
+    sd = float(sigma)
+    radius = int(truncate * sd + 0.5)
+    sigma2 = sigma * sigma
+    x = torch.arange(-radius,
+                     radius + 1,
+                     device=inputs.device,
+                     dtype=inputs.dtype)
+    phi_x = torch.exp(-0.5 / sigma2 * x**2)
+    phi_x = phi_x / phi_x.sum()
+
+    # Conv1d weights
+    groups = inputs.shape[-1]
+    weights = torch.tile(phi_x, (groups, 1, 1))
+    inputs = inputs.transpose(-1, -2)
+    outputs = F.conv1d(inputs, weights, padding="same",
+                       groups=groups).transpose(-1, -2)
+
+    return outputs.reshape(_inputs.shape)

mGPT/data/transforms/joints2rots/config.py

@@ -0,0 +1,119 @@
+import numpy as np
+from mGPT.utils.joints import mmm_joints, smplh2mmm_indexes
+
+# Map joints Name to SMPL joints idx
+JOINT_MAP = {
+    'MidHip': 0,
+    'LHip': 1,
+    'LKnee': 4,
+    'LAnkle': 7,
+    'LFoot': 10,
+    'RHip': 2,
+    'RKnee': 5,
+    'RAnkle': 8,
+    'RFoot': 11,
+    'LShoulder': 16,
+    'LElbow': 18,
+    'LWrist': 20,
+    'LHand': 22,
+    'RShoulder': 17,
+    'RElbow': 19,
+    'RWrist': 21,
+    'RHand': 23,
+    'spine1': 3,
+    'spine2': 6,
+    'spine3': 9,
+    'Neck': 12,
+    'Head': 15,
+    'LCollar': 13,
+    'Rcollar': 14,
+    'Nose': 24,
+    'REye': 26,
+    'LEye': 26,
+    'REar': 27,
+    'LEar': 28,
+    'LHeel': 31,
+    'RHeel': 34,
+    'OP RShoulder': 17,
+    'OP LShoulder': 16,
+    'OP RHip': 2,
+    'OP LHip': 1,
+    'OP Neck': 12,
+}
+
+mmm2smpl_correspondence = {
+    "root": "MidHip",
+    "BP": "spine1",
+    "BT": "spine3",
+    "BLN": "Neck",
+    "BUN": "Head",
+    "LS": "LShoulder",
+    "LE": "LElbow",
+    "LW": "LWrist",
+    "RS": "RShoulder",
+    "RE": "RElbow",
+    "RW": "RWrist",
+    "LH": "LHip",
+    "LK": "LKnee",
+    "LA": "LAnkle",
+    "LMrot": "LHeel",
+    "LF": "LFoot",
+    "RH": "RHip",
+    "RK": "RKnee",
+    "RA": "RAnkle",
+    "RMrot": "RHeel",
+    "RF": "RFoot"
+}
+
+full_smpl_idx = range(24)
+key_smpl_idx = [0, 1, 4, 7, 2, 5, 8, 17, 19, 21, 16, 18, 20]
+
+AMASS_JOINT_MAP = {
+    'MidHip': 0,
+    'LHip': 1,
+    'LKnee': 4,
+    'LAnkle': 7,
+    'LFoot': 10,
+    'RHip': 2,
+    'RKnee': 5,
+    'RAnkle': 8,
+    'RFoot': 11,
+    'LShoulder': 16,
+    'LElbow': 18,
+    'LWrist': 20,
+    'RShoulder': 17,
+    'RElbow': 19,
+    'RWrist': 21,
+    'spine1': 3,
+    'spine2': 6,
+    'spine3': 9,
+    'Neck': 12,
+    'Head': 15,
+    'LCollar': 13,
+    'Rcollar': 14,
+}
+amass_idx = range(22)
+amass_smpl_idx = range(22)
+
+# cal mmm in smpl index
+smpl2mmm_correspondence = {
+    val: key
+    for key, val in mmm2smpl_correspondence.items()
+}
+smpl2mmm_indexes = [JOINT_MAP[mmm2smpl_correspondence[x]] for x in mmm_joints]
+
+# cal mmm joints map
+MMM_JOINT_MAP = {
+    val: JOINT_MAP[val]
+    for key, val in mmm2smpl_correspondence.items()
+}
+
+# mmm_idx = range(21)
+# mmm_smpl_dix = smpl2mmm_indexes
+# mmm_smpl_dix = smplh2mmm_indexes
+# todo - configable
+SMPL_MODEL_DIR = "/apdcephfs/share_1227775/shingxchen/AIMotion/TMOSTData/deps/smpl_models/"
+GMM_MODEL_DIR = "/apdcephfs/share_1227775/shingxchen/AIMotion/TMOSTData/deps/smpl_models/"
+SMPL_MEAN_FILE = "/apdcephfs/share_1227775/shingxchen/AIMotion/TMOSTData/deps/smpl_models/neutral_smpl_mean_params.h5"
+# for collsion
+Part_Seg_DIR = "/apdcephfs/share_1227775/shingxchen/AIMotion/TMOSTData/deps/smpl_models/smplx_parts_segm.pkl"

mGPT/data/transforms/joints2rots/customloss.py

@@ -0,0 +1,217 @@
+import torch
+import torch.nn.functional as F
+import config
+
+# Guassian
+def gmof(x, sigma):
+    """
+    Geman-McClure error function
+    """
+    x_squared = x ** 2
+    sigma_squared = sigma ** 2
+    return (sigma_squared * x_squared) / (sigma_squared + x_squared)
+
+# angle prior
+def angle_prior(pose):
+    """
+    Angle prior that penalizes unnatural bending of the knees and elbows
+    """
+    # We subtract 3 because pose does not include the global rotation of the model
+    return torch.exp(
+        pose[:, [55 - 3, 58 - 3, 12 - 3, 15 - 3]] * torch.tensor([1., -1., -1, -1.], device=pose.device)) ** 2
+
+
+def perspective_projection(points, rotation, translation,
+                           focal_length, camera_center):
+    """
+    This function computes the perspective projection of a set of points.
+    Input:
+        points (bs, N, 3): 3D points
+        rotation (bs, 3, 3): Camera rotation
+        translation (bs, 3): Camera translation
+        focal_length (bs,) or scalar: Focal length
+        camera_center (bs, 2): Camera center
+    """
+    batch_size = points.shape[0]
+    K = torch.zeros([batch_size, 3, 3], device=points.device)
+    K[:, 0, 0] = focal_length
+    K[:, 1, 1] = focal_length
+    K[:, 2, 2] = 1.
+    K[:, :-1, -1] = camera_center
+
+    # Transform points
+    points = torch.einsum('bij,bkj->bki', rotation, points)
+    points = points + translation.unsqueeze(1)
+
+    # Apply perspective distortion
+    projected_points = points / points[:, :, -1].unsqueeze(-1)
+
+    # Apply camera intrinsics
+    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)
+
+    return projected_points[:, :, :-1]
+
+
+def body_fitting_loss(body_pose, betas, model_joints, camera_t, camera_center,
+                      joints_2d, joints_conf, pose_prior,
+                      focal_length=5000, sigma=100, pose_prior_weight=4.78,
+                      shape_prior_weight=5, angle_prior_weight=15.2,
+                      output='sum'):
+    """
+    Loss function for body fitting
+    """
+    batch_size = body_pose.shape[0]
+    rotation = torch.eye(3, device=body_pose.device).unsqueeze(0).expand(batch_size, -1, -1)
+
+    projected_joints = perspective_projection(model_joints, rotation, camera_t,
+                                              focal_length, camera_center)
+
+    # Weighted robust reprojection error
+    reprojection_error = gmof(projected_joints - joints_2d, sigma)
+    reprojection_loss = (joints_conf ** 2) * reprojection_error.sum(dim=-1)
+
+    # Pose prior loss
+    pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas)
+
+    # Angle prior for knees and elbows
+    angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1)
+
+    # Regularizer to prevent betas from taking large values
+    shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1)
+
+    total_loss = reprojection_loss.sum(dim=-1) + pose_prior_loss + angle_prior_loss + shape_prior_loss
+
+    if output == 'sum':
+        return total_loss.sum()
+    elif output == 'reprojection':
+        return reprojection_loss
+
+
+# --- get camera fitting loss -----
+def camera_fitting_loss(model_joints, camera_t, camera_t_est, camera_center, 
+                        joints_2d, joints_conf,
+                        focal_length=5000, depth_loss_weight=100):
+    """
+    Loss function for camera optimization.
+    """
+    # Project model joints
+    batch_size = model_joints.shape[0]
+    rotation = torch.eye(3, device=model_joints.device).unsqueeze(0).expand(batch_size, -1, -1)
+    projected_joints = perspective_projection(model_joints, rotation, camera_t,
+                                              focal_length, camera_center)
+
+    # get the indexed four
+    op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder']
+    op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints]
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+
+    reprojection_error_op = (joints_2d[:, op_joints_ind] -
+                             projected_joints[:, op_joints_ind]) ** 2
+    reprojection_error_gt = (joints_2d[:, gt_joints_ind] -
+                             projected_joints[:, gt_joints_ind]) ** 2
+
+    # Check if for each example in the batch all 4 OpenPose detections are valid, otherwise use the GT detections
+    # OpenPose joints are more reliable for this task, so we prefer to use them if possible
+    is_valid = (joints_conf[:, op_joints_ind].min(dim=-1)[0][:, None, None] > 0).float()
+    reprojection_loss = (is_valid * reprojection_error_op + (1 - is_valid) * reprojection_error_gt).sum(dim=(1, 2))
+
+    # Loss that penalizes deviation from depth estimate
+    depth_loss = (depth_loss_weight ** 2) * (camera_t[:, 2] - camera_t_est[:, 2]) ** 2
+
+    total_loss = reprojection_loss + depth_loss
+    return total_loss.sum()
+
+
+
+ # #####--- body fitiing loss -----
+def body_fitting_loss_3d(body_pose, preserve_pose,
+                         betas, model_joints, camera_translation,
+                         j3d, pose_prior,
+                         joints3d_conf,
+                         sigma=100, pose_prior_weight=4.78*1.5,
+                         shape_prior_weight=5.0, angle_prior_weight=15.2,
+                         joint_loss_weight=500.0,
+                         pose_preserve_weight=0.0,
+                         use_collision=False,
+                         model_vertices=None, model_faces=None,
+                         search_tree=None,  pen_distance=None,  filter_faces=None,
+                         collision_loss_weight=1000
+                         ):
+    """
+    Loss function for body fitting
+    """
+    batch_size = body_pose.shape[0]
+
+    #joint3d_loss = (joint_loss_weight ** 2) * gmof((model_joints + camera_translation) - j3d, sigma).sum(dim=-1)
+    
+    joint3d_error = gmof((model_joints + camera_translation) - j3d, sigma)
+    
+    joint3d_loss_part = (joints3d_conf ** 2) * joint3d_error.sum(dim=-1)
+    joint3d_loss = (joint_loss_weight ** 2) * joint3d_loss_part
+    
+    # Pose prior loss
+    pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas)
+    # Angle prior for knees and elbows
+    angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1)
+    # Regularizer to prevent betas from taking large values
+    shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1)
+
+    collision_loss = 0.0
+    # Calculate the loss due to interpenetration
+    if use_collision:
+        triangles = torch.index_select(
+            model_vertices, 1,
+            model_faces).view(batch_size, -1, 3, 3)
+
+        with torch.no_grad():
+            collision_idxs = search_tree(triangles)
+
+        # Remove unwanted collisions
+        if filter_faces is not None:
+            collision_idxs = filter_faces(collision_idxs)
+
+        if collision_idxs.ge(0).sum().item() > 0:
+            collision_loss = torch.sum(collision_loss_weight * pen_distance(triangles, collision_idxs))
+    
+    pose_preserve_loss = (pose_preserve_weight ** 2) * ((body_pose - preserve_pose) ** 2).sum(dim=-1)
+
+    total_loss = joint3d_loss + pose_prior_loss + angle_prior_loss + shape_prior_loss + collision_loss + pose_preserve_loss
+
+    return total_loss.sum()
+
+
+# #####--- get camera fitting loss -----
+def camera_fitting_loss_3d(model_joints, camera_t, camera_t_est,
+                           j3d, joints_category="orig", depth_loss_weight=100.0):
+    """
+    Loss function for camera optimization.
+    """
+    model_joints = model_joints + camera_t
+    # # get the indexed four
+    # op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder']
+    # op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints]
+    #
+    # j3d_error_loss = (j3d[:, op_joints_ind] -
+    #                          model_joints[:, op_joints_ind]) ** 2
+
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    
+    if joints_category=="orig":
+        select_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    elif joints_category=="AMASS":
+        select_joints_ind = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints]
+    elif joints_category=="MMM":
+        select_joints_ind = [config.MMM_JOINT_MAP[joint] for joint in gt_joints]
+    else:
+        print("NO SUCH JOINTS CATEGORY!")
+
+    j3d_error_loss = (j3d[:, select_joints_ind] -
+                      model_joints[:, gt_joints_ind]) ** 2
+
+    # Loss that penalizes deviation from depth estimate
+    depth_loss = (depth_loss_weight**2) *  (camera_t - camera_t_est)**2
+
+    total_loss = j3d_error_loss +  depth_loss
+    return total_loss.sum()

mGPT/data/transforms/joints2rots/prior.py

@@ -0,0 +1,229 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import sys
+import os
+
+import time
+import pickle
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+DEFAULT_DTYPE = torch.float32
+
+
+def create_prior(prior_type, **kwargs):
+    if prior_type == 'gmm':
+        prior = MaxMixturePrior(**kwargs)
+    elif prior_type == 'l2':
+        return L2Prior(**kwargs)
+    elif prior_type == 'angle':
+        return SMPLifyAnglePrior(**kwargs)
+    elif prior_type == 'none' or prior_type is None:
+        # Don't use any pose prior
+        def no_prior(*args, **kwargs):
+            return 0.0
+        prior = no_prior
+    else:
+        raise ValueError('Prior {}'.format(prior_type) + ' is not implemented')
+    return prior
+
+
+class SMPLifyAnglePrior(nn.Module):
+    def __init__(self, dtype=torch.float32, **kwargs):
+        super(SMPLifyAnglePrior, self).__init__()
+
+        # Indices for the roration angle of
+        # 55: left elbow,  90deg bend at -np.pi/2
+        # 58: right elbow, 90deg bend at np.pi/2
+        # 12: left knee,   90deg bend at np.pi/2
+        # 15: right knee,  90deg bend at np.pi/2
+        angle_prior_idxs = np.array([55, 58, 12, 15], dtype=np.int64)
+        angle_prior_idxs = torch.tensor(angle_prior_idxs, dtype=torch.long)
+        self.register_buffer('angle_prior_idxs', angle_prior_idxs)
+
+        angle_prior_signs = np.array([1, -1, -1, -1],
+                                     dtype=np.float6432 if dtype == torch.float32
+                                     else np.float6464)
+        angle_prior_signs = torch.tensor(angle_prior_signs,
+                                         dtype=dtype)
+        self.register_buffer('angle_prior_signs', angle_prior_signs)
+
+    def forward(self, pose, with_global_pose=False):
+        ''' Returns the angle prior loss for the given pose
+        Args:
+            pose: (Bx[23 + 1] * 3) torch tensor with the axis-angle
+            representation of the rotations of the joints of the SMPL model.
+        Kwargs:
+            with_global_pose: Whether the pose vector also contains the global
+            orientation of the SMPL model. If not then the indices must be
+            corrected.
+        Returns:
+            A sze (B) tensor containing the angle prior loss for each element
+            in the batch.
+        '''
+        angle_prior_idxs = self.angle_prior_idxs - (not with_global_pose) * 3
+        return torch.exp(pose[:, angle_prior_idxs] *
+                         self.angle_prior_signs).pow(2)
+
+
+class L2Prior(nn.Module):
+    def __init__(self, dtype=DEFAULT_DTYPE, reduction='sum', **kwargs):
+        super(L2Prior, self).__init__()
+
+    def forward(self, module_input, *args):
+        return torch.sum(module_input.pow(2))
+
+
+class MaxMixturePrior(nn.Module):
+
+    def __init__(self, prior_folder='prior',
+                 num_gaussians=6, dtype=DEFAULT_DTYPE, epsilon=1e-16,
+                 use_merged=True,
+                 **kwargs):
+        super(MaxMixturePrior, self).__init__()
+
+        if dtype == DEFAULT_DTYPE:
+            np_dtype = np.float6432
+        elif dtype == torch.float64:
+            np_dtype = np.float6464
+        else:
+            print('Unknown float type {}, exiting!'.format(dtype))
+            sys.exit(-1)
+
+        self.num_gaussians = num_gaussians
+        self.epsilon = epsilon
+        self.use_merged = use_merged
+        gmm_fn = 'gmm_{:02d}.pkl'.format(num_gaussians)
+
+        full_gmm_fn = os.path.join(prior_folder, gmm_fn)
+        if not os.path.exists(full_gmm_fn):
+            print('The path to the mixture prior "{}"'.format(full_gmm_fn) +
+                  ' does not exist, exiting!')
+            sys.exit(-1)
+
+        with open(full_gmm_fn, 'rb') as f:
+            gmm = pickle.load(f, encoding='latin1')
+
+        if type(gmm) == dict:
+            means = gmm['means'].astype(np_dtype)
+            covs = gmm['covars'].astype(np_dtype)
+            weights = gmm['weights'].astype(np_dtype)
+        elif 'sklearn.mixture.gmm.GMM' in str(type(gmm)):
+            means = gmm.means_.astype(np_dtype)
+            covs = gmm.covars_.astype(np_dtype)
+            weights = gmm.weights_.astype(np_dtype)
+        else:
+            print('Unknown type for the prior: {}, exiting!'.format(type(gmm)))
+            sys.exit(-1)
+
+        self.register_buffer('means', torch.tensor(means, dtype=dtype))
+
+        self.register_buffer('covs', torch.tensor(covs, dtype=dtype))
+
+        precisions = [np.linalg.inv(cov) for cov in covs]
+        precisions = np.stack(precisions).astype(np_dtype)
+
+        self.register_buffer('precisions',
+                             torch.tensor(precisions, dtype=dtype))
+
+        # The constant term:
+        sqrdets = np.array([(np.sqrt(np.linalg.det(c)))
+                            for c in gmm['covars']])
+        const = (2 * np.pi)**(69 / 2.)
+
+        nll_weights = np.asarray(gmm['weights'] / (const *
+                                                   (sqrdets / sqrdets.min())))
+        nll_weights = torch.tensor(nll_weights, dtype=dtype).unsqueeze(dim=0)
+        self.register_buffer('nll_weights', nll_weights)
+
+        weights = torch.tensor(gmm['weights'], dtype=dtype).unsqueeze(dim=0)
+        self.register_buffer('weights', weights)
+
+        self.register_buffer('pi_term',
+                             torch.log(torch.tensor(2 * np.pi, dtype=dtype)))
+
+        cov_dets = [np.log(np.linalg.det(cov.astype(np_dtype)) + epsilon)
+                    for cov in covs]
+        self.register_buffer('cov_dets',
+                             torch.tensor(cov_dets, dtype=dtype))
+
+        # The dimensionality of the random variable
+        self.random_var_dim = self.means.shape[1]
+
+    def get_mean(self):
+        ''' Returns the mean of the mixture '''
+        mean_pose = torch.matmul(self.weights, self.means)
+        return mean_pose
+
+    def merged_log_likelihood(self, pose, betas):
+        diff_from_mean = pose.unsqueeze(dim=1) - self.means
+
+        prec_diff_prod = torch.einsum('mij,bmj->bmi',
+                                      [self.precisions, diff_from_mean])
+        diff_prec_quadratic = (prec_diff_prod * diff_from_mean).sum(dim=-1)
+
+        curr_loglikelihood = 0.5 * diff_prec_quadratic - \
+            torch.log(self.nll_weights)
+        #  curr_loglikelihood = 0.5 * (self.cov_dets.unsqueeze(dim=0) +
+        #  self.random_var_dim * self.pi_term +
+        #  diff_prec_quadratic
+        #  ) - torch.log(self.weights)
+
+        min_likelihood, _ = torch.min(curr_loglikelihood, dim=1)
+        return min_likelihood
+
+    def log_likelihood(self, pose, betas, *args, **kwargs):
+        ''' Create graph operation for negative log-likelihood calculation
+        '''
+        likelihoods = []
+
+        for idx in range(self.num_gaussians):
+            mean = self.means[idx]
+            prec = self.precisions[idx]
+            cov = self.covs[idx]
+            diff_from_mean = pose - mean
+
+            curr_loglikelihood = torch.einsum('bj,ji->bi',
+                                              [diff_from_mean, prec])
+            curr_loglikelihood = torch.einsum('bi,bi->b',
+                                              [curr_loglikelihood,
+                                               diff_from_mean])
+            cov_term = torch.log(torch.det(cov) + self.epsilon)
+            curr_loglikelihood += 0.5 * (cov_term +
+                                         self.random_var_dim *
+                                         self.pi_term)
+            likelihoods.append(curr_loglikelihood)
+
+        log_likelihoods = torch.stack(likelihoods, dim=1)
+        min_idx = torch.argmin(log_likelihoods, dim=1)
+        weight_component = self.nll_weights[:, min_idx]
+        weight_component = -torch.log(weight_component)
+
+        return weight_component + log_likelihoods[:, min_idx]
+
+    def forward(self, pose, betas):
+        if self.use_merged:
+            return self.merged_log_likelihood(pose, betas)
+        else:
+            return self.log_likelihood(pose, betas)

mGPT/data/transforms/joints2rots/smplify.py

@@ -0,0 +1,284 @@
+import torch
+import os, sys
+import pickle
+import smplx
+import numpy as np
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(__file__))
+from customloss import (camera_fitting_loss, 
+                        body_fitting_loss, 
+                        camera_fitting_loss_3d,
+                        body_fitting_loss_3d, 
+                        )
+from prior import MaxMixturePrior
+import config
+
+
+
+@torch.no_grad()
+def guess_init_3d(model_joints, 
+                  j3d, 
+                  joints_category="orig"):
+    """Initialize the camera translation via triangle similarity, by using the torso joints        .
+    :param model_joints: SMPL model with pre joints
+    :param j3d: 25x3 array of Kinect Joints
+    :returns: 3D vector corresponding to the estimated camera translation
+    """
+    # get the indexed four
+    gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder']
+    gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints]
+    
+    if joints_category=="orig":
+        joints_ind_category = [config.JOINT_MAP[joint] for joint in gt_joints]
+    elif joints_category=="AMASS":
+        joints_ind_category = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints] 
+    elif joints_category=="MMM":
+        joints_ind_category = [config.MMM_JOINT_MAP[joint] for joint in gt_joints] 
+    else:
+        print("NO SUCH JOINTS CATEGORY!") 
+
+    sum_init_t = (j3d[:, joints_ind_category] - model_joints[:, gt_joints_ind]).sum(dim=1)
+    init_t = sum_init_t / 4.0
+    return init_t
+
+
+# SMPLIfy 3D
+class SMPLify3D():
+    """Implementation of SMPLify, use 3D joints."""
+
+    def __init__(self,
+                 smplxmodel,
+                 step_size=1e-2,
+                 batch_size=1,
+                 num_iters=100,
+                 use_collision=False,
+                 use_lbfgs=True,
+                 joints_category="orig",
+                 device=torch.device('cuda:0'),
+                 ):
+
+        # Store options
+        self.batch_size = batch_size
+        self.device = device
+        self.step_size = step_size
+
+        self.num_iters = num_iters
+        # --- choose optimizer
+        self.use_lbfgs = use_lbfgs
+        # GMM pose prior
+        self.pose_prior = MaxMixturePrior(prior_folder=config.GMM_MODEL_DIR,
+                                          num_gaussians=8,
+                                          dtype=torch.float32).to(device)
+        # collision part
+        self.use_collision = use_collision
+        if self.use_collision:
+            self.part_segm_fn = config.Part_Seg_DIR
+        
+        # reLoad SMPL-X model
+        self.smpl = smplxmodel
+
+        self.model_faces = smplxmodel.faces_tensor.view(-1)
+
+        # select joint joint_category
+        self.joints_category = joints_category
+        
+        if joints_category=="orig":
+            self.smpl_index = config.full_smpl_idx
+            self.corr_index = config.full_smpl_idx 
+        elif joints_category=="AMASS":
+            self.smpl_index = config.amass_smpl_idx
+            self.corr_index = config.amass_idx
+        # elif joints_category=="MMM":
+        #     self.smpl_index = config.mmm_smpl_dix
+        #     self.corr_index = config.mmm_idx
+        else:
+            self.smpl_index = None 
+            self.corr_index = None
+            print("NO SUCH JOINTS CATEGORY!")
+
+    # ---- get the man function here ------
+    def __call__(self, init_pose, init_betas, init_cam_t, j3d, conf_3d=1.0, seq_ind=0):
+        """Perform body fitting.
+        Input:
+            init_pose: SMPL pose estimate
+            init_betas: SMPL betas estimate
+            init_cam_t: Camera translation estimate
+            j3d: joints 3d aka keypoints
+            conf_3d: confidence for 3d joints
+			seq_ind: index of the sequence
+        Returns:
+            vertices: Vertices of optimized shape
+            joints: 3D joints of optimized shape
+            pose: SMPL pose parameters of optimized shape
+            betas: SMPL beta parameters of optimized shape
+            camera_translation: Camera translation
+        """
+
+        # # # add the mesh inter-section to avoid
+        search_tree = None
+        pen_distance = None
+        filter_faces = None
+        
+        if self.use_collision:
+            from mesh_intersection.bvh_search_tree import BVH
+            import mesh_intersection.loss as collisions_loss
+            from mesh_intersection.filter_faces import FilterFaces
+
+            search_tree = BVH(max_collisions=8)
+
+            pen_distance = collisions_loss.DistanceFieldPenetrationLoss(
+                           sigma=0.5, point2plane=False, vectorized=True, penalize_outside=True)
+
+            if self.part_segm_fn:
+                # Read the part segmentation
+                part_segm_fn = os.path.expandvars(self.part_segm_fn)
+                with open(part_segm_fn, 'rb') as faces_parents_file:
+                    face_segm_data = pickle.load(faces_parents_file,  encoding='latin1')
+                faces_segm = face_segm_data['segm']
+                faces_parents = face_segm_data['parents']
+                # Create the module used to filter invalid collision pairs
+                filter_faces = FilterFaces(
+                    faces_segm=faces_segm, faces_parents=faces_parents,
+                    ign_part_pairs=None).to(device=self.device)
+                    
+                    
+        # Split SMPL pose to body pose and global orientation
+        body_pose = init_pose[:, 3:].detach().clone()
+        global_orient = init_pose[:, :3].detach().clone()
+        betas = init_betas.detach().clone()
+
+        # use guess 3d to get the initial
+        smpl_output = self.smpl(global_orient=global_orient,
+                                body_pose=body_pose,
+                                betas=betas)
+        model_joints = smpl_output.joints
+
+        init_cam_t = guess_init_3d(model_joints, j3d, self.joints_category).detach()
+        camera_translation = init_cam_t.clone()
+        
+        preserve_pose = init_pose[:, 3:].detach().clone()
+       # -------------Step 1: Optimize camera translation and body orientation--------
+        # Optimize only camera translation and body orientation
+        body_pose.requires_grad = False
+        betas.requires_grad = False
+        global_orient.requires_grad = True
+        camera_translation.requires_grad = True
+
+        camera_opt_params = [global_orient, camera_translation]
+
+        if self.use_lbfgs:
+            camera_optimizer = torch.optim.LBFGS(camera_opt_params, max_iter=self.num_iters,
+                                                 lr=self.step_size, line_search_fn='strong_wolfe')
+            for i in range(10):
+                def closure():
+                    camera_optimizer.zero_grad()
+                    smpl_output = self.smpl(global_orient=global_orient,
+                                            body_pose=body_pose,
+                                            betas=betas)
+                    model_joints = smpl_output.joints
+
+                    loss = camera_fitting_loss_3d(model_joints, camera_translation,
+                                                  init_cam_t, j3d, self.joints_category)
+                    loss.backward()
+                    return loss
+
+                camera_optimizer.step(closure)
+        else:
+            camera_optimizer = torch.optim.Adam(camera_opt_params, lr=self.step_size, betas=(0.9, 0.999))
+
+            for i in range(20):
+                smpl_output = self.smpl(global_orient=global_orient,
+                                        body_pose=body_pose,
+                                        betas=betas)
+                model_joints = smpl_output.joints
+
+                loss = camera_fitting_loss_3d(model_joints[:, self.smpl_index], camera_translation,
+                                              init_cam_t,  j3d[:, self.corr_index], self.joints_category)
+                camera_optimizer.zero_grad()
+                loss.backward()
+                camera_optimizer.step()
+
+        # Fix camera translation after optimizing camera
+        # --------Step 2: Optimize body joints --------------------------
+        # Optimize only the body pose and global orientation of the body
+        body_pose.requires_grad = True
+        global_orient.requires_grad = True
+        camera_translation.requires_grad = True
+
+        # --- if we use the sequence, fix the shape
+        if seq_ind == 0:
+            betas.requires_grad = True
+            body_opt_params = [body_pose, betas, global_orient, camera_translation]
+        else:
+            betas.requires_grad = False
+            body_opt_params = [body_pose, global_orient, camera_translation]
+
+        if self.use_lbfgs:
+            body_optimizer = torch.optim.LBFGS(body_opt_params, max_iter=self.num_iters,
+                                               lr=self.step_size, line_search_fn='strong_wolfe')
+            
+            for i in tqdm(range(self.num_iters), desc=f"LBFGS iter: "):
+            # for i in range(self.num_iters):
+                def closure():
+                    body_optimizer.zero_grad()
+                    smpl_output = self.smpl(global_orient=global_orient,
+                                            body_pose=body_pose,
+                                            betas=betas)
+                    model_joints = smpl_output.joints
+                    model_vertices = smpl_output.vertices
+
+                    loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                                j3d[:, self.corr_index], self.pose_prior,
+                                                joints3d_conf=conf_3d,
+                                                joint_loss_weight=600.0,
+                                                pose_preserve_weight=5.0,
+                                                use_collision=self.use_collision, 
+                                                model_vertices=model_vertices, model_faces=self.model_faces,
+                                                search_tree=search_tree, pen_distance=pen_distance, filter_faces=filter_faces)
+                    loss.backward()
+                    return loss
+
+                body_optimizer.step(closure)
+        else:
+            body_optimizer = torch.optim.Adam(body_opt_params, lr=self.step_size, betas=(0.9, 0.999))
+
+            for i in range(self.num_iters):
+                smpl_output = self.smpl(global_orient=global_orient,
+                                        body_pose=body_pose,
+                                        betas=betas)
+                model_joints = smpl_output.joints
+                model_vertices = smpl_output.vertices
+
+                loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                            j3d[:, self.corr_index], self.pose_prior,
+                                            joints3d_conf=conf_3d,
+                                            joint_loss_weight=600.0,
+                                            use_collision=self.use_collision, 
+                                            model_vertices=model_vertices, model_faces=self.model_faces,
+                                            search_tree=search_tree,  pen_distance=pen_distance,  filter_faces=filter_faces)
+                body_optimizer.zero_grad()
+                loss.backward()
+                body_optimizer.step()
+
+        # Get final loss value
+        with torch.no_grad():
+            smpl_output = self.smpl(global_orient=global_orient,
+                                    body_pose=body_pose,
+                                    betas=betas, return_full_pose=True)
+            model_joints = smpl_output.joints
+            model_vertices = smpl_output.vertices
+
+            final_loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation,
+                                              j3d[:, self.corr_index], self.pose_prior,
+                                              joints3d_conf=conf_3d,
+                                              joint_loss_weight=600.0,
+                                              use_collision=self.use_collision, model_vertices=model_vertices, model_faces=self.model_faces,
+                                              search_tree=search_tree,  pen_distance=pen_distance,  filter_faces=filter_faces)
+
+        vertices = smpl_output.vertices.detach()
+        joints = smpl_output.joints.detach()
+        pose = torch.cat([global_orient, body_pose], dim=-1).detach()
+        betas = betas.detach()
+
+        return vertices, joints, pose, betas, camera_translation, final_loss

mGPT/data/transforms/rots2joints/__init__.py

@@ -0,0 +1,3 @@
+from .base import Rots2Joints
+from .smplh import SMPLH
+from .smplx import SMPLX

mGPT/data/transforms/rots2joints/base.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+from pathlib import Path
+import os
+# import hydra
+
+class Rots2Joints(nn.Module):
+    def __init__(self, path: Optional[str] = None,
+                 normalization: bool = False,
+                 eps: float = 1e-12,
+                 **kwargs) -> None:
+        if normalization and path is None:
+            raise TypeError("You should provide a path if normalization is on.")
+
+        super().__init__()
+        self.normalization = normalization
+        self.eps = eps
+        # workaround for cluster local/sync
+        if path is not None:
+            rel_p = path.split('/')
+            rel_p = rel_p[rel_p.index('deps'):]
+            rel_p = '/'.join(rel_p)
+            # path = hydra.utils.get_original_cwd() + '/' + rel_p
+        if normalization:
+            mean_path = Path(path) / "mean.pt"
+            std_path = Path(path) / "std.pt"
+            self.register_buffer('mean', torch.load(mean_path))
+            self.register_buffer('std', torch.load(std_path))
+
+    def normalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = (features - self.mean)/(self.std + self.eps)
+        return features
+
+    def unnormalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = features * self.std + self.mean
+        return features

mGPT/data/transforms/rots2joints/smplh.py

@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import contextlib
+from typing import Optional
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from mGPT.utils.joints import smplh_to_mmm_scaling_factor
+from mGPT.utils.joints import smplh2mmm_indexes
+from .base import Rots2Joints
+
+
+def slice_or_none(data, cslice):
+    if data is None:
+        return data
+    else:
+        return data[cslice]
+
+
+class SMPLH(Rots2Joints):
+
+    def __init__(self,
+                 path: str,
+                 jointstype: str = "mmm",
+                 input_pose_rep: str = "matrix",
+                 batch_size: int = 512,
+                 gender="neutral",
+                 **kwargs) -> None:
+        super().__init__(path=None, normalization=False)
+        self.batch_size = batch_size
+        self.input_pose_rep = input_pose_rep
+        self.jointstype = jointstype
+        self.training = False
+
+        from smplx.body_models import SMPLHLayer
+        import os
+        # rel_p = path.split('/')
+        # rel_p = rel_p[rel_p.index('data'):]
+        # rel_p = '/'.join(rel_p)
+
+        # Remove annoying print
+        with contextlib.redirect_stdout(None):
+            self.smplh = SMPLHLayer(path, ext="pkl", gender=gender).eval()
+
+        self.faces = self.smplh.faces
+        for p in self.parameters():
+            p.requires_grad = False
+
+    def train(self, *args, **kwargs):
+        return self
+
+    def forward(self,
+                smpl_data: dict,
+                jointstype: Optional[str] = None,
+                input_pose_rep: Optional[str] = None,
+                batch_size: Optional[int] = None) -> Tensor:
+
+        # Take values from init if not specified there
+        jointstype = self.jointstype if jointstype is None else jointstype
+        batch_size = self.batch_size if batch_size is None else batch_size
+        input_pose_rep = self.input_pose_rep if input_pose_rep is None else input_pose_rep
+
+        if input_pose_rep == "xyz":
+            raise NotImplementedError(
+                "You should use identity pose2joints instead")
+
+        poses = smpl_data.rots
+        trans = smpl_data.trans
+
+        from functools import reduce
+        import operator
+        save_shape_bs_len = poses.shape[:-3]
+        nposes = reduce(operator.mul, save_shape_bs_len, 1)
+
+        if poses.shape[-3] == 52:
+            nohands = False
+        elif poses.shape[-3] == 22:
+            nohands = True
+        else:
+            raise NotImplementedError("Could not parse the poses.")
+
+        # Convert any rotations to matrix
+        # from temos.tools.easyconvert import to_matrix
+        # matrix_poses = to_matrix(input_pose_rep, poses)
+        matrix_poses = poses
+
+        # Reshaping
+        matrix_poses = matrix_poses.reshape((nposes, *matrix_poses.shape[-3:]))
+        global_orient = matrix_poses[:, 0]
+
+        if trans is None:
+            trans = torch.zeros((*save_shape_bs_len, 3),
+                                dtype=poses.dtype,
+                                device=poses.device)
+
+        trans_all = trans.reshape((nposes, *trans.shape[-1:]))
+
+        body_pose = matrix_poses[:, 1:22]
+        if nohands:
+            left_hand_pose = None
+            right_hand_pose = None
+        else:
+            hand_pose = matrix_poses[:, 22:]
+            left_hand_pose = hand_pose[:, :15]
+            right_hand_pose = hand_pose[:, 15:]
+
+        n = len(body_pose)
+        outputs = []
+        for chunk in range(int((n - 1) / batch_size) + 1):
+            chunk_slice = slice(chunk * batch_size, (chunk + 1) * batch_size)
+            smpl_output = self.smplh(
+                global_orient=slice_or_none(global_orient, chunk_slice),
+                body_pose=slice_or_none(body_pose, chunk_slice),
+                left_hand_pose=slice_or_none(left_hand_pose, chunk_slice),
+                right_hand_pose=slice_or_none(right_hand_pose, chunk_slice),
+                transl=slice_or_none(trans_all, chunk_slice))
+
+            if jointstype == "vertices":
+                output_chunk = smpl_output.vertices
+            else:
+                joints = smpl_output.joints
+                output_chunk = joints
+
+            outputs.append(output_chunk)
+
+        outputs = torch.cat(outputs)
+        outputs = outputs.reshape((*save_shape_bs_len, *outputs.shape[1:]))
+
+        # Change topology if needed
+        outputs = smplh_to(jointstype, outputs, trans)
+
+        return outputs
+
+    def inverse(self, joints: Tensor) -> Tensor:
+        raise NotImplementedError("Cannot inverse SMPLH layer.")
+
+
+def smplh_to(jointstype, data, trans):
+    from mGPT.utils.joints import get_root_idx
+
+    if "mmm" in jointstype:
+        from mGPT.utils.joints import smplh2mmm_indexes
+        indexes = smplh2mmm_indexes
+        data = data[..., indexes, :]
+
+        # make it compatible with mmm
+        if jointstype == "mmm":
+            from mGPT.utils.joints import smplh_to_mmm_scaling_factor
+            data *= smplh_to_mmm_scaling_factor
+
+        if jointstype == "smplmmm":
+            pass
+        elif jointstype in ["mmm", "mmmns"]:
+            # swap axis
+            data = data[..., [1, 2, 0]]
+            # revert left and right
+            data[..., 2] = -data[..., 2]
+
+    elif jointstype == "smplnh":
+        from mGPT.utils.joints import smplh2smplnh_indexes
+        indexes = smplh2smplnh_indexes
+        data = data[..., indexes, :]
+    elif jointstype == "smplh":
+        pass
+    elif jointstype == "vertices":
+        pass
+    else:
+        raise NotImplementedError(f"SMPLH to {jointstype} is not implemented.")
+
+    if jointstype != "vertices":
+        # shift the output in each batch
+        # such that it is centered on the pelvis/root on the first frame
+        root_joint_idx = get_root_idx(jointstype)
+        shift = trans[..., 0, :] - data[..., 0, root_joint_idx, :]
+        data += shift[..., None, None, :]
+
+    return data

mGPT/data/transforms/rots2joints/smplx.py

@@ -0,0 +1,201 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import contextlib
+from typing import Optional
+import torch
+from torch import Tensor
+from mGPT.utils.joints import smplh_to_mmm_scaling_factor, smplh2mmm_indexes, get_root_idx
+from mGPT.utils.easyconvert import rep_to_rep
+from .base import Rots2Joints
+
+
+def slice_or_none(data, cslice):
+    if data is None:
+        return data
+    else:
+        return data[cslice]
+
+
+class SMPLX(Rots2Joints):
+    def __init__(self,
+                 path: str,
+                 jointstype: str = "mmm",
+                 input_pose_rep: str = "matrix",
+                 batch_size: int = 512,
+                 gender="neutral",
+                 **kwargs) -> None:
+        super().__init__(path=None, normalization=False)
+        self.batch_size = batch_size
+        self.input_pose_rep = input_pose_rep
+        self.jointstype = jointstype
+        self.training = False
+
+        from smplx.body_models import SMPLXLayer
+        import os
+        # rel_p = path.split('/')
+        # rel_p = rel_p[rel_p.index('data'):]
+        # rel_p = '/'.join(rel_p)
+
+        # Remove annoying print
+        with contextlib.redirect_stdout(None):
+            self.smplx = SMPLXLayer(path,
+                                    ext="npz",
+                                    gender=gender,
+                                    batch_size=batch_size).eval()
+
+        self.faces = self.smplx.faces
+        for p in self.parameters():
+            p.requires_grad = False
+
+    def train(self, *args, **kwargs):
+        return self
+
+    def forward(self,
+                smpl_data: dict,
+                jointstype: Optional[str] = None,
+                input_pose_rep: Optional[str] = None,
+                batch_size: Optional[int] = None) -> Tensor:
+
+        # Take values from init if not specified there
+        jointstype = self.jointstype if jointstype is None else jointstype
+        batch_size = self.batch_size if batch_size is None else batch_size
+        input_pose_rep = self.input_pose_rep if input_pose_rep is None else input_pose_rep
+
+        poses = smpl_data.rots
+        trans = smpl_data.trans
+
+        from functools import reduce
+        import operator
+        save_shape_bs_len = poses.shape[:-3]
+        nposes = reduce(operator.mul, save_shape_bs_len, 1)
+
+        
+        matrix_poses = rep_to_rep(self.input_pose_rep, input_pose_rep, poses)
+
+        # Reshaping
+        matrix_poses = matrix_poses.reshape((nposes, *matrix_poses.shape[-3:]))
+
+        global_orient = matrix_poses[:, 0]
+
+        if trans is None:
+            trans = torch.zeros((*save_shape_bs_len, 3),
+                                dtype=poses.dtype,
+                                device=poses.device)
+
+        trans_all = trans.reshape((nposes, *trans.shape[-1:]))
+
+        body_pose = matrix_poses[:, 1:22]
+
+        if poses.shape[-3] == 55:
+            nohands = False
+            nofaces = False
+        elif poses.shape[-3] == 52:
+            nohands = False
+            nofaces = True
+        elif poses.shape[-3] == 22:
+            nohands = True
+            nofaces = True
+        else:
+            raise NotImplementedError("Could not parse the poses.")
+
+        if nohands:
+            left_hand_pose = None
+            right_hand_pose = None
+        else:
+            left_hand_pose = matrix_poses[:, 25:40]
+            right_hand_pose = matrix_poses[:, 40:55]
+
+        if nofaces:
+            jaw_pose = None
+            leye_pose = None
+            reye_pose = None
+        else:
+            jaw_pose = matrix_poses[:, 22:23]
+            leye_pose = matrix_poses[:, 23:24]
+            reye_pose = matrix_poses[:, 24:25]
+
+        n = len(body_pose)
+        outputs = []
+        for chunk in range(int((n - 1) / batch_size) + 1):
+            chunk_slice = slice(chunk * batch_size, (chunk + 1) * batch_size)
+            smpl_output = self.smplx(
+                global_orient=slice_or_none(global_orient, chunk_slice),
+                body_pose=slice_or_none(body_pose, chunk_slice),
+                left_hand_pose=slice_or_none(left_hand_pose, chunk_slice),
+                right_hand_pose=slice_or_none(right_hand_pose, chunk_slice),
+                jaw_pose=slice_or_none(jaw_pose, chunk_slice),
+                leye_pose=slice_or_none(leye_pose, chunk_slice),
+                reye_pose=slice_or_none(reye_pose, chunk_slice),
+                transl=slice_or_none(trans_all, chunk_slice))
+
+            if jointstype == "vertices":
+                output_chunk = smpl_output.vertices
+            else:
+                joints = smpl_output.joints
+                output_chunk = joints
+
+            outputs.append(output_chunk)
+
+        outputs = torch.cat(outputs)
+        outputs = outputs.reshape((*save_shape_bs_len, *outputs.shape[1:]))
+
+        # Change topology if needed
+        outputs = smplx_to(jointstype, outputs, trans)
+
+        return outputs
+
+    def inverse(self, joints: Tensor) -> Tensor:
+        raise NotImplementedError("Cannot inverse SMPLX layer.")
+
+
+def smplx_to(jointstype, data, trans):
+
+    if "mmm" in jointstype:
+        indexes = smplh2mmm_indexes
+        data = data[..., indexes, :]
+
+        # make it compatible with mmm
+        if jointstype == "mmm":
+            data *= smplh_to_mmm_scaling_factor
+
+        if jointstype == "smplmmm":
+            pass
+        elif jointstype in ["mmm", "mmmns"]:
+            # swap axis
+            data = data[..., [1, 2, 0]]
+            # revert left and right
+            data[..., 2] = -data[..., 2]
+
+    elif jointstype == "smplnh":
+        from mGPT.utils.joints import smplh2smplnh_indexes
+        indexes = smplh2smplnh_indexes
+        data = data[..., indexes, :]
+    elif jointstype == "smplh":
+        pass
+    elif jointstype == "vertices":
+        pass
+    else:
+        raise NotImplementedError(f"SMPLX to {jointstype} is not implemented.")
+
+    if jointstype != "vertices":
+        # shift the output in each batch
+        # such that it is centered on the pelvis/root on the first frame
+        root_joint_idx = get_root_idx(jointstype)
+        shift = trans[..., 0, :] - data[..., 0, root_joint_idx, :]
+        data += shift[..., None, None, :]
+
+    return data

mGPT/data/transforms/rots2rfeats/__init__.py

@@ -0,0 +1,5 @@
+from .base import Rots2Rfeats
+# from .globvel import Globalvel
+
+from .globvelandy import Globalvelandy
+# from .rifeats import Rifeats

mGPT/data/transforms/rots2rfeats/base.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+from pathlib import Path
+import os
+
+class Rots2Rfeats(nn.Module):
+    def __init__(self, path: Optional[str] = None,
+                 normalization: bool = True,
+                 eps: float = 1e-12,
+                 **kwargs) -> None:
+        if normalization and path is None:
+            raise TypeError("You should provide a path if normalization is on.")
+
+        super().__init__()
+        self.normalization = normalization
+        self.eps = eps
+        if normalization:
+            # workaround for cluster local/sync
+            rel_p = path.split('/')
+            # superhacky it is for the datatype ugly stuff change it, copy the main stuff to seperate_pairs dict
+            if rel_p[-1] == 'separate_pairs':
+                rel_p.remove('separate_pairs') 
+            ########################################################
+            # rel_p = rel_p[rel_p.index('deps'):]
+            rel_p = '/'.join(rel_p)
+            # path = hydra.utils.get_original_cwd() + '/' + rel_p
+            path = rel_p
+            mean_path = Path(path) / "rfeats_mean.pt"
+            std_path = Path(path) / "rfeats_std.pt"
+
+            self.register_buffer('mean', torch.load(mean_path))
+            self.register_buffer('std', torch.load(std_path))
+
+    def normalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = (features - self.mean)/(self.std + self.eps)
+        return features
+
+    def unnormalize(self, features: Tensor) -> Tensor:
+        if self.normalization:
+            features = features * self.std + self.mean
+        return features

mGPT/data/transforms/rots2rfeats/globvelandy.py

@@ -0,0 +1,128 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+from einops import rearrange
+
+from mGPT.utils.easyconvert import rep_to_rep, nfeats_of, to_matrix
+import mGPT.utils.geometry_tools as geometry_tools
+
+from .base import Rots2Rfeats
+
+
+class Globalvelandy(Rots2Rfeats):
+    def __init__(self,
+                 path: Optional[str] = None,
+                 normalization: bool = False,
+                 pose_rep: str = "rot6d",
+                 canonicalize: bool = False,
+                 offset: bool = True,
+                 **kwargs) -> None:
+        super().__init__(path=path, normalization=normalization)
+
+        self.canonicalize = canonicalize
+        self.pose_rep = pose_rep
+        self.nfeats = nfeats_of(pose_rep)
+        self.offset = offset
+
+    def forward(self, data, data_rep='matrix', first_frame=None) -> Tensor:
+
+        poses, trans = data.rots, data.trans
+
+        # extract the root gravity axis
+        # for smpl it is the last coordinate
+        root_y = trans[..., 2]
+        trajectory = trans[..., [0, 1]]
+
+        # Compute the difference of trajectory
+        vel_trajectory = torch.diff(trajectory, dim=-2)
+
+        # 0 for the first one => keep the dimentionality
+        if first_frame is None:
+            first_frame = 0 * vel_trajectory[..., [0], :]
+
+        vel_trajectory = torch.cat((first_frame, vel_trajectory), dim=-2)
+
+        # first normalize the data
+        if self.canonicalize:
+
+            matrix_poses = rep_to_rep(data_rep, 'matrix', poses)
+            global_orient = matrix_poses[..., 0, :, :]
+
+            # remove the rotation
+            rot2d = rep_to_rep(data_rep, 'rotvec', poses[0, 0, ...])
+            
+            # Remove the fist rotation along the vertical axis
+            rot2d[..., :2] = 0
+
+            if self.offset:
+                # add a bit more rotation
+                rot2d[..., 2] += torch.pi / 2
+
+            rot2d = rep_to_rep('rotvec', 'matrix', rot2d)
+            
+            # turn with the same amount all the rotations
+            global_orient = torch.einsum("...kj,...kl->...jl", rot2d,
+                                         global_orient)
+
+            matrix_poses = torch.cat(
+                (global_orient[..., None, :, :], matrix_poses[..., 1:, :, :]),
+                dim=-3)
+
+            poses = rep_to_rep('matrix', data_rep, matrix_poses)
+
+            # Turn the trajectory as well
+            vel_trajectory = torch.einsum("...kj,...lk->...lj",
+                                          rot2d[..., :2, :2], vel_trajectory)
+
+        poses = rep_to_rep(data_rep, self.pose_rep, poses)
+        features = torch.cat(
+            (root_y[..., None], vel_trajectory,
+             rearrange(poses, "... joints rot -> ... (joints rot)")),
+            dim=-1)
+        features = self.normalize(features)
+
+        return features
+
+    def extract(self, features):
+        root_y = features[..., 0]
+        vel_trajectory = features[..., 1:3]
+        poses_features = features[..., 3:]
+        poses = rearrange(poses_features,
+                          "... (joints rot) -> ... joints rot",
+                          rot=self.nfeats)
+        return root_y, vel_trajectory, poses
+
+    def inverse(self, features, last_frame=None):
+        features = self.unnormalize(features)
+        root_y, vel_trajectory, poses = self.extract(features)
+
+        # integrate the trajectory
+        trajectory = torch.cumsum(vel_trajectory, dim=-2)
+        if last_frame is None:
+            pass
+        # First frame should be 0, but if infered it is better to ensure it
+        trajectory = trajectory - trajectory[..., [0], :]
+
+        # Get back the translation
+        trans = torch.cat([trajectory, root_y[..., None]], dim=-1)
+        matrix_poses = rep_to_rep(self.pose_rep, 'matrix',  poses)
+
+        from ..smpl import RotTransDatastruct
+        return RotTransDatastruct(rots=matrix_poses, trans=trans)

mGPT/data/transforms/smpl.py

@@ -0,0 +1,191 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+from torch import Tensor
+import smplx
+
+from .base import Datastruct, dataclass, Transform
+
+from .rots2rfeats import Rots2Rfeats
+from .rots2joints import Rots2Joints
+from .joints2jfeats import Joints2Jfeats
+
+
+class SMPLTransform(Transform):
+    def __init__(self, rots2rfeats: Rots2Rfeats,
+                 rots2joints: Rots2Joints,
+                 joints2jfeats: Joints2Jfeats,
+                 **kwargs):
+        self.rots2rfeats = rots2rfeats
+        self.rots2joints = rots2joints
+        self.joints2jfeats = joints2jfeats
+
+    def Datastruct(self, **kwargs):
+        return SMPLDatastruct(_rots2rfeats=self.rots2rfeats,
+                              _rots2joints=self.rots2joints,
+                              _joints2jfeats=self.joints2jfeats,
+                              transforms=self,
+                              **kwargs)
+
+    def __repr__(self):
+        return "SMPLTransform()"
+
+
+class RotIdentityTransform(Transform):
+    def __init__(self, **kwargs):
+        return
+
+    def Datastruct(self, **kwargs):
+        return RotTransDatastruct(**kwargs)
+
+    def __repr__(self):
+        return "RotIdentityTransform()"
+
+
+@dataclass
+class RotTransDatastruct(Datastruct):
+    rots: Tensor
+    trans: Tensor
+
+    transforms: RotIdentityTransform = RotIdentityTransform()
+
+    def __post_init__(self):
+        self.datakeys = ["rots", "trans"]
+
+    def __len__(self):
+        return len(self.rots)
+
+
+@dataclass
+class SMPLDatastruct(Datastruct):
+    transforms: SMPLTransform
+    _rots2rfeats: Rots2Rfeats
+    _rots2joints: Rots2Joints
+    _joints2jfeats: Joints2Jfeats
+
+    features: Optional[Tensor] = None
+    rots_: Optional[RotTransDatastruct] = None
+    rfeats_: Optional[Tensor] = None
+    joints_: Optional[Tensor] = None
+    jfeats_: Optional[Tensor] = None
+    vertices_: Optional[Tensor] = None
+
+    def __post_init__(self):
+        self.datakeys = ['features', 'rots_', 'rfeats_',
+                         'joints_', 'jfeats_', 'vertices_']
+        # starting point
+        if self.features is not None and self.rfeats_ is None:
+            self.rfeats_ = self.features
+
+    @property
+    def rots(self):
+        # Cached value
+        if self.rots_ is not None:
+            return self.rots_
+
+        # self.rfeats_ should be defined
+        assert self.rfeats_ is not None
+
+        self._rots2rfeats.to(self.rfeats.device)
+        self.rots_ = self._rots2rfeats.inverse(self.rfeats)
+        return self.rots_
+
+    @property
+    def rfeats(self):
+        # Cached value
+        if self.rfeats_ is not None:
+            return self.rfeats_
+
+        # self.rots_ should be defined
+        assert self.rots_ is not None
+
+        self._rots2rfeats.to(self.rots.device)
+        self.rfeats_ = self._rots2rfeats(self.rots)
+        return self.rfeats_
+
+    @property
+    def joints(self):
+        # Cached value
+        if self.joints_ is not None:
+            return self.joints_
+
+        self._rots2joints.to(self.rots.device)
+        self.joints_ = self._rots2joints(self.rots)
+        return self.joints_
+
+    @property
+    def jfeats(self):
+        # Cached value
+        if self.jfeats_ is not None:
+            return self.jfeats_
+
+        self._joints2jfeats.to(self.joints.device)
+        self.jfeats_ = self._joints2jfeats(self.joints)
+        return self.jfeats_
+    
+    @property
+    def vertices(self):
+        # Cached value
+        if self.vertices_ is not None:
+            return self.vertices_
+
+        self._rots2joints.to(self.rots.device)
+        self.vertices_ = self._rots2joints(self.rots, jointstype="vertices")
+        return self.vertices_
+    
+    def __len__(self):
+        return len(self.rfeats)
+
+
+def get_body_model(model_type, gender, batch_size, device='cpu', ext='pkl'):
+    '''
+    type: smpl, smplx smplh and others. Refer to smplx tutorial
+    gender: male, female, neutral
+    batch_size: an positive integar
+    '''
+    mtype = model_type.upper()
+    if gender != 'neutral':
+        if not isinstance(gender, str):
+            gender = str(gender.astype(str)).upper()
+        else:
+            gender = gender.upper()
+    else:
+        gender = gender.upper()
+        ext = 'npz'
+    body_model_path = f'data/smpl_models/{model_type}/{mtype}_{gender}.{ext}'
+
+    body_model = smplx.create(body_model_path, model_type=type,
+                              gender=gender, ext=ext,
+                              use_pca=False,
+                              num_pca_comps=12,
+                              create_global_orient=True,
+                              create_body_pose=True,
+                              create_betas=True,
+                              create_left_hand_pose=True,
+                              create_right_hand_pose=True,
+                              create_expression=True,
+                              create_jaw_pose=True,
+                              create_leye_pose=True,
+                              create_reye_pose=True,
+                              create_transl=True,
+                              batch_size=batch_size)
+    
+    if device == 'cuda':
+        return body_model.cuda()
+    else:
+        return body_model
+

mGPT/data/transforms/xyz.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional
+from torch import Tensor
+
+from .base import Datastruct, dataclass, Transform
+from ..tools import collate_tensor_with_padding
+
+from .joints2jfeats import Joints2Jfeats
+
+
+class XYZTransform(Transform):
+    def __init__(self, joints2jfeats: Joints2Jfeats, **kwargs):
+        self.joints2jfeats = joints2jfeats
+
+    def Datastruct(self, **kwargs):
+        return XYZDatastruct(_joints2jfeats=self.joints2jfeats,
+                             transforms=self,
+                             **kwargs)
+
+    def __repr__(self):
+        return "XYZTransform()"
+
+
+@dataclass
+class XYZDatastruct(Datastruct):
+    transforms: XYZTransform
+    _joints2jfeats: Joints2Jfeats
+
+    features: Optional[Tensor] = None
+    joints_: Optional[Tensor] = None
+    jfeats_: Optional[Tensor] = None
+
+    def __post_init__(self):
+        self.datakeys = ["features", "joints_", "jfeats_"]
+        # starting point
+        if self.features is not None and self.jfeats_ is None:
+            self.jfeats_ = self.features
+
+    @property
+    def joints(self):
+        # Cached value
+        if self.joints_ is not None:
+            return self.joints_
+
+        # self.jfeats_ should be defined
+        assert self.jfeats_ is not None
+
+        self._joints2jfeats.to(self.jfeats.device)
+        self.joints_ = self._joints2jfeats.inverse(self.jfeats)
+        return self.joints_
+
+    @property
+    def jfeats(self):
+        # Cached value
+        if self.jfeats_ is not None:
+            return self.jfeats_
+
+        # self.joints_ should be defined
+        assert self.joints_ is not None
+
+        self._joints2jfeats.to(self.joints.device)
+        self.jfeats_ = self._joints2jfeats(self.joints)
+        return self.jfeats_
+
+    def __len__(self):
+        return len(self.jfeats)

mGPT/data/utils.py

@@ -0,0 +1,81 @@
+import torch
+import rich
+import pickle
+import numpy as np
+
+
+def lengths_to_mask(lengths):
+    max_len = max(lengths)
+    mask = torch.arange(max_len, device=lengths.device).expand(
+        len(lengths), max_len) < lengths.unsqueeze(1)
+    return mask
+
+
+# padding to max length in one batch
+def collate_tensors(batch):
+    if isinstance(batch[0], np.ndarray):
+        batch = [torch.tensor(b).float() for b in batch]
+
+    dims = batch[0].dim()
+    max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
+    size = (len(batch), ) + tuple(max_size)
+    canvas = batch[0].new_zeros(size=size)
+    for i, b in enumerate(batch):
+        sub_tensor = canvas[i]
+        for d in range(dims):
+            sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
+        sub_tensor.add_(b)
+    return canvas
+
+def humanml3d_collate(batch):
+    notnone_batches = [b for b in batch if b is not None]
+    EvalFlag = False if notnone_batches[0][5] is None else True
+
+    # Sort by text length
+    if EvalFlag:
+        notnone_batches.sort(key=lambda x: x[5], reverse=True)
+
+    # Motion only
+    adapted_batch = {
+        "motion":
+        collate_tensors([torch.tensor(b[1]).float() for b in notnone_batches]),
+        "length": [b[2] for b in notnone_batches],
+    }
+
+    # Text and motion
+    if notnone_batches[0][0] is not None:
+        adapted_batch.update({
+            "text": [b[0] for b in notnone_batches],
+            "all_captions": [b[7] for b in notnone_batches],
+        })
+
+    # Evaluation related
+    if EvalFlag:
+        adapted_batch.update({
+            "text": [b[0] for b in notnone_batches],
+            "word_embs":
+            collate_tensors(
+                [torch.tensor(b[3]).float() for b in notnone_batches]),
+            "pos_ohot":
+            collate_tensors(
+                [torch.tensor(b[4]).float() for b in notnone_batches]),
+            "text_len":
+            collate_tensors([torch.tensor(b[5]) for b in notnone_batches]),
+            "tokens": [b[6] for b in notnone_batches],
+        })
+
+    # Tasks
+    if len(notnone_batches[0]) == 9:
+        adapted_batch.update({"tasks": [b[8] for b in notnone_batches]})
+
+    return adapted_batch
+
+
+def load_pkl(path, description=None, progressBar=False):
+    if progressBar:
+        with rich.progress.open(path, 'rb', description=description) as file:
+            data = pickle.load(file)
+    else:
+        with open(path, 'rb') as file:
+            data = pickle.load(file)
+    return data