From e2c67e5abe528ef23f3f21cf020875bf5d7038f4 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 29 Mar 2024 18:38:28 +0200 Subject: [PATCH] include data folder --- .gitignore | 2 +- mGPT/data/HumanML3D.py | 119 ++++ mGPT/data/Kit.py | 88 +++ mGPT/data/__init__.py | 103 ++++ mGPT/data/build_data.py | 15 + mGPT/data/humanml/README.md | 1 + mGPT/data/humanml/__init__.py | 7 + mGPT/data/humanml/common/quaternion.py | 423 +++++++++++++ mGPT/data/humanml/common/skeleton.py | 199 ++++++ mGPT/data/humanml/dataset_m.py | 156 +++++ mGPT/data/humanml/dataset_m_vq.py | 54 ++ mGPT/data/humanml/dataset_t2m.py | 211 +++++++ mGPT/data/humanml/dataset_t2m_cb.py | 211 +++++++ mGPT/data/humanml/dataset_t2m_eval.py | 92 +++ mGPT/data/humanml/dataset_t2m_m2t.py | 119 ++++ mGPT/data/humanml/dataset_t2m_token.py | 86 +++ mGPT/data/humanml/scripts/motion_process.py | 529 ++++++++++++++++ mGPT/data/humanml/utils/paramUtil.py | 63 ++ mGPT/data/humanml/utils/word_vectorizer.py | 79 +++ mGPT/data/mean.npy | Bin 0 -> 2232 bytes mGPT/data/mean_eval.npy | Bin 0 -> 2232 bytes mGPT/data/std.npy | Bin 0 -> 2232 bytes mGPT/data/std_eval.npy | Bin 0 -> 2232 bytes mGPT/data/tools/__init__.py | 2 + mGPT/data/tools/collate.py | 99 +++ mGPT/data/tools/easyconvert.py | 72 +++ mGPT/data/tools/geometry.py | 566 ++++++++++++++++++ mGPT/data/tools/tensors.py | 26 + mGPT/data/transforms/__init__.py | 15 + mGPT/data/transforms/base.py | 84 +++ mGPT/data/transforms/identity.py | 44 ++ .../data/transforms/joints2jfeats/__init__.py | 2 + mGPT/data/transforms/joints2jfeats/base.py | 59 ++ mGPT/data/transforms/joints2jfeats/rifke.py | 159 +++++ mGPT/data/transforms/joints2jfeats/tools.py | 97 +++ mGPT/data/transforms/joints2rots/config.py | 119 ++++ .../data/transforms/joints2rots/customloss.py | 217 +++++++ mGPT/data/transforms/joints2rots/prior.py | 229 +++++++ mGPT/data/transforms/joints2rots/smplify.py | 284 +++++++++ mGPT/data/transforms/rots2joints/__init__.py | 3 + mGPT/data/transforms/rots2joints/base.py | 56 ++ mGPT/data/transforms/rots2joints/smplh.py | 192 ++++++ mGPT/data/transforms/rots2joints/smplx.py | 201 +++++++ mGPT/data/transforms/rots2rfeats/__init__.py | 5 + mGPT/data/transforms/rots2rfeats/base.py | 60 ++ .../transforms/rots2rfeats/globvelandy.py | 128 ++++ mGPT/data/transforms/smpl.py | 191 ++++++ mGPT/data/transforms/xyz.py | 81 +++ mGPT/data/utils.py | 81 +++ 49 files changed, 5628 insertions(+), 1 deletion(-) create mode 100644 mGPT/data/HumanML3D.py create mode 100644 mGPT/data/Kit.py create mode 100644 mGPT/data/__init__.py create mode 100644 mGPT/data/build_data.py create mode 100644 mGPT/data/humanml/README.md create mode 100644 mGPT/data/humanml/__init__.py create mode 100644 mGPT/data/humanml/common/quaternion.py create mode 100644 mGPT/data/humanml/common/skeleton.py create mode 100644 mGPT/data/humanml/dataset_m.py create mode 100644 mGPT/data/humanml/dataset_m_vq.py create mode 100644 mGPT/data/humanml/dataset_t2m.py create mode 100644 mGPT/data/humanml/dataset_t2m_cb.py create mode 100644 mGPT/data/humanml/dataset_t2m_eval.py create mode 100644 mGPT/data/humanml/dataset_t2m_m2t.py create mode 100644 mGPT/data/humanml/dataset_t2m_token.py create mode 100644 mGPT/data/humanml/scripts/motion_process.py create mode 100644 mGPT/data/humanml/utils/paramUtil.py create mode 100644 mGPT/data/humanml/utils/word_vectorizer.py create mode 100644 mGPT/data/mean.npy create mode 100644 mGPT/data/mean_eval.npy create mode 100644 mGPT/data/std.npy create mode 100644 mGPT/data/std_eval.npy create mode 100644 mGPT/data/tools/__init__.py create mode 100644 mGPT/data/tools/collate.py create mode 100644 mGPT/data/tools/easyconvert.py create mode 100644 mGPT/data/tools/geometry.py create mode 100644 mGPT/data/tools/tensors.py create mode 100644 mGPT/data/transforms/__init__.py create mode 100644 mGPT/data/transforms/base.py create mode 100644 mGPT/data/transforms/identity.py create mode 100644 mGPT/data/transforms/joints2jfeats/__init__.py create mode 100644 mGPT/data/transforms/joints2jfeats/base.py create mode 100644 mGPT/data/transforms/joints2jfeats/rifke.py create mode 100644 mGPT/data/transforms/joints2jfeats/tools.py create mode 100644 mGPT/data/transforms/joints2rots/config.py create mode 100644 mGPT/data/transforms/joints2rots/customloss.py create mode 100644 mGPT/data/transforms/joints2rots/prior.py create mode 100644 mGPT/data/transforms/joints2rots/smplify.py create mode 100644 mGPT/data/transforms/rots2joints/__init__.py create mode 100644 mGPT/data/transforms/rots2joints/base.py create mode 100644 mGPT/data/transforms/rots2joints/smplh.py create mode 100644 mGPT/data/transforms/rots2joints/smplx.py create mode 100644 mGPT/data/transforms/rots2rfeats/__init__.py create mode 100644 mGPT/data/transforms/rots2rfeats/base.py create mode 100644 mGPT/data/transforms/rots2rfeats/globvelandy.py create mode 100644 mGPT/data/transforms/smpl.py create mode 100644 mGPT/data/transforms/xyz.py create mode 100644 mGPT/data/utils.py diff --git a/.gitignore b/.gitignore index dddf596..2a6d887 100644 --- a/.gitignore +++ b/.gitignore @@ -105,7 +105,7 @@ venv.bak/ .mypy_cache/ # custom -data +# data # data for pytest moved to http server # !tests/data .vscode diff --git a/mGPT/data/HumanML3D.py b/mGPT/data/HumanML3D.py new file mode 100644 index 0000000..8bec575 --- /dev/null +++ b/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 diff --git a/mGPT/data/Kit.py b/mGPT/data/Kit.py new file mode 100644 index 0000000..1eecaa1 --- /dev/null +++ b/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 diff --git a/mGPT/data/__init__.py b/mGPT/data/__init__.py new file mode 100644 index 0000000..948f49c --- /dev/null +++ b/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, + ) diff --git a/mGPT/data/build_data.py b/mGPT/data/build_data.py new file mode 100644 index 0000000..88b493e --- /dev/null +++ b/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) diff --git a/mGPT/data/humanml/README.md b/mGPT/data/humanml/README.md new file mode 100644 index 0000000..4bf224f --- /dev/null +++ b/mGPT/data/humanml/README.md @@ -0,0 +1 @@ +This code is based on https://github.com/EricGuo5513/text-to-motion.git \ No newline at end of file diff --git a/mGPT/data/humanml/__init__.py b/mGPT/data/humanml/__init__.py new file mode 100644 index 0000000..8e75358 --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/common/quaternion.py b/mGPT/data/humanml/common/quaternion.py new file mode 100644 index 0000000..dca3d89 --- /dev/null +++ b/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) diff --git a/mGPT/data/humanml/common/skeleton.py b/mGPT/data/humanml/common/skeleton.py new file mode 100644 index 0000000..b2ae85a --- /dev/null +++ b/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 + + + + + diff --git a/mGPT/data/humanml/dataset_m.py b/mGPT/data/humanml/dataset_m.py new file mode 100644 index 0000000..241cabc --- /dev/null +++ b/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, diff --git a/mGPT/data/humanml/dataset_m_vq.py b/mGPT/data/humanml/dataset_m_vq.py new file mode 100644 index 0000000..dfa3ae5 --- /dev/null +++ b/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, diff --git a/mGPT/data/humanml/dataset_t2m.py b/mGPT/data/humanml/dataset_t2m.py new file mode 100644 index 0000000..3460681 --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/dataset_t2m_cb.py b/mGPT/data/humanml/dataset_t2m_cb.py new file mode 100644 index 0000000..6f7e1fe --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/dataset_t2m_eval.py b/mGPT/data/humanml/dataset_t2m_eval.py new file mode 100644 index 0000000..4162981 --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/dataset_t2m_m2t.py b/mGPT/data/humanml/dataset_t2m_m2t.py new file mode 100644 index 0000000..259078c --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/dataset_t2m_token.py b/mGPT/data/humanml/dataset_t2m_token.py new file mode 100644 index 0000000..d9f1e54 --- /dev/null +++ b/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 diff --git a/mGPT/data/humanml/scripts/motion_process.py b/mGPT/data/humanml/scripts/motion_process.py new file mode 100644 index 0000000..8b3395c --- /dev/null +++ b/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)) diff --git a/mGPT/data/humanml/utils/paramUtil.py b/mGPT/data/humanml/utils/paramUtil.py new file mode 100644 index 0000000..a9f1708 --- /dev/null +++ b/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' + diff --git a/mGPT/data/humanml/utils/word_vectorizer.py b/mGPT/data/humanml/utils/word_vectorizer.py new file mode 100644 index 0000000..d272058 --- /dev/null +++ b/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 diff --git a/mGPT/data/mean.npy b/mGPT/data/mean.npy new file mode 100644 index 0000000000000000000000000000000000000000..e754cba2fee3e20e6b0b11fe8d3ac24ae7b0214c GIT binary patch literal 2232 zcmbW!`9IYA7r=4S$Xbe9k)npCkjhSRNV!>(EXi7AmuQ3{%b<`{8iS-l))Bh+R)#+B zP*M$M$dY|B*~StQ6Y1*y3E%VEc|D##okKRaIbmVP%M;8KMsRZvJWnL($Px~FXc3fT z2_F8$%S4xpXZ?w8?ti+G%Y{Jq)gI{O65zhN?@`fERobg08!r2QkC21jC^lNziETRI zlwFeRCjC@yEQ++AA zVlc1KMHvU_+f!jH;ZbPFu6vskbqg|n5S4f+>l>tRYwc zZm_;Lq44vEDzFH8mA~Xj1H9$$K9li$q)+8E+&xc27*l~FWLXUhYbGqdy~;=4OZ&2q zwW^`D{W^lDFV2Cm@~2qFlT@g($})Z1RR)(T%Bx3G^C1YND(`DFCNEJy&*R_yg_vO=a?;MO!x;R653u6Hm#b~6;_s=0V zslVbzd?Vbj@rhE#&qItvH-m4H04l<-1|OV^kO1MP{}Rv!v%t;vY#)@A@( zGe^pcH3MR=^iFC$IR`y2{dE`l7^pJu5t}wR1L6L23l3k*Au}JhMM#bb)I76-L)WK) zV>7UDt(FE)51N)teq*4IJCjGYU7m(2H_4UF?`U8Wu`EF(eunm+8_lm6jG=yiub!HR5jvg-yx+@~bW{#zFEvv?)=M})IXGBqImp$4*lT9X=%0nb;Wj2mH z33#83zO=jVPOB-ZrJX>$# zT8~D_LY#ouHhA^V*b4(44q7n%-M`x*1to^g{_)_~b&DX-)&`gMDQHew)FY{HeJ)Y? zQ&5rJu}Rw65D5?+O>{0uBLm}H{+=L1sI9HCQhhxNqNHI1%7IQurqOAYWs_iQKGd7~ zrXDqD-`p4EK|$fy56H#W)qzV%-JnVED6lfrHDyHT@GPf*weV~bR;>NsJC@XgI8Gav zX-$EpJ+X~hMRn-$TIzqEe7pHejU*k$k|Sg;g-wG*dNWlwMinXb;O&m?l|Vcq)N2p> zRUtQ_(XVfL65@U3I45n(!B(2oG(jB!-Uj9AG@bikT=$^hm@*Sp6DXZzg*-@5QVIA% z7zZ4Ueo=Dn0|;u0c34!5fS0PL&s(o_bYETd6~7P@?9^rj{yCV3sK+E!F!^zC)4BWl z+0+LhehoTIvz>-;lH2>YQwTx!~mX&N~6>$_0-+trX*0Y_^ ztaJH~N9LrIUoR>$;L&nyZ*%|)wx&5bCHt8NTZ(MyKaRS%_4m*}3bd8yc`U>y9Nxqsy&pH?DeE z&?Q)zd$Qp>qb0u2_X(N~7%GmScEeE>dbwaC8;v?W?VSsbhwLW;Un8}LfTus+NKLsN z{`avzAxnn}_FnoTkIK5if644`ERGF4Jb(F~IeP=`sL3?V&>TXfYdGUSKJIwVkw!u> z=q45Nr&y0dn^!uqC9(?ji(X(((2LLp=U)5Q57I%J5&qI@0}Etse(E``)`i}LMvw5F z8igSC$(v?xs$lVEW(wzB5r}2~kg3+pKuu#;|MRZ5lXg1q??R4t{&Mh>5x#!iyQsU}VPC48tK3u;%_MFAe#L6)5pB7d2tz;i_E)Q6>pfoD7pblq4MGc8 z<%2t$P->vny|P6-Mf(&nI!mJwk}4_QVb5( zoF7!lSOtc- z7~b7b)wG2KA~9AvPj5?rrztyMtbi8^-HIjM?W3Wo;<-AzPZnwPD5_ETcwt9-j?nJa zUNf2#oOX|kH5OY*EHv|?s1VgaMLG$R?icol$FD&SCSVoe&4X-DtLJZOCIR0Zy2V_> z!8#scQDs+o(018CEt%EcOerb})#PHX__)J|)7BvG?v3Lo@=5Ta-i=yH6+_|;(sIJj zeqlUse=h~LlHh>c25HaZ6M8;F$MNF06|8o8tGE4XxAq-*EM+W>T%_+0uf4jAotj!q z4V@xES*U1&@$`E5+Kw9zZCJ+a-sc_`pC*CcYYUCF2|JNzWM=aI?JO)>T`^>8kp%n^ zmO&?sH^XZwE%qXwg$czd&>373%sRYmID+RuWV=Y9d?E*gzh|0u5jeni)^K}Q5(oYT DDM3hl literal 0 HcmV?d00001 diff --git a/mGPT/data/mean_eval.npy b/mGPT/data/mean_eval.npy new file mode 100644 index 0000000000000000000000000000000000000000..e754cba2fee3e20e6b0b11fe8d3ac24ae7b0214c GIT binary patch literal 2232 zcmbW!`9IYA7r=4S$Xbe9k)npCkjhSRNV!>(EXi7AmuQ3{%b<`{8iS-l))Bh+R)#+B zP*M$M$dY|B*~StQ6Y1*y3E%VEc|D##okKRaIbmVP%M;8KMsRZvJWnL($Px~FXc3fT z2_F8$%S4xpXZ?w8?ti+G%Y{Jq)gI{O65zhN?@`fERobg08!r2QkC21jC^lNziETRI zlwFeRCjC@yEQ++AA zVlc1KMHvU_+f!jH;ZbPFu6vskbqg|n5S4f+>l>tRYwc zZm_;Lq44vEDzFH8mA~Xj1H9$$K9li$q)+8E+&xc27*l~FWLXUhYbGqdy~;=4OZ&2q zwW^`D{W^lDFV2Cm@~2qFlT@g($})Z1RR)(T%Bx3G^C1YND(`DFCNEJy&*R_yg_vO=a?;MO!x;R653u6Hm#b~6;_s=0V zslVbzd?Vbj@rhE#&qItvH-m4H04l<-1|OV^kO1MP{}Rv!v%t;vY#)@A@( zGe^pcH3MR=^iFC$IR`y2{dE`l7^pJu5t}wR1L6L23l3k*Au}JhMM#bb)I76-L)WK) zV>7UDt(FE)51N)teq*4IJCjGYU7m(2H_4UF?`U8Wu`EF(eunm+8_lm6jG=yiub!HR5jvg-yx+@~bW{#zFEvv?)=M})IXGBqImp$4*lT9X=%0nb;Wj2mH z33#83zO=jVPOB-ZrJX>$# zT8~D_LY#ouHhA^V*b4(44q7n%-M`x*1to^g{_)_~b&DX-)&`gMDQHew)FY{HeJ)Y? zQ&5rJu}Rw65D5?+O>{0uBLm}H{+=L1sI9HCQhhxNqNHI1%7IQurqOAYWs_iQKGd7~ zrXDqD-`p4EK|$fy56H#W)qzV%-JnVED6lfrHDyHT@GPf*weV~bR;>NsJC@XgI8Gav zX-$EpJ+X~hMRn-$TIzqEe7pHejU*k$k|Sg;g-wG*dNWlwMinXb;O&m?l|Vcq)N2p> zRUtQ_(XVfL65@U3I45n(!B(2oG(jB!-Uj9AG@bikT=$^hm@*Sp6DXZzg*-@5QVIA% z7zZ4Ueo=Dn0|;u0c34!5fS0PL&s(o_bYETd6~7P@?9^rj{yCV3sK+E!F!^zC)4BWl z+0+LhehoTIvz>-;lH2>YQwTx!~mX&N~6>$_0-+trX*0Y_^ ztaJH~N9LrIUoR>$;L&nyZ*%|)wx&5bCHt8NTZ(MyKaRS%_4m*}3bd8yc`U>y9Nxqsy&pH?DeE z&?Q)zd$Qp>qb0u2_X(N~7%GmScEeE>dbwaC8;v?W?VSsbhwLW;Un8}LfTus+NKLsN z{`avzAxnn}_FnoTkIK5if644`ERGF4Jb(F~IeP=`sL3?V&>TXfYdGUSKJIwVkw!u> z=q45Nr&y0dn^!uqC9(?ji(X(((2LLp=U)5Q57I%J5&qI@0}Etse(E``)`i}LMvw5F z8igSC$(v?xs$lVEW(wzB5r}2~kg3+pKuu#;|MRZ5lXg1q??R4t{&Mh>5x#!iyQsU}VPC48tK3u;%_MFAe#L6)5pB7d2tz;i_E)Q6>pfoD7pblq4MGc8 z<%2t$P->vny|P6-Mf(&nI!mJwk}4_QVb5( zoF7!lSOtc- z7~b7b)wG2KA~9AvPj5?rrztyMtbi8^-HIjM?W3Wo;<-AzPZnwPD5_ETcwt9-j?nJa zUNf2#oOX|kH5OY*EHv|?s1VgaMLG$R?icol$FD&SCSVoe&4X-DtLJZOCIR0Zy2V_> z!8#scQDs+o(018CEt%EcOerb})#PHX__)J|)7BvG?v3Lo@=5Ta-i=yH6+_|;(sIJj zeqlUse=h~LlHh>c25HaZ6M8;F$MNF06|8o8tGE4XxAq-*EM+W>T%_+0uf4jAotj!q z4V@xES*U1&@$`E5+Kw9zZCJ+a-sc_`pC*CcYYUCF2|JNzWM=aI?JO)>T`^>8kp%n^ zmO&?sH^XZwE%qXwg$czd&>373%sRYmID+RuWV=Y9d?E*gzh|0u5jeni)^K}Q5(oYT DDM3hl literal 0 HcmV?d00001 diff --git a/mGPT/data/std.npy b/mGPT/data/std.npy new file mode 100644 index 0000000000000000000000000000000000000000..c0ae831306536a16da5c499325588eb521fd3255 GIT binary patch literal 2232 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$7X2v?2ItsN4WCJdt@VQ%3?mF1JEOYOy(Q|{-nFfagRfH<-{YqQ4PQ~vZ zrAJsel-HfL@36dVKT3~=!)Q7f{`r6*)~K?OX^Q 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 diff --git a/mGPT/data/tools/easyconvert.py b/mGPT/data/tools/easyconvert.py new file mode 100644 index 0000000..3649a93 --- /dev/null +++ b/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 diff --git a/mGPT/data/tools/geometry.py b/mGPT/data/tools/geometry.py new file mode 100644 index 0000000..e6eafa2 --- /dev/null +++ b/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) diff --git a/mGPT/data/tools/tensors.py b/mGPT/data/tools/tensors.py new file mode 100644 index 0000000..6bcc051 --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/__init__.py b/mGPT/data/transforms/__init__.py new file mode 100644 index 0000000..394dfc5 --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/base.py b/mGPT/data/transforms/base.py new file mode 100644 index 0000000..1c60a60 --- /dev/null +++ b/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) diff --git a/mGPT/data/transforms/identity.py b/mGPT/data/transforms/identity.py new file mode 100644 index 0000000..ec12e7f --- /dev/null +++ b/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) diff --git a/mGPT/data/transforms/joints2jfeats/__init__.py b/mGPT/data/transforms/joints2jfeats/__init__.py new file mode 100644 index 0000000..0a924e8 --- /dev/null +++ b/mGPT/data/transforms/joints2jfeats/__init__.py @@ -0,0 +1,2 @@ +from .base import Joints2Jfeats +from .rifke import Rifke diff --git a/mGPT/data/transforms/joints2jfeats/base.py b/mGPT/data/transforms/joints2jfeats/base.py new file mode 100644 index 0000000..03d6f5f --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/joints2jfeats/rifke.py b/mGPT/data/transforms/joints2jfeats/rifke.py new file mode 100644 index 0000000..c6f2a8e --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/joints2jfeats/tools.py b/mGPT/data/transforms/joints2jfeats/tools.py new file mode 100644 index 0000000..734e109 --- /dev/null +++ b/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) diff --git a/mGPT/data/transforms/joints2rots/config.py b/mGPT/data/transforms/joints2rots/config.py new file mode 100644 index 0000000..9014bef --- /dev/null +++ b/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" diff --git a/mGPT/data/transforms/joints2rots/customloss.py b/mGPT/data/transforms/joints2rots/customloss.py new file mode 100644 index 0000000..2c3c3a5 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/mGPT/data/transforms/joints2rots/prior.py b/mGPT/data/transforms/joints2rots/prior.py new file mode 100644 index 0000000..d85debd --- /dev/null +++ b/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) diff --git a/mGPT/data/transforms/joints2rots/smplify.py b/mGPT/data/transforms/joints2rots/smplify.py new file mode 100644 index 0000000..7df5150 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/mGPT/data/transforms/rots2joints/__init__.py b/mGPT/data/transforms/rots2joints/__init__.py new file mode 100644 index 0000000..7719c70 --- /dev/null +++ b/mGPT/data/transforms/rots2joints/__init__.py @@ -0,0 +1,3 @@ +from .base import Rots2Joints +from .smplh import SMPLH +from .smplx import SMPLX diff --git a/mGPT/data/transforms/rots2joints/base.py b/mGPT/data/transforms/rots2joints/base.py new file mode 100644 index 0000000..524f830 --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/rots2joints/smplh.py b/mGPT/data/transforms/rots2joints/smplh.py new file mode 100644 index 0000000..90efa4f --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/rots2joints/smplx.py b/mGPT/data/transforms/rots2joints/smplx.py new file mode 100644 index 0000000..107eb57 --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/rots2rfeats/__init__.py b/mGPT/data/transforms/rots2rfeats/__init__.py new file mode 100644 index 0000000..29b206c --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/rots2rfeats/base.py b/mGPT/data/transforms/rots2rfeats/base.py new file mode 100644 index 0000000..98c33bd --- /dev/null +++ b/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 diff --git a/mGPT/data/transforms/rots2rfeats/globvelandy.py b/mGPT/data/transforms/rots2rfeats/globvelandy.py new file mode 100644 index 0000000..fe223af --- /dev/null +++ b/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) diff --git a/mGPT/data/transforms/smpl.py b/mGPT/data/transforms/smpl.py new file mode 100644 index 0000000..fc46b11 --- /dev/null +++ b/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 + diff --git a/mGPT/data/transforms/xyz.py b/mGPT/data/transforms/xyz.py new file mode 100644 index 0000000..7add165 --- /dev/null +++ b/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) diff --git a/mGPT/data/utils.py b/mGPT/data/utils.py new file mode 100644 index 0000000..30714ff --- /dev/null +++ b/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