├── teach ├── __init__.py ├── data │ ├── __init__.py │ ├── sampling │ │ ├── __init__.py │ │ ├── framerate.py │ │ ├── base.py │ │ └── frames.py │ ├── tools │ │ ├── __init__.py │ │ ├── tensors.py │ │ └── smpl.py │ ├── text2motion.py │ └── base.py ├── model │ ├── __init__.py │ ├── motionencoder │ │ ├── __init__.py │ │ ├── gru.py │ │ └── actor.py │ ├── utils │ │ ├── __init__.py │ │ ├── tools.py │ │ ├── vae.py │ │ └── positional_encoding.py │ ├── metrics │ │ └── __init__.py │ ├── motiondecoder │ │ ├── __init__.py │ │ ├── gru.py │ │ ├── actor.py │ │ └── metactor.py │ ├── losses │ │ ├── __init__.py │ │ ├── recons.py │ │ ├── kl.py │ │ └── actioncompute.py │ ├── textencoder │ │ ├── __init__.py │ │ ├── distilbert_linear.py │ │ └── distilbert.py │ ├── dummy.py │ └── base.py ├── tools │ ├── __init__.py │ ├── runid.py │ ├── logging.py │ └── easyconvert.py ├── utils │ └── __init__.py ├── render │ ├── blender │ │ ├── __init__.py │ │ ├── data.py │ │ ├── vertices.py │ │ ├── tools.py │ │ ├── sampler.py │ │ ├── camera.py │ │ ├── materials.py │ │ ├── scene.py │ │ ├── floor.py │ │ └── meshes.py │ └── __init__.py ├── transforms │ ├── rots2joints │ │ ├── __init__.py │ │ └── base.py │ ├── joints2jfeats │ │ ├── __init__.py │ │ ├── base.py │ │ └── tools.py │ ├── __init__.py │ ├── rots2rfeats │ │ ├── __init__.py │ │ └── base.py │ ├── identity.py │ ├── xyz.py │ └── base.py ├── callback │ ├── __init__.py │ └── progress.py ├── launch │ ├── blender.py │ └── prepare.py └── logger │ ├── __init__.py │ ├── wandb_log.py │ └── tools.py ├── data_loaders ├── humanml │ ├── data │ │ └── __init__.py │ ├── networks │ │ └── __init__.py │ ├── motion_loaders │ │ ├── __init__.py │ │ └── dataset_motion_loader.py │ ├── README.md │ └── utils │ │ ├── paramUtil.py │ │ ├── get_opt.py │ │ └── word_vectorizer.py ├── a2m │ └── humanact12poses.py ├── humanml_utils.py ├── tensors.py └── get_data.py ├── transforms ├── rots2joints │ ├── __init__.py │ └── base.py ├── joints2jfeats │ ├── __init__.py │ ├── base.py │ └── tools.py ├── __init__.py ├── rots2rfeats │ ├── __init__.py │ └── base.py ├── identity.py ├── xyz.py └── base.py ├── utils ├── fixseed.py ├── config.py ├── misc.py ├── PYTORCH3D_LICENSE ├── dist_util.py ├── seq_process.py ├── smpl.py └── model_util.py ├── model ├── cfg_sampler.py ├── text_encoder.py ├── rotation2xyz.py ├── motion_encoder.py ├── motion_clip.py └── smpl.py ├── train ├── train_platforms.py └── train_mdm.py ├── train_diffusion.py ├── README.md ├── .gitignore ├── train_contrastive.py ├── diffusion └── losses.py └── motionclip_data_loaders └── pretrain_datamodule.py /teach/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teach/data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teach/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teach/tools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teach/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data_loaders/humanml/data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data_loaders/humanml/networks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data_loaders/humanml/motion_loaders/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teach/render/blender/__init__.py: -------------------------------------------------------------------------------- 1 | from .render import render -------------------------------------------------------------------------------- /transforms/rots2joints/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Rots2Joints 2 | from .smplh import SMPLH 3 | -------------------------------------------------------------------------------- /data_loaders/humanml/README.md: -------------------------------------------------------------------------------- 1 | This code is based on https://github.com/EricGuo5513/text-to-motion.git -------------------------------------------------------------------------------- /teach/transforms/rots2joints/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Rots2Joints 2 | from .smplh import SMPLH 3 | -------------------------------------------------------------------------------- /transforms/joints2jfeats/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Joints2Jfeats 2 | from .rifke import Rifke 3 | -------------------------------------------------------------------------------- /teach/callback/__init__.py: -------------------------------------------------------------------------------- 1 | from .render import RenderCallback 2 | from .progress import ProgressLogger 3 | -------------------------------------------------------------------------------- /teach/transforms/joints2jfeats/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Joints2Jfeats 2 | from .rifke import Rifke 3 | -------------------------------------------------------------------------------- /teach/data/sampling/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import FrameSampler 2 | from .framerate import subsample, upsample 3 | -------------------------------------------------------------------------------- /teach/model/motionencoder/__init__.py: -------------------------------------------------------------------------------- 1 | from .actor import ActorAgnosticEncoder 2 | from .gru import GRUEncoder 3 | -------------------------------------------------------------------------------- /teach/model/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .positional_encoding import PositionalEncoding 2 | from .vae import reparameterize 3 | -------------------------------------------------------------------------------- /transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Transform 2 | from .smpl import SMPLTransform 3 | from .xyz import XYZTransform 4 | -------------------------------------------------------------------------------- /teach/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Transform 2 | from .smpl import SMPLTransform 3 | from .xyz import XYZTransform 4 | -------------------------------------------------------------------------------- /teach/render/__init__.py: -------------------------------------------------------------------------------- 1 | from .anim import render_animation 2 | import sys 3 | if 'blender' not in sys.executable: 4 | from .mesh_viz import visualize_meshes 5 | -------------------------------------------------------------------------------- /teach/model/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .compute import ComputeMetrics 2 | from .compute_best import ComputeMetricsBest 3 | from .compute_teach import ComputeMetricsTeach 4 | -------------------------------------------------------------------------------- /transforms/rots2rfeats/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Rots2Rfeats 2 | # from .globvel import Globalvel 3 | 4 | from .globvelandy import Globalvelandy 5 | # from .rifeats import Rifeats 6 | -------------------------------------------------------------------------------- /teach/transforms/rots2rfeats/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Rots2Rfeats 2 | # from .globvel import Globalvel 3 | 4 | from .globvelandy import Globalvelandy 5 | # from .rifeats import Rifeats 6 | -------------------------------------------------------------------------------- /teach/data/tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .tensors import lengths_to_mask 2 | from .collate import collate_text_and_length, collate_pairs_and_text, collate_datastruct_and_text, collate_tensor_with_padding 3 | -------------------------------------------------------------------------------- /teach/model/motiondecoder/__init__.py: -------------------------------------------------------------------------------- 1 | from .actor import ActorAgnosticDecoder 2 | from .metactor import MetaActorDecoder 3 | from .metactor2 import MetaActorDecoder2 4 | from .gru import GRUDecoder 5 | -------------------------------------------------------------------------------- /teach/model/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .compute import TemosComputeLosses 2 | from .actioncompute import ActionComputeLosses 3 | from .compute_teach import TeachComputeLosses 4 | from .kl import KLLoss, KLLossMulti 5 | -------------------------------------------------------------------------------- /teach/model/textencoder/__init__.py: -------------------------------------------------------------------------------- 1 | from .distilbert_linear import DistilbertEncoderLinear 2 | from .distilbert_transformer import DistilbertEncoderTransformer 3 | from .actionlevel import ActionLevelTextEncoder 4 | -------------------------------------------------------------------------------- /utils/fixseed.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import random 4 | 5 | 6 | def fixseed(seed): 7 | torch.backends.cudnn.benchmark = False 8 | random.seed(seed) 9 | np.random.seed(seed) 10 | torch.manual_seed(seed) 11 | 12 | 13 | # SEED = 10 14 | # EVALSEED = 0 15 | # # Provoc warning: not fully functionnal yet 16 | # # torch.set_deterministic(True) 17 | # torch.backends.cudnn.benchmark = False 18 | # fixseed(SEED) 19 | -------------------------------------------------------------------------------- /utils/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | SMPL_DATA_PATH = "./body_models/smpl" 4 | 5 | SMPL_KINTREE_PATH = os.path.join(SMPL_DATA_PATH, "kintree_table.pkl") 6 | SMPL_MODEL_PATH = os.path.join(SMPL_DATA_PATH, "SMPL_NEUTRAL.pkl") 7 | JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(SMPL_DATA_PATH, 'J_regressor_extra.npy') 8 | 9 | ROT_CONVENTION_TO_ROT_NUMBER = { 10 | 'legacy': 23, 11 | 'no_hands': 21, 12 | 'full_hands': 51, 13 | 'mitten_hands': 33, 14 | } 15 | 16 | GENDERS = ['neutral', 'male', 'female'] 17 | NUM_BETAS = 10 -------------------------------------------------------------------------------- /teach/render/blender/data.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | class Data: 18 | def __len__(self): 19 | return self.N -------------------------------------------------------------------------------- /teach/model/utils/tools.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | def detach_to_numpy(tensor): 18 | return tensor.detach().cpu().numpy() 19 | 20 | 21 | def remove_padding(tensors, lengths): 22 | return [tensor[:tensor_length] for tensor, tensor_length in zip(tensors, lengths)] 23 | -------------------------------------------------------------------------------- /teach/tools/runid.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | """ 18 | runid util. 19 | Taken from wandb.sdk.lib.runid 20 | """ 21 | 22 | import shortuuid # type: ignore 23 | 24 | 25 | def generate_id() -> str: 26 | # ~3t run ids (36**8) 27 | run_gen = shortuuid.ShortUUID(alphabet=list("0123456789abcdefghijklmnopqrstuvwxyz")) 28 | return run_gen.random(8) -------------------------------------------------------------------------------- /teach/data/tools/tensors.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import List, Dict 18 | import torch 19 | from torch import Tensor 20 | 21 | 22 | def lengths_to_mask(lengths: List[int], device: torch.device) -> Tensor: 23 | lengths = torch.tensor(lengths, device=device) 24 | max_len = max(lengths) 25 | mask = torch.arange(max_len, device=device).expand(len(lengths), max_len) < lengths.unsqueeze(1) 26 | return mask 27 | -------------------------------------------------------------------------------- /teach/model/losses/recons.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | from torch.nn.functional import smooth_l1_loss 19 | 20 | 21 | class Recons: 22 | def __call__(self, input_motion_feats_lst, output_features_lst): 23 | recons = torch.stack([smooth_l1_loss(x, y, reduce="mean") for x,y in zip(input_motion_feats_lst, output_features_lst)]).mean() 24 | return recons 25 | 26 | def __repr__(self): 27 | return "Recons()" 28 | -------------------------------------------------------------------------------- /teach/model/utils/vae.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | 19 | 20 | def reparameterize(mu, logvar, seed=None): 21 | std = torch.exp(logvar / 2) 22 | 23 | if seed is None: 24 | eps = std.data.new(std.size()).normal_() 25 | else: 26 | generator = torch.Generator(device=mu.device) 27 | generator.manual_seed(seed) 28 | eps = std.data.new(std.size()).normal_(generator=generator) 29 | 30 | return eps.mul(std).add_(mu) -------------------------------------------------------------------------------- /teach/render/blender/vertices.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | 19 | 20 | def prepare_vertices(vertices, canonicalize=True): 21 | data = vertices 22 | # Swap axis (gravity=Z instead of Y) 23 | # data = data[..., [2, 0, 1]] 24 | 25 | # Make left/right correct 26 | # data[..., [1]] = -data[..., [1]] 27 | 28 | # Center the first root to the first frame 29 | data -= data[[0], [0], :] 30 | 31 | # Remove the floor 32 | data[..., 2] -= np.min(data[..., 2]) 33 | return data -------------------------------------------------------------------------------- /model/cfg_sampler.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | from copy import deepcopy 5 | 6 | # A wrapper model for Classifier-free guidance **SAMPLING** only 7 | # https://arxiv.org/abs/2207.12598 8 | class ClassifierFreeSampleModel(nn.Module): 9 | 10 | def __init__(self, model): 11 | super().__init__() 12 | self.model = model # model is the actual model to run 13 | 14 | assert self.model.cond_mask_prob > 0, 'Cannot run a guided diffusion on a model that has not been trained with no conditions' 15 | 16 | # pointers to inner model 17 | # self.rot2xyz = self.model.rot2xyz 18 | self.translation = self.model.translation 19 | self.njoints = self.model.njoints 20 | self.nfeats = self.model.nfeats 21 | self.data_rep = self.model.data_rep 22 | self.cond_mode = self.model.cond_mode 23 | 24 | def forward(self, x, timesteps, y=None): 25 | cond_mode = self.model.cond_mode 26 | assert cond_mode in ['text', 'action'] 27 | y_uncond = deepcopy(y) 28 | y_uncond['uncond'] = True 29 | out = self.model(x, timesteps, y) 30 | out_uncond = self.model(x, timesteps, y_uncond) 31 | return out_uncond + (y['scale'].view(-1, 1, 1, 1) * (out - out_uncond)) 32 | 33 | -------------------------------------------------------------------------------- /data_loaders/humanml/motion_loaders/dataset_motion_loader.py: -------------------------------------------------------------------------------- 1 | from t2m.data.dataset import Text2MotionDatasetV2, collate_fn 2 | from t2m.utils.word_vectorizer import WordVectorizer 3 | import numpy as np 4 | from os.path import join as pjoin 5 | from torch.utils.data import DataLoader 6 | from t2m.utils.get_opt import get_opt 7 | 8 | def get_dataset_motion_loader(opt_path, batch_size, device): 9 | opt = get_opt(opt_path, device) 10 | 11 | # Configurations of T2M dataset and KIT dataset is almost the same 12 | if opt.dataset_name == 't2m' or opt.dataset_name == 'kit': 13 | print('Loading dataset %s ...' % opt.dataset_name) 14 | 15 | mean = np.load(pjoin(opt.meta_dir, 'mean.npy')) 16 | std = np.load(pjoin(opt.meta_dir, 'std.npy')) 17 | 18 | w_vectorizer = WordVectorizer('./glove', 'our_vab') 19 | split_file = pjoin(opt.data_root, 'test.txt') 20 | dataset = Text2MotionDatasetV2(opt, mean, std, split_file, w_vectorizer) 21 | dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=4, drop_last=True, 22 | collate_fn=collate_fn, shuffle=True) 23 | else: 24 | raise KeyError('Dataset not Recognized !!') 25 | 26 | print('Ground Truth Dataset Loading Completed!!!') 27 | return dataloader, dataset -------------------------------------------------------------------------------- /utils/misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def to_numpy(tensor): 5 | if torch.is_tensor(tensor): 6 | return tensor.cpu().numpy() 7 | elif type(tensor).__module__ != 'numpy': 8 | raise ValueError("Cannot convert {} to numpy array".format( 9 | type(tensor))) 10 | return tensor 11 | 12 | 13 | def to_torch(ndarray): 14 | if type(ndarray).__module__ == 'numpy': 15 | return torch.from_numpy(ndarray) 16 | elif not torch.is_tensor(ndarray): 17 | raise ValueError("Cannot convert {} to torch tensor".format( 18 | type(ndarray))) 19 | return ndarray 20 | 21 | 22 | def cleanexit(): 23 | import sys 24 | import os 25 | try: 26 | sys.exit(0) 27 | except SystemExit: 28 | os._exit(0) 29 | 30 | def load_model_wo_clip(model, state_dict): 31 | missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) 32 | assert len(unexpected_keys) == 0 33 | assert all([k.startswith('clip_model.') for k in missing_keys]) 34 | 35 | def freeze_joints(x, joints_to_freeze): 36 | # Freezes selected joint *rotations* as they appear in the first frame 37 | # x [bs, [root+n_joints], joint_dim(6), seqlen] 38 | frozen = x.detach().clone() 39 | frozen[:, joints_to_freeze, :, :] = frozen[:, joints_to_freeze, :, :1] 40 | return frozen 41 | -------------------------------------------------------------------------------- /teach/model/losses/kl.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | 19 | 20 | class KLLoss: 21 | def __init__(self): 22 | pass 23 | 24 | def __call__(self, q, p): 25 | div = torch.distributions.kl_divergence(q, p) 26 | return div.mean() 27 | 28 | def __repr__(self): 29 | return "KLLoss()" 30 | 31 | 32 | class KLLossMulti: 33 | def __init__(self): 34 | self.klloss = KLLoss() 35 | 36 | def __call__(self, qlist, plist): 37 | return sum([self.klloss(q, p) 38 | for q, p in zip(qlist, plist)]) 39 | 40 | def __repr__(self): 41 | return "KLLossMulti()" 42 | -------------------------------------------------------------------------------- /teach/data/text2motion.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from torch.utils.data import Dataset 18 | 19 | 20 | class Text2Motion(Dataset): 21 | def __init__(self, texts, lengths): 22 | if not isinstance(lengths, list): 23 | raise NotImplementedError("Texts and lengths should be batched.") 24 | 25 | self.texts = texts 26 | self.lengths = lengths 27 | 28 | self.N = len(self.lengths) 29 | 30 | 31 | def __getitem__(self, index): 32 | return {"text": self.texts[index], 33 | "length": self.lengths[index]} 34 | 35 | def __len__(self): 36 | return self.N 37 | 38 | def __repr__(self): 39 | return f"Text2Motion dataset: {len(self)} data" 40 | -------------------------------------------------------------------------------- /model/text_encoder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from transformers import RobertaTokenizer, RobertaModel 3 | from typing import List 4 | import torch 5 | import numpy as np 6 | from torch.nn import TransformerEncoder, TransformerEncoderLayer 7 | import torch.nn.functional as F 8 | import math 9 | from transformers import CLIPTokenizer, CLIPTextModel 10 | 11 | class TextEncoder(nn.Module): 12 | def __init__( 13 | self, 14 | # clip_version="ViT-B/32" 15 | ): 16 | super().__init__() 17 | # self.clip_version = clip_version 18 | try: 19 | self.clip_model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") 20 | self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") 21 | except: 22 | self.clip_model = CLIPTextModel.from_pretrained("./clip") 23 | self.tokenizer = CLIPTokenizer.from_pretrained("./clip") 24 | 25 | def forward(self, raw_text): 26 | 27 | device = next(self.parameters()).device 28 | text = self.tokenizer(raw_text, padding='max_length', return_tensors="pt", max_length=22, truncation=True) 29 | input_ids = text["input_ids"] 30 | attention_mask = text["attention_mask"] 31 | output = self.clip_model(input_ids.to(device), attention_mask.to(device))['pooler_output'] 32 | return output 33 | 34 | if __name__ == '__main__': 35 | encoder = TextEncoder() 36 | -------------------------------------------------------------------------------- /transforms/identity.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | from torch import Tensor 19 | 20 | from .base import Datastruct, dataclass, Transform 21 | 22 | 23 | class IdentityTransform(Transform): 24 | def __init__(self, **kwargs): 25 | return 26 | 27 | def Datastruct(self, **kwargs): 28 | return IdentityDatastruct(**kwargs) 29 | 30 | def __repr__(self): 31 | return "IdentityTransform()" 32 | 33 | 34 | @dataclass 35 | class IdentityDatastruct(Datastruct): 36 | transforms: IdentityTransform 37 | 38 | features: Optional[Tensor] = None 39 | 40 | def __post_init__(self): 41 | self.datakeys = ["features"] 42 | 43 | def __len__(self): 44 | return len(self.rfeats) 45 | -------------------------------------------------------------------------------- /teach/transforms/identity.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | from torch import Tensor 19 | 20 | from .base import Datastruct, dataclass, Transform 21 | 22 | 23 | class IdentityTransform(Transform): 24 | def __init__(self, **kwargs): 25 | return 26 | 27 | def Datastruct(self, **kwargs): 28 | return IdentityDatastruct(**kwargs) 29 | 30 | def __repr__(self): 31 | return "IdentityTransform()" 32 | 33 | 34 | @dataclass 35 | class IdentityDatastruct(Datastruct): 36 | transforms: IdentityTransform 37 | 38 | features: Optional[Tensor] = None 39 | 40 | def __post_init__(self): 41 | self.datakeys = ["features"] 42 | 43 | def __len__(self): 44 | return len(self.rfeats) 45 | -------------------------------------------------------------------------------- /train/train_platforms.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | class TrainPlatform: 4 | def __init__(self, save_dir): 5 | pass 6 | 7 | def report_scalar(self, name, value, iteration, group_name=None): 8 | pass 9 | 10 | def report_args(self, args, name): 11 | pass 12 | 13 | def close(self): 14 | pass 15 | 16 | 17 | class ClearmlPlatform(TrainPlatform): 18 | def __init__(self, save_dir): 19 | from clearml import Task 20 | path, name = os.path.split(save_dir) 21 | self.task = Task.init(project_name='motion_diffusion', 22 | task_name=name, 23 | output_uri=path) 24 | self.logger = self.task.get_logger() 25 | 26 | def report_scalar(self, name, value, iteration, group_name): 27 | self.logger.report_scalar(title=group_name, series=name, iteration=iteration, value=value) 28 | 29 | def report_args(self, args, name): 30 | self.task.connect(args, name=name) 31 | 32 | def close(self): 33 | self.task.close() 34 | 35 | 36 | class TensorboardPlatform(TrainPlatform): 37 | def __init__(self, save_dir): 38 | from torch.utils.tensorboard import SummaryWriter 39 | self.writer = SummaryWriter(log_dir=save_dir) 40 | 41 | def report_scalar(self, name, value, iteration, group_name=None): 42 | self.writer.add_scalar(f'{group_name}/{name}', value, iteration) 43 | 44 | def close(self): 45 | self.writer.close() 46 | 47 | 48 | class NoPlatform(TrainPlatform): 49 | def __init__(self, save_dir): 50 | pass 51 | 52 | 53 | -------------------------------------------------------------------------------- /utils/PYTORCH3D_LICENSE: -------------------------------------------------------------------------------- 1 | BSD License 2 | 3 | For PyTorch3D software 4 | 5 | Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. 6 | 7 | Redistribution and use in source and binary forms, with or without modification, 8 | are permitted provided that the following conditions are met: 9 | 10 | * Redistributions of source code must retain the above copyright notice, this 11 | list of conditions and the following disclaimer. 12 | 13 | * Redistributions in binary form must reproduce the above copyright notice, 14 | this list of conditions and the following disclaimer in the documentation 15 | and/or other materials provided with the distribution. 16 | 17 | * Neither the name Facebook nor the names of its contributors may be used to 18 | endorse or promote products derived from this software without specific 19 | prior written permission. 20 | 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 22 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 23 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 25 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 26 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 28 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 30 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /data_loaders/a2m/humanact12poses.py: -------------------------------------------------------------------------------- 1 | import pickle as pkl 2 | import numpy as np 3 | import os 4 | from .dataset import Dataset 5 | 6 | 7 | class HumanAct12Poses(Dataset): 8 | dataname = "humanact12" 9 | 10 | def __init__(self, datapath="dataset/HumanAct12Poses", split="train", **kargs): 11 | self.datapath = datapath 12 | 13 | super().__init__(**kargs) 14 | 15 | pkldatafilepath = os.path.join(datapath, "humanact12poses.pkl") 16 | data = pkl.load(open(pkldatafilepath, "rb")) 17 | 18 | self._pose = [x for x in data["poses"]] 19 | self._num_frames_in_video = [p.shape[0] for p in self._pose] 20 | self._joints = [x for x in data["joints3D"]] 21 | 22 | self._actions = [x for x in data["y"]] 23 | 24 | total_num_actions = 12 25 | self.num_actions = total_num_actions 26 | 27 | self._train = list(range(len(self._pose))) 28 | 29 | keep_actions = np.arange(0, total_num_actions) 30 | 31 | self._action_to_label = {x: i for i, x in enumerate(keep_actions)} 32 | self._label_to_action = {i: x for i, x in enumerate(keep_actions)} 33 | 34 | self._action_classes = humanact12_coarse_action_enumerator 35 | 36 | def _load_joints3D(self, ind, frame_ix): 37 | return self._joints[ind][frame_ix] 38 | 39 | def _load_rotvec(self, ind, frame_ix): 40 | pose = self._pose[ind][frame_ix].reshape(-1, 24, 3) 41 | return pose 42 | 43 | 44 | humanact12_coarse_action_enumerator = { 45 | 0: "warm_up", 46 | 1: "walk", 47 | 2: "run", 48 | 3: "jump", 49 | 4: "drink", 50 | 5: "lift_dumbbell", 51 | 6: "sit", 52 | 7: "eat", 53 | 8: "turn steering wheel", 54 | 9: "phone", 55 | 10: "boxing", 56 | 11: "throw", 57 | } 58 | -------------------------------------------------------------------------------- /teach/data/sampling/framerate.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | 19 | 20 | # TODO: use a real subsampler.. 21 | def subsample(num_frames, last_framerate, new_framerate): 22 | step = int(last_framerate / new_framerate) 23 | assert step >= 1 24 | frames = np.arange(0, num_frames, step) 25 | return frames 26 | 27 | 28 | # TODO: use a real upsampler.. 29 | def upsample(motion, last_framerate, new_framerate): 30 | step = int(new_framerate / last_framerate) 31 | assert step >= 1 32 | 33 | # Alpha blending => interpolation 34 | alpha = np.linspace(0, 1, step+1) 35 | last = np.einsum("l,...->l...", 1-alpha, motion[:-1]) 36 | new = np.einsum("l,...->l...", alpha, motion[1:]) 37 | 38 | chuncks = (last + new)[:-1] 39 | output = np.concatenate(chuncks.swapaxes(1, 0)) 40 | # Don't forget the last one 41 | output = np.concatenate((output, motion[[-1]])) 42 | return output 43 | 44 | 45 | if __name__ == "__main__": 46 | motion = np.arange(105) 47 | submotion = motion[subsample(len(motion), 100.0, 12.5)] 48 | newmotion = upsample(submotion, 12.5, 100) 49 | 50 | print(newmotion) 51 | -------------------------------------------------------------------------------- /teach/tools/logging.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import logging 18 | import tqdm 19 | 20 | 21 | class LevelsFilter(logging.Filter): 22 | def __init__(self, levels): 23 | self.levels = [getattr(logging, level) for level in levels] 24 | 25 | def filter(self, record): 26 | return record.levelno in self.levels 27 | 28 | 29 | class StreamToLogger(object): 30 | """ 31 | Fake file-like stream object that redirects writes to a logger instance. 32 | """ 33 | def __init__(self, logger, level): 34 | self.logger = logger 35 | self.level = level 36 | self.linebuf = '' 37 | 38 | def write(self, buf): 39 | for line in buf.rstrip().splitlines(): 40 | self.logger.log(self.level, line.rstrip()) 41 | 42 | def flush(self): 43 | pass 44 | 45 | 46 | class TqdmLoggingHandler(logging.Handler): 47 | def __init__(self, level=logging.NOTSET): 48 | super().__init__(level) 49 | 50 | def emit(self, record): 51 | try: 52 | msg = self.format(record) 53 | tqdm.tqdm.write(msg) 54 | self.flush() 55 | except Exception: 56 | self.handleError(record) 57 | -------------------------------------------------------------------------------- /train/train_mdm.py: -------------------------------------------------------------------------------- 1 | # This code is based on https://github.com/openai/guided-diffusion 2 | """ 3 | Train a diffusion model on images. 4 | """ 5 | 6 | import os 7 | import json 8 | from utils.fixseed import fixseed 9 | from utils.parser_util import train_args 10 | from utils import dist_util 11 | from train.training_loop import TrainLoop 12 | from data_loaders.get_data import get_dataset_loader 13 | from utils.model_util import create_model_and_diffusion 14 | from train.train_platforms import ClearmlPlatform, TensorboardPlatform, NoPlatform # required for the eval operation 15 | 16 | def main(): 17 | args = train_args() 18 | fixseed(args.seed) 19 | train_platform_type = eval(args.train_platform_type) 20 | train_platform = train_platform_type(args.save_dir) 21 | train_platform.report_args(args, name='Args') 22 | 23 | if args.save_dir is None: 24 | raise FileNotFoundError('save_dir was not specified.') 25 | elif os.path.exists(args.save_dir) and not args.overwrite: 26 | raise FileExistsError('save_dir [{}] already exists.'.format(args.save_dir)) 27 | elif not os.path.exists(args.save_dir): 28 | os.makedirs(args.save_dir) 29 | args_path = os.path.join(args.save_dir, 'args.json') 30 | with open(args_path, 'w') as fw: 31 | json.dump(vars(args), fw, indent=4, sort_keys=True) 32 | 33 | dist_util.setup_dist(args.device) 34 | 35 | print("creating data loader...") 36 | data = get_dataset_loader(args, name=args.dataset, batch_size=args.batch_size, num_frames=args.num_frames) 37 | 38 | print("creating model and diffusion...") 39 | model, diffusion = create_model_and_diffusion(args, data) 40 | model.to(dist_util.dev()) 41 | model.rot2xyz.smpl_model.eval() 42 | 43 | print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters_wo_clip()) / 1000000.0)) 44 | print("Training...") 45 | loop = TrainLoop(args, train_platform, model, diffusion, data) 46 | loop.run_loop() 47 | train_platform.close() 48 | 49 | if __name__ == "__main__": 50 | main() 51 | -------------------------------------------------------------------------------- /train_diffusion.py: -------------------------------------------------------------------------------- 1 | # This code is based on https://github.com/openai/guided-diffusion 2 | """ 3 | Train a diffusion model on images. 4 | """ 5 | 6 | import os 7 | import json 8 | from utils.fixseed import fixseed 9 | from utils.parser_util import train_args 10 | from utils import dist_util 11 | from train.training_loop import TrainLoop 12 | from data_loaders.get_data import get_dataset_loader 13 | from utils.model_util import create_model_and_diffusion 14 | from train.train_platforms import ClearmlPlatform, TensorboardPlatform, NoPlatform # required for the eval operation 15 | 16 | def main(): 17 | args = train_args() 18 | fixseed(args.seed) 19 | train_platform_type = eval(args.train_platform_type) 20 | train_platform = train_platform_type(args.save_dir) 21 | train_platform.report_args(args, name='Args') 22 | 23 | if args.save_dir is None: 24 | raise FileNotFoundError('save_dir was not specified.') 25 | elif os.path.exists(args.save_dir) and not args.overwrite: 26 | raise FileExistsError('save_dir [{}] already exists.'.format(args.save_dir)) 27 | elif not os.path.exists(args.save_dir): 28 | os.makedirs(args.save_dir) 29 | args_path = os.path.join(args.save_dir, 'args.json') 30 | with open(args_path, 'w') as fw: 31 | json.dump(vars(args), fw, indent=4, sort_keys=True) 32 | 33 | dist_util.setup_dist(args.device) 34 | 35 | print("creating data loader...") 36 | data = get_dataset_loader(args, name=args.dataset, batch_size=args.batch_size, num_frames=args.num_frames) 37 | 38 | print("creating model and diffusion...") 39 | model, diffusion = create_model_and_diffusion(args, data) 40 | model.to(dist_util.dev()) 41 | # model.rot2xyz.smpl_model.eval() 42 | 43 | print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters_wo_clip()) / 1000000.0)) 44 | print("Training...") 45 | loop = TrainLoop(args, train_platform, model, diffusion, data) 46 | loop.run_loop() 47 | train_platform.close() 48 | 49 | if __name__ == "__main__": 50 | main() 51 | -------------------------------------------------------------------------------- /teach/launch/blender.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | # Fix blender path 18 | import sys 19 | sys.path.append("/linkhome/rech/genlgm01/uwm78rj/.local/lib/python3.9/site-packages") 20 | 21 | import bpy 22 | import os 23 | DIR = os.path.dirname(bpy.data.filepath) 24 | if DIR not in sys.path: 25 | sys.path.append(DIR) 26 | 27 | from argparse import ArgumentParser 28 | 29 | # Workaround temorary for cluster vs local 30 | # TODO fix it 31 | import socket 32 | if socket.gethostname() == 'ps018': 33 | packages_path = '/home/nathanasiou/.local/lib/python3.10/site-packages' 34 | sys.path.insert(0, packages_path) 35 | sys.path.append("/home/nathanasiou/.venvs/teach/lib/python3.10/site-packages") 36 | sys.path.append('/usr/lib/python3/dist-packages') 37 | 38 | # Monkey patch argparse such that 39 | # blender / python / hydra parsing works 40 | def parse_args(self, args=None, namespace=None): 41 | if args is not None: 42 | return self.parse_args_bak(args=args, namespace=namespace) 43 | try: 44 | idx = sys.argv.index("--") 45 | args = sys.argv[idx+1:] # the list after '--' 46 | except ValueError as e: # '--' not in the list: 47 | args = [] 48 | return self.parse_args_bak(args=args, namespace=namespace) 49 | 50 | setattr(ArgumentParser, 'parse_args_bak', ArgumentParser.parse_args) 51 | setattr(ArgumentParser, 'parse_args', parse_args) 52 | -------------------------------------------------------------------------------- /data_loaders/humanml/utils/paramUtil.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | # Define a kinematic tree for the skeletal struture 4 | 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]] 5 | 6 | kit_raw_offsets = np.array( 7 | [ 8 | [0, 0, 0], 9 | [0, 1, 0], 10 | [0, 1, 0], 11 | [0, 1, 0], 12 | [0, 1, 0], 13 | [1, 0, 0], 14 | [0, -1, 0], 15 | [0, -1, 0], 16 | [-1, 0, 0], 17 | [0, -1, 0], 18 | [0, -1, 0], 19 | [1, 0, 0], 20 | [0, -1, 0], 21 | [0, -1, 0], 22 | [0, 0, 1], 23 | [0, 0, 1], 24 | [-1, 0, 0], 25 | [0, -1, 0], 26 | [0, -1, 0], 27 | [0, 0, 1], 28 | [0, 0, 1] 29 | ] 30 | ) 31 | 32 | t2m_raw_offsets = np.array([[0,0,0], 33 | [1,0,0], 34 | [-1,0,0], 35 | [0,1,0], 36 | [0,-1,0], 37 | [0,-1,0], 38 | [0,1,0], 39 | [0,-1,0], 40 | [0,-1,0], 41 | [0,1,0], 42 | [0,0,1], 43 | [0,0,1], 44 | [0,1,0], 45 | [1,0,0], 46 | [-1,0,0], 47 | [0,0,1], 48 | [0,-1,0], 49 | [0,-1,0], 50 | [0,-1,0], 51 | [0,-1,0], 52 | [0,-1,0], 53 | [0,-1,0]]) 54 | 55 | 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]] 56 | t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]] 57 | t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]] 58 | 59 | 60 | kit_tgt_skel_id = '03950' 61 | 62 | t2m_tgt_skel_id = '000021' 63 | 64 | -------------------------------------------------------------------------------- /data_loaders/humanml_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | HML_JOINT_NAMES = [ 4 | 'pelvis', 5 | 'left_hip', 6 | 'right_hip', 7 | 'spine1', 8 | 'left_knee', 9 | 'right_knee', 10 | 'spine2', 11 | 'left_ankle', 12 | 'right_ankle', 13 | 'spine3', 14 | 'left_foot', 15 | 'right_foot', 16 | 'neck', 17 | 'left_collar', 18 | 'right_collar', 19 | 'head', 20 | 'left_shoulder', 21 | 'right_shoulder', 22 | 'left_elbow', 23 | 'right_elbow', 24 | 'left_wrist', 25 | 'right_wrist', 26 | ] 27 | 28 | NUM_HML_JOINTS = len(HML_JOINT_NAMES) # 22 SMPLH body joints 29 | 30 | HML_LOWER_BODY_JOINTS = [HML_JOINT_NAMES.index(name) for name in ['pelvis', 'left_hip', 'right_hip', 'left_knee', 'right_knee', 'left_ankle', 'right_ankle', 'left_foot', 'right_foot',]] 31 | SMPL_UPPER_BODY_JOINTS = [i for i in range(len(HML_JOINT_NAMES)) if i not in HML_LOWER_BODY_JOINTS] 32 | 33 | 34 | # Recover global angle and positions for rotation data 35 | # root_rot_velocity (B, seq_len, 1) 36 | # root_linear_velocity (B, seq_len, 2) 37 | # root_y (B, seq_len, 1) 38 | # ric_data (B, seq_len, (joint_num - 1)*3) 39 | # rot_data (B, seq_len, (joint_num - 1)*6) 40 | # local_velocity (B, seq_len, joint_num*3) 41 | # foot contact (B, seq_len, 4) 42 | HML_ROOT_BINARY = np.array([True] + [False] * (NUM_HML_JOINTS-1)) 43 | HML_ROOT_MASK = np.concatenate(([True]*(1+2+1), 44 | HML_ROOT_BINARY[1:].repeat(3), 45 | HML_ROOT_BINARY[1:].repeat(6), 46 | HML_ROOT_BINARY.repeat(3), 47 | [False] * 4)) 48 | HML_LOWER_BODY_JOINTS_BINARY = np.array([i in HML_LOWER_BODY_JOINTS for i in range(NUM_HML_JOINTS)]) 49 | HML_LOWER_BODY_MASK = np.concatenate(([True]*(1+2+1), 50 | HML_LOWER_BODY_JOINTS_BINARY[1:].repeat(3), 51 | HML_LOWER_BODY_JOINTS_BINARY[1:].repeat(6), 52 | HML_LOWER_BODY_JOINTS_BINARY.repeat(3), 53 | [True]*4)) 54 | HML_UPPER_BODY_MASK = ~HML_LOWER_BODY_MASK -------------------------------------------------------------------------------- /utils/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import socket 6 | 7 | import torch as th 8 | import torch.distributed as dist 9 | 10 | # Change this to reflect your cluster layout. 11 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 12 | GPUS_PER_NODE = 8 13 | 14 | SETUP_RETRY_COUNT = 3 15 | 16 | used_device = 0 17 | 18 | def setup_dist(device=0): 19 | """ 20 | Setup a distributed process group. 21 | """ 22 | global used_device 23 | used_device = device 24 | if dist.is_initialized(): 25 | return 26 | # os.environ["CUDA_VISIBLE_DEVICES"] = str(device) # f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 27 | 28 | # comm = MPI.COMM_WORLD 29 | # backend = "gloo" if not th.cuda.is_available() else "nccl" 30 | 31 | # if backend == "gloo": 32 | # hostname = "localhost" 33 | # else: 34 | # hostname = socket.gethostbyname(socket.getfqdn()) 35 | # os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 36 | # os.environ["RANK"] = str(comm.rank) 37 | # os.environ["WORLD_SIZE"] = str(comm.size) 38 | 39 | # port = comm.bcast(_find_free_port(), root=used_device) 40 | # os.environ["MASTER_PORT"] = str(port) 41 | # dist.init_process_group(backend=backend, init_method="env://") 42 | 43 | 44 | def dev(): 45 | """ 46 | Get the device to use for torch.distributed. 47 | """ 48 | global used_device 49 | if th.cuda.is_available() and used_device>=0: 50 | return th.device(f"cuda:{used_device}") 51 | return th.device("cpu") 52 | 53 | 54 | def load_state_dict(path, **kwargs): 55 | """ 56 | Load a PyTorch file without redundant fetches across MPI ranks. 57 | """ 58 | return th.load(path, **kwargs) 59 | 60 | 61 | def sync_params(params): 62 | """ 63 | Synchronize a sequence of Tensors across ranks from rank 0. 64 | """ 65 | for p in params: 66 | with th.no_grad(): 67 | dist.broadcast(p, 0) 68 | 69 | 70 | def _find_free_port(): 71 | try: 72 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 73 | s.bind(("", 0)) 74 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 75 | return s.getsockname()[1] 76 | finally: 77 | s.close() 78 | -------------------------------------------------------------------------------- /teach/render/blender/tools.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import bpy 18 | import numpy as np 19 | 20 | 21 | def mesh_detect(data): 22 | # heuristic 23 | if data.shape[1] > 1000: 24 | return True 25 | return False 26 | 27 | 28 | # see this for more explanation 29 | # https://gist.github.com/iyadahmed/7c7c0fae03c40bd87e75dc7059e35377 30 | # This should be solved with new version of blender 31 | class ndarray_pydata(np.ndarray): 32 | def __bool__(self) -> bool: 33 | return len(self) > 0 34 | 35 | 36 | def load_numpy_vertices_into_blender(vertices, faces, name, mat): 37 | mesh = bpy.data.meshes.new(name) 38 | mesh.from_pydata(vertices, [], faces.view(ndarray_pydata)) 39 | mesh.validate() 40 | 41 | obj = bpy.data.objects.new(name, mesh) 42 | bpy.context.scene.collection.objects.link(obj) 43 | 44 | bpy.ops.object.select_all(action='DESELECT') 45 | obj.select_set(True) 46 | obj.active_material = mat 47 | bpy.context.view_layer.objects.active = obj 48 | bpy.ops.object.shade_smooth() 49 | bpy.ops.object.select_all(action='DESELECT') 50 | return True 51 | 52 | 53 | def delete_objs(names): 54 | if not isinstance(names, list): 55 | names = [names] 56 | # bpy.ops.object.mode_set(mode='OBJECT') 57 | bpy.ops.object.select_all(action='DESELECT') 58 | for obj in bpy.context.scene.objects: 59 | for name in names: 60 | if obj.name.startswith(name) or obj.name.endswith(name): 61 | obj.select_set(True) 62 | bpy.ops.object.delete() 63 | bpy.ops.object.select_all(action='DESELECT') -------------------------------------------------------------------------------- /transforms/rots2joints/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | import hydra 24 | 25 | class Rots2Joints(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = False, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | # workaround for cluster local/sync 37 | if path is not None: 38 | rel_p = path.split('/') 39 | rel_p = rel_p[rel_p.index('deps'):] 40 | rel_p = '/'.join(rel_p) 41 | path = hydra.utils.get_original_cwd() + '/' + rel_p 42 | if normalization: 43 | mean_path = Path(path) / "mean.pt" 44 | std_path = Path(path) / "std.pt" 45 | self.register_buffer('mean', torch.load(mean_path)) 46 | self.register_buffer('std', torch.load(std_path)) 47 | 48 | def normalize(self, features: Tensor) -> Tensor: 49 | if self.normalization: 50 | features = (features - self.mean)/(self.std + self.eps) 51 | return features 52 | 53 | def unnormalize(self, features: Tensor) -> Tensor: 54 | if self.normalization: 55 | features = features * self.std + self.mean 56 | return features 57 | -------------------------------------------------------------------------------- /teach/transforms/rots2joints/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | # import hydra 24 | 25 | class Rots2Joints(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = False, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | # workaround for cluster local/sync 37 | if path is not None: 38 | rel_p = path.split('/') 39 | rel_p = rel_p[rel_p.index('deps'):] 40 | rel_p = '/'.join(rel_p) 41 | path = hydra.utils.get_original_cwd() + '/' + rel_p 42 | if normalization: 43 | mean_path = Path(path) / "mean.pt" 44 | std_path = Path(path) / "std.pt" 45 | self.register_buffer('mean', torch.load(mean_path)) 46 | self.register_buffer('std', torch.load(std_path)) 47 | 48 | def normalize(self, features: Tensor) -> Tensor: 49 | if self.normalization: 50 | features = (features - self.mean)/(self.std + self.eps) 51 | return features 52 | 53 | def unnormalize(self, features: Tensor) -> Tensor: 54 | if self.normalization: 55 | features = features * self.std + self.mean 56 | return features 57 | -------------------------------------------------------------------------------- /transforms/joints2jfeats/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | import hydra 24 | 25 | class Joints2Jfeats(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = False, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | # workaround for cluster local/sync 37 | if path is not None: 38 | rel_p = path.split('/') 39 | rel_p = rel_p[rel_p.index('deps'):] 40 | rel_p = '/'.join(rel_p) 41 | path = hydra.utils.get_original_cwd() + '/' + rel_p 42 | if normalization: 43 | mean_path = Path(path) / "jfeats_mean.pt" 44 | std_path = Path(path) / "jfeats_std.pt" 45 | self.register_buffer('mean', torch.load(mean_path)) 46 | self.register_buffer('std', torch.load(std_path)) 47 | 48 | def normalize(self, features: Tensor) -> Tensor: 49 | if self.normalization: 50 | features = (features - self.mean)/(self.std + self.eps) 51 | return features 52 | 53 | def unnormalize(self, features: Tensor) -> Tensor: 54 | if self.normalization: 55 | features = features * self.std + self.mean 56 | return features 57 | -------------------------------------------------------------------------------- /teach/transforms/joints2jfeats/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | # import hydra 24 | 25 | class Joints2Jfeats(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = False, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | # workaround for cluster local/sync 37 | if path is not None: 38 | pass 39 | # rel_p = path.split('/') 40 | # rel_p = rel_p[rel_p.index('deps'):] 41 | # rel_p = '/'.join(rel_p) 42 | # path = hydra.utils.get_original_cwd() + '/' + rel_p 43 | if normalization: 44 | mean_path = Path(path) / "jfeats_mean.pt" 45 | std_path = Path(path) / "jfeats_std.pt" 46 | self.register_buffer('mean', torch.load(mean_path)) 47 | self.register_buffer('std', torch.load(std_path)) 48 | 49 | def normalize(self, features: Tensor) -> Tensor: 50 | if self.normalization: 51 | features = (features - self.mean)/(self.std + self.eps) 52 | return features 53 | 54 | def unnormalize(self, features: Tensor) -> Tensor: 55 | if self.normalization: 56 | features = features * self.std + self.mean 57 | return features 58 | -------------------------------------------------------------------------------- /teach/data/sampling/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | 18 | from typing import Optional 19 | from dataclasses import dataclass 20 | 21 | 22 | @dataclass 23 | class FrameSampler: 24 | sampling: str = "conseq" 25 | sampling_step: int = 1 26 | request_frames: Optional[int] = None 27 | threshold_reject: int = 0.75 28 | # max_len: int = 1000 29 | max_len: int = 750 30 | # min_len: int = 10 31 | min_len: int = 15 32 | 33 | def __call__(self, num_frames): 34 | from .frames import get_frameix_from_data_index 35 | return get_frameix_from_data_index(num_frames, 36 | self.max_len, 37 | self.request_frames, 38 | self.sampling, 39 | self.sampling_step) 40 | 41 | def accept(self, duration): 42 | # Outputs have original lengths 43 | # Check if it is too long 44 | if self.request_frames is None: 45 | if duration > self.max_len: 46 | return False 47 | if duration < self.min_len: 48 | return False 49 | else: 50 | # Reject sample if the length is 51 | # too little relative to 52 | # the request frames 53 | 54 | # min_number = self.threshold_reject * self.request_frames 55 | if duration < self.min_len: # min_number: 56 | return False 57 | return True 58 | 59 | def get(self, key, default=None): 60 | return getattr(self, key, default) 61 | 62 | def __getitem__(self, key): 63 | return getattr(self, key) 64 | -------------------------------------------------------------------------------- /teach/model/utils/positional_encoding.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | import torch 19 | from torch import nn 20 | 21 | 22 | class PositionalEncoding(nn.Module): 23 | def __init__(self, d_model, dropout=0.1, 24 | max_len=5000, batch_first=False, negative=False): 25 | super().__init__() 26 | self.batch_first = batch_first 27 | 28 | self.dropout = nn.Dropout(p=dropout) 29 | self.max_len = max_len 30 | 31 | self.negative = negative 32 | 33 | if negative: 34 | pe = torch.zeros(2*max_len, d_model) 35 | position = torch.arange(-max_len, max_len, dtype=torch.float).unsqueeze(1) 36 | else: 37 | pe = torch.zeros(max_len, d_model) 38 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 39 | 40 | div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)) 41 | pe[:, 0::2] = torch.sin(position * div_term) 42 | pe[:, 1::2] = torch.cos(position * div_term) 43 | pe = pe.unsqueeze(0).transpose(0, 1) 44 | 45 | self.register_buffer('pe', pe) 46 | 47 | def forward(self, x, hist_frames=0): 48 | if not self.negative: 49 | center = 0 50 | assert hist_frames == 0 51 | first = 0 52 | else: 53 | center = self.max_len 54 | first = center-hist_frames 55 | if self.batch_first: 56 | last = first + x.shape[1] 57 | x = x + self.pe.permute(1, 0, 2)[:, first:last, :] 58 | else: 59 | last = first + x.shape[0] 60 | x = x + self.pe[first:last, :] 61 | return self.dropout(x) -------------------------------------------------------------------------------- /teach/model/dummy.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import pytorch_lightning as pl 18 | import torch 19 | from torch import nn 20 | 21 | from teach.data.tools import PoseData 22 | 23 | 24 | class Dummy(pl.LightningModule): 25 | def __init__(self, *args, **kwargs): 26 | super().__init__() 27 | self.linear = nn.Linear(10, 10) 28 | self.relu = nn.ReLU() 29 | self.store_examples = {"train": None, 30 | "val": None, 31 | "test": None} 32 | 33 | def forward(self, batch: dict) -> PoseData: 34 | return batch["pose_data"] 35 | 36 | def allsplit_step(self, split: str, batch, batch_idx): 37 | joints = batch["pose_data"].joints 38 | if batch_idx == 0: 39 | self.store_examples[split] = { 40 | "text": batch["text"], 41 | "length": batch["length"], 42 | "ref": joints, 43 | "from_text": joints, 44 | "from_motion": joints 45 | } 46 | 47 | x = batch["pose_data"].poses 48 | x = torch.rand((5, 10), device=x.device) 49 | xhat = self.linear(x) 50 | 51 | loss = torch.nn.functional.mse_loss(x, xhat) 52 | return loss 53 | 54 | def training_step(self, batch, batch_idx): 55 | return self.allsplit_step("train", batch, batch_idx) 56 | 57 | def validation_step(self, batch, batch_idx): 58 | self.allsplit_step("val", batch, batch_idx) 59 | 60 | def test_step(self, batch, batch_idx): 61 | self.allsplit_step("test", batch, batch_idx) 62 | 63 | def configure_optimizers(self): 64 | optimizer = torch.optim.AdamW(lr=1e-4, params=self.parameters()) 65 | return {"optimizer": optimizer} -------------------------------------------------------------------------------- /utils/seq_process.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | # seq's shape [bs 263 1 T] 4 | def get_canavas(seq_0, args_0, seq_1, args_1, inter_frames): 5 | # seq_canavas is the background 6 | bs, njoints, nfeats, max_frames_0 = seq_0.shape 7 | _, _, _, max_frames_1 = seq_1.shape 8 | max_frames = max_frames_0 + max_frames_1 - inter_frames 9 | # len = args_0['length'] + args_1['length'] - inter_frames 10 | seq_canavas = torch.zeros((bs, njoints, nfeats, max_frames), device=seq_0.device) 11 | # count record the weight 12 | count = torch.zeros((bs, max_frames), device=seq_0.device) 13 | 14 | # TODO convert to tensor operation 15 | for idx in range(bs): 16 | len_0 = args_0['length'][idx] 17 | len_1 = args_1['length'][idx] 18 | len = len_0 + len_1 - inter_frames 19 | seq_canavas[idx,:,:,:len_0] += seq_0[idx,:,:,:len_0] 20 | seq_canavas[idx,:,:,len_0-inter_frames:len] += seq_1[idx,:,:,:len_1] 21 | count[idx,:len_0] += 1 22 | count[idx, len_0-inter_frames:len] += 1 23 | 24 | count = count.unsqueeze(1).unsqueeze(1) 25 | seq_comp = seq_canavas / (count + 1e-5) 26 | 27 | return seq_comp 28 | 29 | def sync_fn(seq_0, args_0, seq_1, args_1, inter_frames): 30 | 31 | seq_comp = get_canavas(seq_0, args_0, seq_1, args_1, inter_frames) 32 | 33 | # ret_seq_0 = torch.zeros_like(seq_0) 34 | # ret_seq_1 = torch.zeros_like(seq_1) 35 | 36 | # for idx in range(bs): 37 | # len_0 = args_0['length'][idx] 38 | # len_1 = args_1['length'][idx] 39 | # len = len_0 + len_1 - inter_frames 40 | # ret_seq_0[idx,:,:,:len_0] = seq_comp[idx,:,:,:len_0] 41 | # ret_seq_1[idx,:,:,:len_1] = seq_comp[idx,:,:,len_0-inter_frames:len] 42 | 43 | return extract_fn(seq_comp, args_0, args_1, inter_frames) 44 | 45 | def extract_fn(seq_comp, args_0, args_1, inter_frames): 46 | bs, njoints, nfeats, _ = seq_comp.shape 47 | 48 | ret_seq_0 = torch.zeros((bs, njoints, nfeats, max(args_0['length'])), device=seq_comp.device) 49 | ret_seq_1 = torch.zeros((bs, njoints, nfeats, max(args_1['length'])), device=seq_comp.device) 50 | 51 | for idx in range(bs): 52 | len_0 = args_0['length'][idx] 53 | len_1 = args_1['length'][idx] 54 | len = len_0 + len_1 - inter_frames 55 | ret_seq_0[idx,:,:,:len_0] = seq_comp[idx,:,:,:len_0] 56 | ret_seq_1[idx,:,:,:len_1] = seq_comp[idx,:,:,len_0-inter_frames:len] 57 | 58 | return ret_seq_0, ret_seq_1 59 | -------------------------------------------------------------------------------- /teach/render/blender/sampler.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | 19 | def get_frameidx(*, mode, nframes, exact_frame, frames_to_keep, 20 | lengths=None, return_lists=True): 21 | if mode == "sequence": 22 | if lengths is not None: 23 | frameidx = [] 24 | cumlen = np.cumsum(lengths) 25 | for i, cum_i in enumerate(cumlen): 26 | if i == 0 : 27 | frameidx_i = np.linspace(0, cum_i - 1, frames_to_keep) 28 | else: 29 | frameidx_i = np.linspace(cumlen[i-1] + 1, cum_i - 1, frames_to_keep) 30 | frameidx_i = np.round(frameidx_i).astype(int) 31 | frameidx_i = list(frameidx_i) 32 | 33 | if return_lists: 34 | frameidx.append(frameidx_i) 35 | else: 36 | frameidx.extend(frameidx_i) 37 | else: 38 | frameidx_t = np.linspace(0, nframes-1, frames_to_keep) 39 | frameidx_t = np.round(frameidx_t).astype(int) 40 | frameidx_t = list(frameidx_t) 41 | frameidx = [frameidx_t] 42 | # exit() 43 | elif mode == "frame": 44 | index_frame = int(exact_frame*nframes) 45 | frameidx = [index_frame] 46 | elif mode == "video": 47 | frameidx = [] 48 | cumlen = np.cumsum(lengths) 49 | for i, cum_i in enumerate(cumlen): 50 | if i == 0 : 51 | frameidx_i = list(range(0, cum_i)) 52 | else: 53 | frameidx_i = list(range(cumlen[i-1], cum_i)) 54 | 55 | if return_lists: 56 | frameidx.append(frameidx_i) 57 | else: 58 | frameidx.extend(frameidx_i) 59 | 60 | return frameidx 61 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # README 2 | 3 | This repository contains official code for our research paper "Synthesizing Long-Term Human Motions with Diffusion Models via Coherent Sampling" (https://dl.acm.org/doi/10.1145/3581783.3611887). 4 | 5 | ## Environment setup 6 | To set up the environment, we provide an `requirements.txt` file that you can use with pip: 7 | ``` 8 | pip install -r requirements.txt 9 | ``` 10 | ## Data preparation 11 | Due to licensing restrictions, we are unable to provide pre-processed data directly. However, you can refer to [TEACH](https://github.com/athn-nik/teach#data) for specific data processing methods. 12 | 13 | Eventually, you should have a folder `data` with such a structure: 14 | ``` 15 | data 16 | |-- babel 17 | | `-- babel_v2.1 18 | | `... 19 | | `-- babel-smplh-30fps-male 20 | | `... 21 | | 22 | |-- smpl_models 23 | | `-- smplh 24 | | `--SMPLH_MALE.pkl 25 | ``` 26 | 27 | Besides, you should download the folder `deps` in [TEACH](https://github.com/athn-nik/teach/tree/main/deps) to this project. 28 | 29 | ## Pre-trained weights 30 | We provide the pretrained models here: [pretrained models](https://drive.google.com/drive/folders/1Lrj5FEt7bFFiv_VnfoDFoQgZzfF4X6RJ?usp=sharing). The 'pretrained.zip' file contains the pretrained model and training configurations used to report metrics in our paper, while 'MotionCLIP.zip' contains the model used for evaluation. You can put the pretrained model and training configuration file under `./save/pcmdm` and put the `MotipnClip.ckpt` under `./motionclip_save`. 31 | 32 | ## Running the code 33 | You can use the following three commands to obtain the results for the last three rows of the experimental results table in our paper: 34 | ``` 35 | # No special sampling 36 | python eval.py --model_path ./save/pcmdm/model000600000.pt --guidance_param 2 --inpainting_frames 0 37 | 38 | # Past inpainting sampling 39 | python eval.py --model_path ./save/pcmdm/model000600000.pt --guidance_param 2 --inpainting_frames 2 40 | 41 | # Compositional transition sampling 42 | python eval.py --model_path ./save/pcmdm/model000600000.pt --guidance_param 2 --composition True --inter_frames 2 43 | ``` 44 | 45 | Besides, if you want to train a model from scratch, you can use this comman: 46 | ``` 47 | python train_diffusion.py --save_dir ./save/pcmdm --dataset babel --hist_frames 5 48 | ``` 49 | 50 | ## Acknowledgments 51 | Our code is based on [TEACH](https://github.com/athn-nik/teach) and [MDM](https://github.com/GuyTevet/motion-diffusion-model). Thanks for their greate work! 52 | -------------------------------------------------------------------------------- /utils/smpl.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | import utls.geometry as geometry 19 | import torch 20 | from torch import Tensor 21 | from utils.easyconvert import matrix_to, axis_angle_to 22 | from teach.transforms.smpl import RotTransDatastruct 23 | 24 | 25 | def canonicalize_smplh(poses: Tensor, trans: Optional[Tensor] = None): 26 | bs, nframes, njoints = poses.shape[:3] 27 | 28 | global_orient = poses[:, :, 0] 29 | 30 | # first global rotations 31 | rot2d = geometry.matrix_to_axis_angle(global_orient[:, 0]) 32 | rot2d[:, :2] = 0 # Remove the rotation along the vertical axis 33 | rot2d = geometry.axis_angle_to_matrix(rot2d) 34 | 35 | # Rotate the global rotation to eliminate Z rotations 36 | global_orient = torch.einsum("ikj,imkl->imjl", rot2d, global_orient) 37 | 38 | # Construct canonicalized version of x 39 | xc = torch.cat((global_orient[:, :, None], poses[:, :, 1:]), dim=2) 40 | 41 | if trans is not None: 42 | vel = trans[:, 1:] - trans[:, :-1] 43 | # Turn the translation as well 44 | vel = torch.einsum("ikj,ilk->ilj", rot2d, vel) 45 | trans = torch.cat((torch.zeros(bs, 1, 3, device=vel.device), 46 | torch.cumsum(vel, 1)), 1) 47 | return xc, trans 48 | else: 49 | return xc 50 | 51 | 52 | def smpl_data_to_matrix_and_trans(data, nohands=True): 53 | trans = data['trans'] 54 | nframes = len(trans) 55 | try: 56 | axis_angle_poses = data['poses'] 57 | axis_angle_poses = data['poses'].reshape(nframes, -1, 3) 58 | except: 59 | breakpoint() 60 | 61 | if nohands: 62 | axis_angle_poses = axis_angle_poses[:, :22] 63 | 64 | matrix_poses = axis_angle_to("matrix", axis_angle_poses) 65 | 66 | return RotTransDatastruct(rots=matrix_poses, trans=trans) 67 | -------------------------------------------------------------------------------- /teach/data/tools/smpl.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | import teach.tools.geometry as geometry 19 | import torch 20 | from torch import Tensor 21 | from teach.tools.easyconvert import matrix_to, axis_angle_to 22 | from teach.transforms.smpl import RotTransDatastruct 23 | 24 | 25 | def canonicalize_smplh(poses: Tensor, trans: Optional[Tensor] = None): 26 | bs, nframes, njoints = poses.shape[:3] 27 | 28 | global_orient = poses[:, :, 0] 29 | 30 | # first global rotations 31 | rot2d = geometry.matrix_to_axis_angle(global_orient[:, 0]) 32 | rot2d[:, :2] = 0 # Remove the rotation along the vertical axis 33 | rot2d = geometry.axis_angle_to_matrix(rot2d) 34 | 35 | # Rotate the global rotation to eliminate Z rotations 36 | global_orient = torch.einsum("ikj,imkl->imjl", rot2d, global_orient) 37 | 38 | # Construct canonicalized version of x 39 | xc = torch.cat((global_orient[:, :, None], poses[:, :, 1:]), dim=2) 40 | 41 | if trans is not None: 42 | vel = trans[:, 1:] - trans[:, :-1] 43 | # Turn the translation as well 44 | vel = torch.einsum("ikj,ilk->ilj", rot2d, vel) 45 | trans = torch.cat((torch.zeros(bs, 1, 3, device=vel.device), 46 | torch.cumsum(vel, 1)), 1) 47 | return xc, trans 48 | else: 49 | return xc 50 | 51 | 52 | def smpl_data_to_matrix_and_trans(data, nohands=True): 53 | trans = data['trans'] 54 | nframes = len(trans) 55 | try: 56 | axis_angle_poses = data['poses'] 57 | axis_angle_poses = data['poses'].reshape(nframes, -1, 3) 58 | except: 59 | breakpoint() 60 | 61 | if nohands: 62 | axis_angle_poses = axis_angle_poses[:, :22] 63 | 64 | matrix_poses = axis_angle_to("matrix", axis_angle_poses) 65 | 66 | return RotTransDatastruct(rots=matrix_poses, trans=trans) 67 | -------------------------------------------------------------------------------- /transforms/rots2rfeats/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | import hydra 24 | 25 | class Rots2Rfeats(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = True, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | if normalization: 37 | # workaround for cluster local/sync 38 | rel_p = path.split('/') 39 | # superhacky it is for the datatype ugly stuff change it, copy the main stuff to seperate_pairs dict 40 | if rel_p[-1] == 'separate_pairs': 41 | rel_p.remove('separate_pairs') 42 | ######################################################## 43 | rel_p = rel_p[rel_p.index('deps'):] 44 | rel_p = '/'.join(rel_p) 45 | path = hydra.utils.get_original_cwd() + '/' + rel_p 46 | mean_path = Path(path) / "rfeats_mean.pt" 47 | std_path = Path(path) / "rfeats_std.pt" 48 | 49 | self.register_buffer('mean', torch.load(mean_path)) 50 | self.register_buffer('std', torch.load(std_path)) 51 | 52 | def normalize(self, features: Tensor) -> Tensor: 53 | if self.normalization: 54 | features = (features - self.mean)/(self.std + self.eps) 55 | return features 56 | 57 | def unnormalize(self, features: Tensor) -> Tensor: 58 | if self.normalization: 59 | features = features * self.std + self.mean 60 | return features 61 | -------------------------------------------------------------------------------- /teach/transforms/rots2rfeats/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import torch 20 | from torch import Tensor, nn 21 | from pathlib import Path 22 | import os 23 | # import hydra 24 | 25 | class Rots2Rfeats(nn.Module): 26 | def __init__(self, path: Optional[str] = None, 27 | normalization: bool = True, 28 | eps: float = 1e-12, 29 | **kwargs) -> None: 30 | if normalization and path is None: 31 | raise TypeError("You should provide a path if normalization is on.") 32 | 33 | super().__init__() 34 | self.normalization = normalization 35 | self.eps = eps 36 | if normalization: 37 | # workaround for cluster local/sync 38 | rel_p = path.split('/') 39 | # superhacky it is for the datatype ugly stuff change it, copy the main stuff to seperate_pairs dict 40 | if rel_p[-1] == 'separate_pairs': 41 | rel_p.remove('separate_pairs') 42 | ######################################################## 43 | rel_p = rel_p[rel_p.index('deps'):] 44 | rel_p = '/'.join(rel_p) 45 | # path = hydra.utils.get_original_cwd() + '/' + rel_p 46 | mean_path = Path(path) / "rfeats_mean.pt" 47 | std_path = Path(path) / "rfeats_std.pt" 48 | 49 | self.register_buffer('mean', torch.load(mean_path)) 50 | self.register_buffer('std', torch.load(std_path)) 51 | 52 | def normalize(self, features: Tensor) -> Tensor: 53 | if self.normalization: 54 | features = (features - self.mean)/(self.std + self.eps) 55 | return features 56 | 57 | def unnormalize(self, features: Tensor) -> Tensor: 58 | if self.normalization: 59 | features = features * self.std + self.mean 60 | return features 61 | -------------------------------------------------------------------------------- /data_loaders/tensors.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def lengths_to_mask(lengths, max_len): 4 | # max_len = max(lengths) 5 | mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1) 6 | return mask 7 | 8 | 9 | def collate_tensors(batch): 10 | dims = batch[0].dim() 11 | max_size = [max([b.size(i) for b in batch]) for i in range(dims)] 12 | size = (len(batch),) + tuple(max_size) 13 | canvas = batch[0].new_zeros(size=size) 14 | for i, b in enumerate(batch): 15 | sub_tensor = canvas[i] 16 | for d in range(dims): 17 | sub_tensor = sub_tensor.narrow(d, 0, b.size(d)) 18 | sub_tensor.add_(b) 19 | return canvas 20 | 21 | 22 | def collate(batch): 23 | notnone_batches = [b for b in batch if b is not None] 24 | databatch = [b['inp'] for b in notnone_batches] 25 | if 'lengths' in notnone_batches[0]: 26 | lenbatch = [b['lengths'] for b in notnone_batches] 27 | else: 28 | lenbatch = [len(b['inp'][0][0]) for b in notnone_batches] 29 | 30 | 31 | databatchTensor = collate_tensors(databatch) 32 | lenbatchTensor = torch.as_tensor(lenbatch) 33 | maskbatchTensor = lengths_to_mask(lenbatchTensor, databatchTensor.shape[-1]).unsqueeze(1).unsqueeze(1) # unqueeze for broadcasting 34 | 35 | motion = databatchTensor 36 | cond = {'y': {'mask': maskbatchTensor, 'lengths': lenbatchTensor}} 37 | 38 | if 'text' in notnone_batches[0]: 39 | textbatch = [b['text'] for b in notnone_batches] 40 | cond['y'].update({'text': textbatch}) 41 | 42 | if 'tokens' in notnone_batches[0]: 43 | textbatch = [b['tokens'] for b in notnone_batches] 44 | cond['y'].update({'tokens': textbatch}) 45 | 46 | if 'action' in notnone_batches[0]: 47 | actionbatch = [b['action'] for b in notnone_batches] 48 | cond['y'].update({'action': torch.as_tensor(actionbatch).unsqueeze(1)}) 49 | 50 | # collate action textual names 51 | if 'action_text' in notnone_batches[0]: 52 | action_text = [b['action_text']for b in notnone_batches] 53 | cond['y'].update({'action_text': action_text}) 54 | 55 | return motion, cond 56 | 57 | # an adapter to our collate func 58 | def t2m_collate(batch): 59 | # batch.sort(key=lambda x: x[3], reverse=True) 60 | adapted_batch = [{ 61 | 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] 62 | 'text': b[2], #b[0]['caption'] 63 | 'tokens': b[6], 64 | 'lengths': b[5], 65 | } for b in batch] 66 | return collate(adapted_batch) 67 | 68 | 69 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | ./dataset 132 | ./data 133 | ./deps 134 | ./save 135 | ./motionclip_save 136 | 137 | *.out 138 | *.json -------------------------------------------------------------------------------- /teach/model/textencoder/distilbert_linear.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import List, Union 18 | import pytorch_lightning as pl 19 | 20 | import torch.nn as nn 21 | import os 22 | 23 | import torch 24 | from torch import Tensor 25 | from torch.distributions.distribution import Distribution 26 | 27 | from .distilbert import DistilbertEncoderBase 28 | 29 | 30 | class DistilbertEncoderLinear(DistilbertEncoderBase): 31 | def __init__(self, modelpath: str, 32 | finetune: bool = False, 33 | vae: bool = True, 34 | latent_dim: int = 256, 35 | **kwargs) -> None: 36 | super().__init__(modelpath=modelpath, finetune=finetune) 37 | self.save_hyperparameters(logger=False) 38 | 39 | encoded_dim = self.text_encoded_dim 40 | # Train a embedding layer 41 | if vae: 42 | self.mu_projection = nn.Sequential(nn.ReLU(), 43 | nn.Linear(encoded_dim, latent_dim)) 44 | self.logvar_projection = nn.Sequential(nn.ReLU(), 45 | nn.Linear(encoded_dim, latent_dim)) 46 | else: 47 | self.projection = nn.Sequential(nn.ReLU(), 48 | nn.Linear(encoded_dim, latent_dim)) 49 | 50 | def forward(self, texts: List[str]) -> Union[Tensor, Distribution]: 51 | last_hidden_state = self.get_last_hidden_state(texts) 52 | latent = last_hidden_state[:, 0] 53 | 54 | if self.hparams.vae: 55 | mu = self.mu_projection(latent) 56 | logvar = self.logvar_projection(latent) 57 | std = logvar.exp().pow(0.5) 58 | # https://github.com/kampta/pytorch-distributions/blob/master/gaussian_vae.py 59 | return torch.distributions.Normal(mu, std) 60 | else: 61 | return self.projection(latent) 62 | -------------------------------------------------------------------------------- /teach/model/motiondecoder/gru.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn as nn 19 | import numpy as np 20 | import pytorch_lightning as pl 21 | 22 | from typing import List, Optional 23 | from torch import nn, Tensor 24 | 25 | from teach.model.utils import PositionalEncoding 26 | from teach.data.tools import lengths_to_mask 27 | 28 | 29 | class GRUDecoder(pl.LightningModule): 30 | def __init__(self, nfeats: int, 31 | latent_dim: int = 256, 32 | num_layers: int = 4, **kwargs) -> None: 33 | 34 | super().__init__() 35 | self.save_hyperparameters(logger=False) 36 | 37 | output_feats = nfeats 38 | # self.feats_embedding = nn.Linear(self.input_feats, latent_dim) 39 | 40 | self.emb_layer = nn.Linear(latent_dim+1, latent_dim) 41 | self.gru = nn.GRU(latent_dim, latent_dim, num_layers=num_layers) 42 | self.final_layer = nn.Linear(latent_dim, output_feats) 43 | 44 | def forward(self, z: Tensor, lengths: List[int]): 45 | mask = lengths_to_mask(lengths, z.device) 46 | latent_dim = z.shape[1] 47 | bs, nframes = mask.shape 48 | nfeats = self.hparams.nfeats 49 | 50 | lengths = torch.tensor(lengths, device=z.device) 51 | 52 | # Repeat the input 53 | z = z[None].repeat((nframes, 1, 1)) 54 | 55 | # Add time information to the input 56 | time = mask * 1/(lengths[..., None]-1) 57 | time = (time[:, None] * torch.arange(time.shape[1], device=z.device))[:, 0] 58 | time = time.T[..., None] 59 | z = torch.cat((z, time), 2) 60 | 61 | # emb to latent space again 62 | z = self.emb_layer(z) 63 | 64 | # pass to gru 65 | z = self.gru(z)[0] 66 | output = self.final_layer(z) 67 | 68 | # zero for padded area 69 | output[~mask.T] = 0 70 | # Pytorch GRU: [Sequence, Batch size, ...] 71 | feats = output.permute(1, 0, 2) 72 | 73 | return feats 74 | -------------------------------------------------------------------------------- /teach/logger/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import os 18 | from pathlib import Path 19 | from .wandb_log import WandbLogger 20 | from pytorch_lightning.loggers import TensorBoardLogger 21 | from pytorch_lightning.loggers.base import DummyLogger 22 | from hydra.utils import to_absolute_path 23 | from omegaconf import DictConfig, OmegaConf 24 | from .tools import cfg_to_flatten_config 25 | import types 26 | import wandb 27 | 28 | def instantiate_logger(cfg: DictConfig): 29 | conf = OmegaConf.to_container(cfg.logger, resolve=True) 30 | name = conf.pop('logger_name') 31 | 32 | if name == 'wandb': 33 | project_save_dir = to_absolute_path(Path(cfg.path.working_dir) / conf['save_dir']) 34 | Path(project_save_dir).mkdir(exist_ok=True) 35 | conf['dir'] = project_save_dir 36 | conf['config'] = cfg_to_flatten_config(cfg) 37 | # maybe do this for connection error in cluster, could be redundant 38 | conf['settings'] = wandb.Settings(start_method="fork") 39 | # conf['mode']= 'online' if not cfg.logger.offline else 'offline' 40 | conf['notes']= cfg.logger.notes if cfg.logger.notes is not None else None 41 | conf['tags'] = cfg.logger.tags.strip().split(',')\ 42 | if cfg.logger.tags is not None else None 43 | logger = WandbLogger(**conf) 44 | # begin / end already defined 45 | 46 | else: 47 | def begin(self, *args, **kwargs): 48 | return 49 | 50 | def end(self, *args, **kwargs): 51 | return 52 | 53 | if name == 'tensorboard': 54 | logger = TensorBoardLogger(**conf) 55 | logger.begin = begin 56 | logger.end = end 57 | elif name in ["none", None]: 58 | logger = DummyLogger() 59 | logger.begin = begin 60 | logger.end = end 61 | else: 62 | raise NotImplementedError("This logger is not recognized.") 63 | 64 | logger.lname = name 65 | return logger 66 | -------------------------------------------------------------------------------- /train_contrastive.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | import random 4 | import numpy as np 5 | from torch import optim 6 | import torch.nn as nn 7 | import os 8 | import time 9 | import pytorch_lightning as pl 10 | from argparse import ArgumentParser 11 | from pytorch_lightning import Trainer 12 | import pytorch_lightning.callbacks as plc 13 | from pytorch_lightning.loggers import TensorBoardLogger 14 | from model.motion_clip import MotionClip 15 | from motionclip_data_loaders.pretrain_datamodule import PretrainDataModule 16 | 17 | 18 | def main(args): 19 | pl.seed_everything(args.seed) 20 | 21 | # data 22 | # train_dataset = GINPretrainDataset(args.root, args.text_max_len, args.graph_aug1, args.graph_aug2) 23 | # dm = LightningDataset(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers) 24 | dm = PretrainDataModule.from_argparse_args(args) 25 | 26 | 27 | # model 28 | model = MotionClip( 29 | temperature=args.temperature, 30 | motion_hidden_dim=args.motion_hidden_dim, 31 | # motion_num_layers=args.gin_num_layers, 32 | # drop_ratio=args.drop_ratio, 33 | # graph_pooling=args.graph_pooling, 34 | # graph_self=args.graph_self, 35 | text_hidden_dim=args.text_hidden_dim, 36 | # bert_pretrain=args.bert_pretrain, 37 | projection_dim=args.projection_dim, 38 | lr=args.lr, 39 | weight_decay=args.weight_decay 40 | ) 41 | print('total params:', sum(p.numel() for p in model.parameters())) 42 | 43 | callbacks = [] 44 | callbacks.append(plc.ModelCheckpoint(dirpath="all_checkpoints/pretrain_gin/", every_n_epochs=5)) 45 | strategy = pl.strategies.DDPSpawnStrategy(find_unused_parameters=False) 46 | trainer = Trainer.from_argparse_args(args, callbacks=callbacks, strategy=strategy) 47 | 48 | trainer.fit(model, datamodule=dm) 49 | 50 | 51 | 52 | if __name__ == '__main__': 53 | parser = argparse.ArgumentParser() 54 | # parser.add_argument('--default_root_dir', type=str, default='./checkpoints/', help='location of model checkpoints') 55 | # parser.add_argument('--max_epochs', type=int, default=500) 56 | 57 | # GPU 58 | parser.add_argument('--seed', type=int, default=100, help='random seed') 59 | # parser.add_argument('--devices', type=str, default='0,1,2,3', help='device ids of multile gpus') 60 | 61 | parser = Trainer.add_argparse_args(parser) 62 | parser = MotionClip.add_model_specific_args(parser) # add model args 63 | parser = PretrainDataModule.add_argparse_args(parser) # add data args 64 | args = parser.parse_args() 65 | 66 | print('Args in experiment:') 67 | print(args) 68 | 69 | main(args) 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | -------------------------------------------------------------------------------- /teach/render/blender/camera.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import bpy 18 | 19 | class Camera: 20 | def __init__(self, *, first_root, mode, is_mesh, fakeinone=False): 21 | camera = bpy.data.objects['Camera'] 22 | 23 | ## initial position 24 | camera.location.x = 7.36 25 | camera.location.y = -6.93 26 | if is_mesh: 27 | # camera.location.z = 5.45 28 | camera.location.z = 5.6 29 | else: 30 | camera.location.z = 5.2 31 | 32 | # wider point of view 33 | if mode == "sequence": 34 | if is_mesh: 35 | if fakeinone: 36 | camera.data.lens = 75 37 | else: 38 | camera.data.lens = 100 39 | 40 | else: 41 | camera.data.lens = 85 42 | elif mode == "frame": 43 | if is_mesh: 44 | camera.data.lens = 130 45 | else: 46 | camera.data.lens = 140 47 | elif mode == "video": 48 | if is_mesh: 49 | camera.data.lens = 90 50 | else: 51 | camera.data.lens = 140 52 | 53 | # camera.location.x += 0.75 54 | 55 | self.mode = mode 56 | self.camera = camera 57 | 58 | self.camera.location.x += first_root[0] 59 | self.camera.location.y += first_root[1] 60 | 61 | self._root = first_root 62 | 63 | def eye_view(self): 64 | 65 | from math import radians 66 | cam_rot = bpy.data.objects["Camera"].rotation_euler 67 | 68 | self.camera.location.z += 0.7 69 | self.camera.location.y += 1.2 70 | self.camera.location.x -= 1.25 71 | self.camera.rotation_euler = (cam_rot[0] + radians(-10), cam_rot[1] + radians(0), cam_rot[2] + radians(0)) 72 | 73 | def update(self, newroot): 74 | delta_root = newroot - self._root 75 | 76 | self.camera.location.x += delta_root[0] 77 | self.camera.location.y += delta_root[1] 78 | 79 | self._root = newroot 80 | -------------------------------------------------------------------------------- /teach/callback/progress.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import logging 18 | 19 | from pytorch_lightning import LightningModule, Trainer 20 | from pytorch_lightning.callbacks import Callback 21 | import psutil 22 | 23 | logger = logging.getLogger(__name__) 24 | 25 | 26 | class ProgressLogger(Callback): 27 | def __init__(self, 28 | metric_monitor: dict, 29 | precision: int = 3): 30 | # Metric to monitor 31 | self.metric_monitor = metric_monitor 32 | self.precision = precision 33 | 34 | def on_train_start(self, trainer: Trainer, pl_module: LightningModule, **kwargs) -> None: 35 | logger.info("Training started") 36 | 37 | def on_train_end(self, trainer: Trainer, pl_module: LightningModule, **kwargs) -> None: 38 | logger.info("Training done") 39 | 40 | def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule, **kwargs) -> None: 41 | if trainer.sanity_checking: 42 | logger.info("Sanity checking ok.") 43 | 44 | def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule, **kwargs) -> None: 45 | metric_format = f"{{:.{self.precision}e}}" 46 | line = f"Epoch {trainer.current_epoch}" 47 | line = f"{line:>{len('Epoch xxxx')}}" # Right padding 48 | metrics_str = [] 49 | 50 | losses_dict = trainer.callback_metrics 51 | for metric_name, dico_name in self.metric_monitor.items(): 52 | if dico_name not in losses_dict: 53 | dico_name = f"losses/{dico_name}" 54 | 55 | if dico_name in losses_dict: 56 | metric = losses_dict[dico_name].item() 57 | metric = metric_format.format(metric) 58 | metric = f"{metric_name} {metric}" 59 | metrics_str.append(metric) 60 | 61 | if len(metrics_str) == 0: 62 | return 63 | 64 | memory = f"Memory {psutil.virtual_memory().percent}%" 65 | line = line + ": " + " ".join(metrics_str) + " " + memory 66 | logger.info(line) 67 | -------------------------------------------------------------------------------- /transforms/xyz.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | from torch import Tensor 19 | 20 | from .base import Datastruct, dataclass, Transform 21 | from teach.data.tools import collate_tensor_with_padding 22 | 23 | from .joints2jfeats import Joints2Jfeats 24 | 25 | 26 | class XYZTransform(Transform): 27 | def __init__(self, joints2jfeats: Joints2Jfeats, **kwargs): 28 | self.joints2jfeats = joints2jfeats 29 | 30 | def Datastruct(self, **kwargs): 31 | return XYZDatastruct(_joints2jfeats=self.joints2jfeats, 32 | transforms=self, 33 | **kwargs) 34 | 35 | def __repr__(self): 36 | return "XYZTransform()" 37 | 38 | 39 | @dataclass 40 | class XYZDatastruct(Datastruct): 41 | transforms: XYZTransform 42 | _joints2jfeats: Joints2Jfeats 43 | 44 | features: Optional[Tensor] = None 45 | joints_: Optional[Tensor] = None 46 | jfeats_: Optional[Tensor] = None 47 | 48 | def __post_init__(self): 49 | self.datakeys = ["features", "joints_", "jfeats_"] 50 | # starting point 51 | if self.features is not None and self.jfeats_ is None: 52 | self.jfeats_ = self.features 53 | 54 | @property 55 | def joints(self): 56 | # Cached value 57 | if self.joints_ is not None: 58 | return self.joints_ 59 | 60 | # self.jfeats_ should be defined 61 | assert self.jfeats_ is not None 62 | 63 | self._joints2jfeats.to(self.jfeats.device) 64 | self.joints_ = self._joints2jfeats.inverse(self.jfeats) 65 | return self.joints_ 66 | 67 | @property 68 | def jfeats(self): 69 | # Cached value 70 | if self.jfeats_ is not None: 71 | return self.jfeats_ 72 | 73 | # self.joints_ should be defined 74 | assert self.joints_ is not None 75 | 76 | self._joints2jfeats.to(self.joints.device) 77 | self.jfeats_ = self._joints2jfeats(self.joints) 78 | return self.jfeats_ 79 | 80 | def __len__(self): 81 | return len(self.jfeats) 82 | -------------------------------------------------------------------------------- /teach/transforms/xyz.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | from torch import Tensor 19 | 20 | from .base import Datastruct, dataclass, Transform 21 | from teach.data.tools import collate_tensor_with_padding 22 | 23 | from .joints2jfeats import Joints2Jfeats 24 | 25 | 26 | class XYZTransform(Transform): 27 | def __init__(self, joints2jfeats: Joints2Jfeats, **kwargs): 28 | self.joints2jfeats = joints2jfeats 29 | 30 | def Datastruct(self, **kwargs): 31 | return XYZDatastruct(_joints2jfeats=self.joints2jfeats, 32 | transforms=self, 33 | **kwargs) 34 | 35 | def __repr__(self): 36 | return "XYZTransform()" 37 | 38 | 39 | @dataclass 40 | class XYZDatastruct(Datastruct): 41 | transforms: XYZTransform 42 | _joints2jfeats: Joints2Jfeats 43 | 44 | features: Optional[Tensor] = None 45 | joints_: Optional[Tensor] = None 46 | jfeats_: Optional[Tensor] = None 47 | 48 | def __post_init__(self): 49 | self.datakeys = ["features", "joints_", "jfeats_"] 50 | # starting point 51 | if self.features is not None and self.jfeats_ is None: 52 | self.jfeats_ = self.features 53 | 54 | @property 55 | def joints(self): 56 | # Cached value 57 | if self.joints_ is not None: 58 | return self.joints_ 59 | 60 | # self.jfeats_ should be defined 61 | assert self.jfeats_ is not None 62 | 63 | self._joints2jfeats.to(self.jfeats.device) 64 | self.joints_ = self._joints2jfeats.inverse(self.jfeats) 65 | return self.joints_ 66 | 67 | @property 68 | def jfeats(self): 69 | # Cached value 70 | if self.jfeats_ is not None: 71 | return self.jfeats_ 72 | 73 | # self.joints_ should be defined 74 | assert self.joints_ is not None 75 | 76 | self._joints2jfeats.to(self.joints.device) 77 | self.jfeats_ = self._joints2jfeats(self.joints) 78 | return self.jfeats_ 79 | 80 | def __len__(self): 81 | return len(self.jfeats) 82 | -------------------------------------------------------------------------------- /transforms/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from dataclasses import dataclass, fields 18 | 19 | 20 | class Transform: 21 | def collate(self, lst_datastruct): 22 | from teach.data.tools import collate_tensor_with_padding 23 | example = lst_datastruct[0] 24 | 25 | def collate_or_none(key): 26 | if example[key] is None: 27 | return None 28 | key_lst = [x[key] for x in lst_datastruct] 29 | return collate_tensor_with_padding(key_lst) 30 | 31 | kwargs = {key: collate_or_none(key) 32 | for key in example.datakeys} 33 | 34 | return self.Datastruct(**kwargs) 35 | 36 | 37 | # Inspired from SMPLX library 38 | # need to define "datakeys" and transforms 39 | @dataclass 40 | class Datastruct: 41 | def __getitem__(self, key): 42 | return getattr(self, key) 43 | 44 | def __setitem__(self, key, value): 45 | self.__dict__[key] = value 46 | 47 | def get(self, key, default=None): 48 | return getattr(self, key, default) 49 | 50 | def __iter__(self): 51 | return self.keys() 52 | 53 | def keys(self): 54 | keys = [t.name for t in fields(self)] 55 | return iter(keys) 56 | 57 | def values(self): 58 | values = [getattr(self, t.name) for t in fields(self)] 59 | return iter(values) 60 | 61 | def items(self): 62 | data = [(t.name, getattr(self, t.name)) for t in fields(self)] 63 | return iter(data) 64 | 65 | def to(self, *args, **kwargs): 66 | for key in self.datakeys: 67 | if self[key] is not None: 68 | self[key] = self[key].to(*args, **kwargs) 69 | return self 70 | 71 | @property 72 | def device(self): 73 | return self[self.datakeys[0]].device 74 | 75 | def detach(self): 76 | def detach_or_none(tensor): 77 | if tensor is not None: 78 | return tensor.detach() 79 | return None 80 | 81 | kwargs = {key: detach_or_none(self[key]) 82 | for key in self.datakeys} 83 | return self.transforms.Datastruct(**kwargs) 84 | -------------------------------------------------------------------------------- /teach/transforms/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from dataclasses import dataclass, fields 18 | 19 | 20 | class Transform: 21 | def collate(self, lst_datastruct): 22 | from teach.data.tools import collate_tensor_with_padding 23 | example = lst_datastruct[0] 24 | 25 | def collate_or_none(key): 26 | if example[key] is None: 27 | return None 28 | key_lst = [x[key] for x in lst_datastruct] 29 | return collate_tensor_with_padding(key_lst) 30 | 31 | kwargs = {key: collate_or_none(key) 32 | for key in example.datakeys} 33 | 34 | return self.Datastruct(**kwargs) 35 | 36 | 37 | # Inspired from SMPLX library 38 | # need to define "datakeys" and transforms 39 | @dataclass 40 | class Datastruct: 41 | def __getitem__(self, key): 42 | return getattr(self, key) 43 | 44 | def __setitem__(self, key, value): 45 | self.__dict__[key] = value 46 | 47 | def get(self, key, default=None): 48 | return getattr(self, key, default) 49 | 50 | def __iter__(self): 51 | return self.keys() 52 | 53 | def keys(self): 54 | keys = [t.name for t in fields(self)] 55 | return iter(keys) 56 | 57 | def values(self): 58 | values = [getattr(self, t.name) for t in fields(self)] 59 | return iter(values) 60 | 61 | def items(self): 62 | data = [(t.name, getattr(self, t.name)) for t in fields(self)] 63 | return iter(data) 64 | 65 | def to(self, *args, **kwargs): 66 | for key in self.datakeys: 67 | if self[key] is not None: 68 | self[key] = self[key].to(*args, **kwargs) 69 | return self 70 | 71 | @property 72 | def device(self): 73 | return self[self.datakeys[0]].device 74 | 75 | def detach(self): 76 | def detach_or_none(tensor): 77 | if tensor is not None: 78 | return tensor.detach() 79 | return None 80 | 81 | kwargs = {key: detach_or_none(self[key]) 82 | for key in self.datakeys} 83 | return self.transforms.Datastruct(**kwargs) 84 | -------------------------------------------------------------------------------- /data_loaders/humanml/utils/get_opt.py: -------------------------------------------------------------------------------- 1 | import os 2 | from argparse import Namespace 3 | import re 4 | from os.path import join as pjoin 5 | from data_loaders.humanml.utils.word_vectorizer import POS_enumerator 6 | 7 | 8 | def is_float(numStr): 9 | flag = False 10 | numStr = str(numStr).strip().lstrip('-').lstrip('+') # 去除正数(+)、负数(-)符号 11 | try: 12 | reg = re.compile(r'^[-+]?[0-9]+\.[0-9]+$') 13 | res = reg.match(str(numStr)) 14 | if res: 15 | flag = True 16 | except Exception as ex: 17 | print("is_float() - error: " + str(ex)) 18 | return flag 19 | 20 | 21 | def is_number(numStr): 22 | flag = False 23 | numStr = str(numStr).strip().lstrip('-').lstrip('+') # 去除正数(+)、负数(-)符号 24 | if str(numStr).isdigit(): 25 | flag = True 26 | return flag 27 | 28 | 29 | def get_opt(opt_path, device): 30 | opt = Namespace() 31 | opt_dict = vars(opt) 32 | 33 | skip = ('-------------- End ----------------', 34 | '------------ Options -------------', 35 | '\n') 36 | print('Reading', opt_path) 37 | with open(opt_path) as f: 38 | for line in f: 39 | if line.strip() not in skip: 40 | # print(line.strip()) 41 | key, value = line.strip().split(': ') 42 | if value in ('True', 'False'): 43 | opt_dict[key] = bool(value) 44 | elif is_float(value): 45 | opt_dict[key] = float(value) 46 | elif is_number(value): 47 | opt_dict[key] = int(value) 48 | else: 49 | opt_dict[key] = str(value) 50 | 51 | # print(opt) 52 | opt_dict['which_epoch'] = 'latest' 53 | opt.save_root = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name) 54 | opt.model_dir = pjoin(opt.save_root, 'model') 55 | opt.meta_dir = pjoin(opt.save_root, 'meta') 56 | 57 | if opt.dataset_name == 't2m': 58 | opt.data_root = './dataset/HumanML3D' 59 | opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs') 60 | opt.text_dir = pjoin(opt.data_root, 'texts') 61 | opt.joints_num = 22 62 | opt.dim_pose = 263 63 | opt.max_motion_length = 196 64 | elif opt.dataset_name == 'kit': 65 | opt.data_root = './dataset/KIT-ML' 66 | opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs') 67 | opt.text_dir = pjoin(opt.data_root, 'texts') 68 | opt.joints_num = 21 69 | opt.dim_pose = 251 70 | opt.max_motion_length = 196 71 | else: 72 | raise KeyError('Dataset not recognized') 73 | 74 | opt.dim_word = 300 75 | opt.num_classes = 200 // opt.unit_length 76 | opt.dim_pos_ohot = len(POS_enumerator) 77 | opt.is_train = False 78 | opt.is_continue = False 79 | opt.device = device 80 | 81 | return opt -------------------------------------------------------------------------------- /diffusion/losses.py: -------------------------------------------------------------------------------- 1 | # This code is based on https://github.com/openai/guided-diffusion 2 | """ 3 | Helpers for various likelihood-based losses. These are ported from the original 4 | Ho et al. diffusion models codebase: 5 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 6 | """ 7 | 8 | import numpy as np 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /data_loaders/humanml/utils/word_vectorizer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pickle 3 | from os.path import join as pjoin 4 | 5 | POS_enumerator = { 6 | 'VERB': 0, 7 | 'NOUN': 1, 8 | 'DET': 2, 9 | 'ADP': 3, 10 | 'NUM': 4, 11 | 'AUX': 5, 12 | 'PRON': 6, 13 | 'ADJ': 7, 14 | 'ADV': 8, 15 | 'Loc_VIP': 9, 16 | 'Body_VIP': 10, 17 | 'Obj_VIP': 11, 18 | 'Act_VIP': 12, 19 | 'Desc_VIP': 13, 20 | 'OTHER': 14, 21 | } 22 | 23 | Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward', 24 | 'up', 'down', 'straight', 'curve') 25 | 26 | Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh') 27 | 28 | Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball') 29 | 30 | Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn', 31 | 'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll', 32 | 'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb') 33 | 34 | Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily', 35 | 'angrily', 'sadly') 36 | 37 | VIP_dict = { 38 | 'Loc_VIP': Loc_list, 39 | 'Body_VIP': Body_list, 40 | 'Obj_VIP': Obj_List, 41 | 'Act_VIP': Act_list, 42 | 'Desc_VIP': Desc_list, 43 | } 44 | 45 | 46 | class WordVectorizer(object): 47 | def __init__(self, meta_root, prefix): 48 | vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix)) 49 | words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb')) 50 | word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb')) 51 | self.word2vec = {w: vectors[word2idx[w]] for w in words} 52 | 53 | def _get_pos_ohot(self, pos): 54 | pos_vec = np.zeros(len(POS_enumerator)) 55 | if pos in POS_enumerator: 56 | pos_vec[POS_enumerator[pos]] = 1 57 | else: 58 | pos_vec[POS_enumerator['OTHER']] = 1 59 | return pos_vec 60 | 61 | def __len__(self): 62 | return len(self.word2vec) 63 | 64 | def __getitem__(self, item): 65 | word, pos = item.split('/') 66 | if word in self.word2vec: 67 | word_vec = self.word2vec[word] 68 | vip_pos = None 69 | for key, values in VIP_dict.items(): 70 | if word in values: 71 | vip_pos = key 72 | break 73 | if vip_pos is not None: 74 | pos_vec = self._get_pos_ohot(vip_pos) 75 | else: 76 | pos_vec = self._get_pos_ohot(pos) 77 | else: 78 | word_vec = self.word2vec['unk'] 79 | pos_vec = self._get_pos_ohot('OTHER') 80 | return word_vec, pos_vec -------------------------------------------------------------------------------- /teach/render/blender/materials.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import bpy 18 | 19 | 20 | def clear_material(material): 21 | if material.node_tree: 22 | material.node_tree.links.clear() 23 | material.node_tree.nodes.clear() 24 | 25 | 26 | def colored_material(r, g, b, a=1, roughness=0.127451): 27 | materials = bpy.data.materials 28 | material = materials.new(name="body") 29 | material.use_nodes = True 30 | clear_material(material) 31 | nodes = material.node_tree.nodes 32 | links = material.node_tree.links 33 | output = nodes.new(type='ShaderNodeOutputMaterial') 34 | diffuse = nodes.new(type='ShaderNodeBsdfDiffuse') 35 | diffuse.inputs["Color"].default_value = (r, g, b, a) 36 | diffuse.inputs["Roughness"].default_value = roughness 37 | links.new(diffuse.outputs['BSDF'], output.inputs['Surface']) 38 | return material 39 | 40 | 41 | def plane_mat(): 42 | materials = bpy.data.materials 43 | material = materials.new(name="plane") 44 | material.use_nodes = True 45 | clear_material(material) 46 | nodes = material.node_tree.nodes 47 | links = material.node_tree.links 48 | output = nodes.new(type='ShaderNodeOutputMaterial') 49 | diffuse = nodes.new(type='ShaderNodeBsdfDiffuse') 50 | checker = nodes.new(type="ShaderNodeTexChecker") 51 | checker.inputs["Scale"].default_value = 1024 52 | checker.inputs["Color1"].default_value = (0.8, 0.8, 0.8, 1) 53 | checker.inputs["Color2"].default_value = (0.3, 0.3, 0.3, 1) 54 | links.new(checker.outputs["Color"], diffuse.inputs['Color']) 55 | links.new(diffuse.outputs['BSDF'], output.inputs['Surface']) 56 | diffuse.inputs["Roughness"].default_value = 0.127451 57 | return material 58 | 59 | 60 | def plane_mat_uni(): 61 | materials = bpy.data.materials 62 | material = materials.new(name="plane_uni") 63 | material.use_nodes = True 64 | clear_material(material) 65 | nodes = material.node_tree.nodes 66 | links = material.node_tree.links 67 | output = nodes.new(type='ShaderNodeOutputMaterial') 68 | diffuse = nodes.new(type='ShaderNodeBsdfDiffuse') 69 | diffuse.inputs["Color"].default_value = (0.8, 0.8, 0.8, 1) 70 | diffuse.inputs["Roughness"].default_value = 0.127451 71 | links.new(diffuse.outputs['BSDF'], output.inputs['Surface']) 72 | return -------------------------------------------------------------------------------- /teach/model/textencoder/distilbert.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import List, Union 18 | import pytorch_lightning as pl 19 | 20 | import torch.nn as nn 21 | import os 22 | 23 | import torch 24 | from torch import Tensor 25 | from torch.distributions.distribution import Distribution 26 | import hydra 27 | 28 | class DistilbertEncoderBase(pl.LightningModule): 29 | def __init__(self, modelpath: str, 30 | finetune: bool = False) -> None: 31 | super().__init__() 32 | 33 | from transformers import AutoTokenizer, AutoModel 34 | from transformers import logging 35 | logging.set_verbosity_error() 36 | # Tokenizer 37 | os.environ["TOKENIZERS_PARALLELISM"] = "false" 38 | # workaround to work from cluster and local 39 | rel_p = modelpath.split('/') 40 | rel_p = rel_p[rel_p.index('deps'):] 41 | rel_p = '/'.join(rel_p) 42 | modelpath = hydra.utils.get_original_cwd() + '/' + rel_p 43 | 44 | self.tokenizer = AutoTokenizer.from_pretrained(modelpath) 45 | 46 | # Text model 47 | self.text_model = AutoModel.from_pretrained(modelpath) 48 | # Don't train the model 49 | if not finetune: 50 | self.text_model.training = False 51 | for p in self.text_model.parameters(): 52 | p.requires_grad = False 53 | 54 | # Then configure the model 55 | self.text_encoded_dim = self.text_model.config.dim 56 | 57 | def train(self, mode: bool = True): 58 | self.training = mode 59 | for module in self.children(): 60 | # Don't put the model in 61 | if module == self.text_model and not self.hparams.finetune: 62 | continue 63 | module.train(mode) 64 | return self 65 | 66 | def get_last_hidden_state(self, texts: List[str], 67 | return_mask: bool = False 68 | ) -> Union[Tensor, tuple[Tensor, Tensor]]: 69 | encoded_inputs = self.tokenizer(texts, return_tensors="pt", padding=True) 70 | output = self.text_model(**encoded_inputs.to(self.text_model.device)) 71 | if not return_mask: 72 | return output.last_hidden_state 73 | return output.last_hidden_state, encoded_inputs.attention_mask.to(dtype=bool) 74 | -------------------------------------------------------------------------------- /teach/launch/prepare.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import sys 18 | import os 19 | # os.environ['HOME']='/home/nathanasiou' 20 | # sys.path.insert(0,'/usr/lib/python3.10/') 21 | # os.environ['PYTHONPATH']='/home/nathanasiou/.venvs/teach/lib/python3.10/site-packages' 22 | import warnings 23 | from pathlib import Path 24 | from omegaconf import OmegaConf 25 | from teach.tools.runid import generate_id 26 | import hydra 27 | 28 | # Local paths 29 | def code_path(path=""): 30 | code_dir = hydra.utils.get_original_cwd() 31 | code_dir = Path(code_dir) 32 | return str(code_dir / path) 33 | 34 | 35 | def working_path(path): 36 | return str(Path(os.getcwd()) / path) 37 | 38 | 39 | # fix the id for this run 40 | ID = generate_id() 41 | def generate_id(): 42 | return ID 43 | 44 | 45 | def get_last_checkpoint(path, ckpt_name="last.ckpt"): 46 | output_dir = Path(hydra.utils.to_absolute_path(path)) 47 | if ckpt_name != 'last.ckpt': 48 | last_ckpt_path = output_dir / "checkpoints" / f'latest-epoch={ckpt_name}.ckpt' 49 | else: 50 | last_ckpt_path = output_dir / "checkpoints" / ckpt_name 51 | return str(last_ckpt_path) 52 | 53 | 54 | def get_samples_folder(path): 55 | output_dir = Path(hydra.utils.to_absolute_path(path)) 56 | samples_path = output_dir / "samples" 57 | return str(samples_path) 58 | 59 | 60 | def get_kitname(load_amass_data: bool, load_with_rot: bool): 61 | if not load_amass_data: 62 | return "kit-mmm-xyz" 63 | if load_amass_data and not load_with_rot: 64 | return "kit-amass-xyz" 65 | if load_amass_data and load_with_rot: 66 | return "kit-amass-rot" 67 | 68 | OmegaConf.register_new_resolver("code_path", code_path) 69 | OmegaConf.register_new_resolver("working_path", working_path) 70 | OmegaConf.register_new_resolver("generate_id", generate_id) 71 | OmegaConf.register_new_resolver("absolute_path", hydra.utils.to_absolute_path) 72 | OmegaConf.register_new_resolver("get_last_checkpoint", get_last_checkpoint) 73 | OmegaConf.register_new_resolver("get_samples_folder", get_samples_folder) 74 | OmegaConf.register_new_resolver("get_kitname", get_kitname) 75 | 76 | 77 | # Remove warnings 78 | warnings.filterwarnings( 79 | "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*" 80 | ) 81 | 82 | warnings.filterwarnings( 83 | "ignore", ".*does not have many workers which may be a bottleneck*" 84 | ) 85 | 86 | warnings.filterwarnings( 87 | "ignore", ".*Our suggested max number of worker in current system is*" 88 | ) 89 | 90 | -------------------------------------------------------------------------------- /teach/model/motiondecoder/actor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn as nn 19 | import numpy as np 20 | import pytorch_lightning as pl 21 | 22 | from typing import List, Optional 23 | from torch import nn, Tensor 24 | 25 | from teach.model.utils import PositionalEncoding 26 | from teach.data.tools import lengths_to_mask 27 | 28 | 29 | class ActorAgnosticDecoder(pl.LightningModule): 30 | def __init__(self, nfeats: int, 31 | latent_dim: int = 256, ff_size: int = 1024, 32 | num_layers: int = 4, num_heads: int = 4, 33 | dropout: float = 0.1, 34 | activation: str = "gelu", **kwargs) -> None: 35 | 36 | super().__init__() 37 | self.save_hyperparameters(logger=False) 38 | 39 | output_feats = nfeats 40 | 41 | 42 | self.sequence_pos_encoding = PositionalEncoding(latent_dim, dropout) 43 | seq_trans_decoder_layer = nn.TransformerDecoderLayer(d_model=latent_dim, 44 | nhead=num_heads, 45 | dim_feedforward=ff_size, 46 | dropout=dropout, 47 | activation=activation) 48 | 49 | self.seqTransDecoder = nn.TransformerDecoder(seq_trans_decoder_layer, 50 | num_layers=num_layers) 51 | 52 | self.final_layer = nn.Linear(latent_dim, output_feats) 53 | 54 | def forward(self, z: Tensor, lengths: List[int]): 55 | mask = lengths_to_mask(lengths, z.device) 56 | latent_dim = z.shape[1] 57 | bs, nframes = mask.shape 58 | nfeats = self.hparams.nfeats 59 | 60 | z = z[None] # sequence of 1 element for the memory 61 | 62 | # Construct time queries 63 | time_queries = torch.zeros(nframes, bs, latent_dim, device=z.device) 64 | time_queries = self.sequence_pos_encoding(time_queries) 65 | 66 | # Pass through the transformer decoder 67 | # with the latent vector for memory 68 | output = self.seqTransDecoder(tgt=time_queries, memory=z, 69 | tgt_key_padding_mask=~mask) 70 | 71 | output = self.final_layer(output) 72 | # zero for padded area 73 | output[~mask.T] = 0 74 | # Pytorch Transformer: [Sequence, Batch size, ...] 75 | feats = output.permute(1, 0, 2) 76 | return feats 77 | -------------------------------------------------------------------------------- /motionclip_data_loaders/pretrain_datamodule.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Microsoft Corporation. 2 | # Licensed under the MIT License. 3 | from pytorch_lightning import LightningDataModule 4 | import torch 5 | from torch.nn import functional as F 6 | from torch.utils.data import DataLoader 7 | from data_loaders.tensors import collate as all_collate 8 | from data_loaders.tensors import t2m_collate 9 | from teach.data.tools.collate import collate_pairs_and_text, collate_datastruct_and_text, collate_contrastive 10 | from teach.data.sampling.base import FrameSampler 11 | 12 | class PretrainDataModule(LightningDataModule): 13 | def __init__( 14 | self, 15 | num_workers: int = 8, 16 | batch_size: int = 32, 17 | *args, 18 | **kwargs, 19 | ): 20 | super().__init__(*args, **kwargs) 21 | self.batch_size = batch_size 22 | self.num_workers = num_workers 23 | # self.dataset = PretrainDataset(root, text_max_len, graph_aug1, graph_aug2) 24 | 25 | from data_loaders.multi_motion.data.dataset import BABEL 26 | datapath = './data/babel/babel-smplh-30fps-male' 27 | framerate = 30 28 | dtype = 'separate_pairs' 29 | from teach.transforms.smpl import SMPLTransform 30 | from teach.transforms.joints2jfeats import Rifke 31 | from teach.transforms.rots2joints import SMPLH 32 | from teach.transforms.rots2rfeats import Globalvelandy 33 | rifke = Rifke(jointstype='mmm', forward_filter=False, 34 | path='./deps/transforms/joints2jfeats/rifke/babel-amass', 35 | normalization=True 36 | ) 37 | smplh = SMPLH(path='./data/smpl_models/smplh', 38 | jointstype='mmm', input_pose_rep='matrix', batch_size=16, gender='male') 39 | globalvelandy = Globalvelandy(canonicalize=True, pose_rep='rot6d', offset=True, 40 | path='./deps/transforms/rots2rfeats/globalvelandy/rot6d/babel-amass', 41 | # path='./deps/transforms/rots2rfeats/globalvelandy/rot6d/babel-amass/separate_pairs', 42 | normalization=True) 43 | transforms = SMPLTransform(rots2rfeats=globalvelandy, rots2joints=smplh, 44 | joints2jfeats=rifke) 45 | frame_sampler = FrameSampler() 46 | frame_sampler.max_len = 256 47 | tiny = False 48 | self.dataset = BABEL(datapath=datapath, framerate=framerate, dtype=dtype, transforms=transforms, tiny=tiny, FrameSampler=frame_sampler) 49 | # datatype = 'separate_pairs' 50 | self.collate = collate_contrastive 51 | 52 | def setup(self, stage: str = None): 53 | self.train_dataset = self.dataset 54 | 55 | def train_dataloader(self): 56 | # loader = DataLoader( 57 | # self.train_dataset, 58 | # batch_size=self.batch_size, 59 | # shuffle=True, 60 | # num_workers=self.num_workers, 61 | # pin_memory=False, 62 | # drop_last=True, 63 | # # persistent_workers = True 64 | # ) 65 | loader = DataLoader( 66 | self.dataset, batch_size=self.batch_size, shuffle=True, 67 | num_workers=self.num_workers, drop_last=True, collate_fn=self.collate 68 | ) 69 | print('len(train_dataset)', len(self.dataset)) 70 | print('len(train_dataloader)', len(loader)) 71 | return loader -------------------------------------------------------------------------------- /teach/model/motionencoder/gru.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn as nn 19 | import numpy as np 20 | import pytorch_lightning as pl 21 | 22 | from typing import List, Optional, Union 23 | from torch import nn, Tensor 24 | from torch.distributions.distribution import Distribution 25 | 26 | from teach.model.utils import PositionalEncoding 27 | from teach.data.tools import lengths_to_mask 28 | 29 | 30 | class GRUEncoder(pl.LightningModule): 31 | def __init__(self, nfeats: int, vae: bool, 32 | latent_dim: int = 256, 33 | num_layers: int = 4, **kwargs) -> None: 34 | super().__init__() 35 | self.save_hyperparameters(logger=False) 36 | 37 | input_feats = nfeats 38 | self.skel_embedding = nn.Linear(input_feats, latent_dim) 39 | 40 | self.gru = nn.GRU(latent_dim, latent_dim, num_layers=num_layers) 41 | 42 | # Action agnostic: only one set of params 43 | if vae: 44 | self.mu = nn.Linear(latent_dim, latent_dim) 45 | self.logvar = nn.Linear(latent_dim, latent_dim) 46 | else: 47 | self.final = nn.Linear(latent_dim, latent_dim) 48 | 49 | def forward(self, features: Tensor, lengths: Optional[List[int]] = None) -> Union[Tensor, Distribution]: 50 | if lengths is None: 51 | lengths = [len(feature) for feature in features] 52 | 53 | device = features.device 54 | 55 | bs, nframes, nfeats = features.shape 56 | mask = lengths_to_mask(lengths, device) 57 | 58 | x = features 59 | # Embed each human poses into latent vectors 60 | x = self.skel_embedding(x) 61 | 62 | # Switch sequence and batch_size because the input of 63 | # Pytorch Transformer is [Sequence, Batch size, ...] 64 | x = x.permute(1, 0, 2) # now it is [nframes, bs, latent_dim] 65 | 66 | # Get all the output of the gru 67 | x = self.gru(x)[0] 68 | 69 | # Put back the batch dimention first 70 | x = x.permute(1, 0, 2) # now it is [nframes, bs, latent_dim] 71 | 72 | # Extract the last valid input 73 | x = x[tuple(torch.stack((torch.arange(bs, device=x.device), 74 | torch.tensor(lengths, device=x.device)-1)))] 75 | 76 | if self.hparams.vae: 77 | mu = self.mu(x) 78 | logvar = self.logvar(x) 79 | std = logvar.exp().pow(0.5) 80 | # https://github.com/kampta/pytorch-distributions/blob/master/gaussian_vae.py 81 | return torch.distributions.Normal(mu, std) 82 | else: 83 | return self.final(x) 84 | -------------------------------------------------------------------------------- /teach/tools/easyconvert.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import teach.tools.geometry as geometry 18 | 19 | def nfeats_of(rottype): 20 | if rottype in ["rotvec", "axisangle"]: 21 | return 3 22 | elif rottype in ["rotquat", "quaternion"]: 23 | return 4 24 | elif rottype in ["rot6d", "6drot", "rotation6d"]: 25 | return 6 26 | elif rottype in ["rotmat"]: 27 | return 9 28 | else: 29 | return TypeError("This rotation type doesn't have features.") 30 | 31 | 32 | def axis_angle_to(newtype, rotations): 33 | if newtype in ["matrix"]: 34 | rotations = geometry.axis_angle_to_matrix(rotations) 35 | return rotations 36 | elif newtype in ["rotmat"]: 37 | rotations = geometry.axis_angle_to_matrix(rotations) 38 | rotations = matrix_to("rotmat", rotations) 39 | return rotations 40 | elif newtype in ["rot6d", "6drot", "rotation6d"]: 41 | rotations = geometry.axis_angle_to_matrix(rotations) 42 | rotations = matrix_to("rot6d", rotations) 43 | return rotations 44 | elif newtype in ["rotquat", "quaternion"]: 45 | rotations = geometry.axis_angle_to_quaternion(rotations) 46 | return rotations 47 | elif newtype in ["rotvec", "axisangle"]: 48 | return rotations 49 | else: 50 | raise NotImplementedError 51 | 52 | 53 | def matrix_to(newtype, rotations): 54 | if newtype in ["matrix"]: 55 | return rotations 56 | if newtype in ["rotmat"]: 57 | rotations = rotations.reshape((*rotations.shape[:-2], 9)) 58 | return rotations 59 | elif newtype in ["rot6d", "6drot", "rotation6d"]: 60 | rotations = geometry.matrix_to_rotation_6d(rotations) 61 | return rotations 62 | elif newtype in ["rotquat", "quaternion"]: 63 | rotations = geometry.matrix_to_quaternion(rotations) 64 | return rotations 65 | elif newtype in ["rotvec", "axisangle"]: 66 | rotations = geometry.matrix_to_axis_angle(rotations) 67 | return rotations 68 | else: 69 | raise NotImplementedError 70 | 71 | 72 | def to_matrix(oldtype, rotations): 73 | if oldtype in ["matrix"]: 74 | return rotations 75 | if oldtype in ["rotmat"]: 76 | rotations = rotations.reshape((*rotations.shape[:-2], 3, 3)) 77 | return rotations 78 | elif oldtype in ["rot6d", "6drot", "rotation6d"]: 79 | rotations = geometry.rotation_6d_to_matrix(rotations) 80 | return rotations 81 | elif oldtype in ["rotquat", "quaternion"]: 82 | rotations = geometry.quaternion_to_matrix(rotations) 83 | return rotations 84 | elif oldtype in ["rotvec", "axisangle"]: 85 | rotations = geometry.axis_angle_to_matrix(rotations) 86 | return rotations 87 | else: 88 | raise NotImplementedError 89 | -------------------------------------------------------------------------------- /utils/model_util.py: -------------------------------------------------------------------------------- 1 | from model.mdm import MDM 2 | from diffusion import gaussian_diffusion as gd 3 | from diffusion.respace import SpacedDiffusion, space_timesteps 4 | 5 | 6 | def load_model_wo_clip(model, state_dict): 7 | missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) 8 | assert len(unexpected_keys) == 0 9 | assert all([k.startswith('clip_model.') for k in missing_keys]) 10 | 11 | 12 | def create_model_and_diffusion(args, data): 13 | # data.dataset = 'babel' 14 | model = MDM(**get_model_args(args, data)) 15 | diffusion = create_gaussian_diffusion(args) 16 | return model, diffusion 17 | 18 | 19 | def get_model_args(args, data): 20 | 21 | # default args 22 | clip_version = 'ViT-B/32' 23 | action_emb = 'tensor' 24 | if args.unconstrained: 25 | cond_mode = 'no_cond' 26 | elif args.dataset in ['kit', 'humanml', 'babel']: 27 | cond_mode = 'text' 28 | else: 29 | cond_mode = 'action' 30 | if hasattr(data.dataset, 'num_actions'): 31 | num_actions = data.dataset.num_actions 32 | else: 33 | num_actions = 1 34 | 35 | # SMPL defaults 36 | data_rep = 'rot6d' 37 | njoints = 25 38 | nfeats = 6 39 | 40 | if args.dataset == 'humanml': 41 | data_rep = 'hml_vec' 42 | njoints = 263 43 | nfeats = 1 44 | elif args.dataset == 'kit': 45 | data_rep = 'hml_vec' 46 | njoints = 251 47 | nfeats = 1 48 | elif args.dataset == 'babel': 49 | data_rep = 'hml_vec' 50 | njoints = 135 51 | nfeats = 1 52 | 53 | return {'modeltype': '', 'njoints': njoints, 'nfeats': nfeats, 'num_actions': num_actions, 54 | 'translation': True, 'pose_rep': 'rot6d', 'glob': True, 'glob_rot': True, 55 | 'latent_dim': args.latent_dim, 'ff_size': 1024, 'num_layers': args.layers, 'num_heads': 4, 56 | 'dropout': 0.1, 'activation': "gelu", 'data_rep': data_rep, 'cond_mode': cond_mode, 57 | 'cond_mask_prob': args.cond_mask_prob, 'action_emb': action_emb, 'arch': args.arch, 58 | 'emb_trans_dec': args.emb_trans_dec, 'clip_version': clip_version, 'dataset': args.dataset, 59 | 'hist_frames': args.hist_frames, 'motion_mask': args.motion_mask} 60 | 61 | 62 | def create_gaussian_diffusion(args): 63 | # default params 64 | predict_xstart = True # we always predict x_start (a.k.a. x0), that's our deal! 65 | steps = 1000 66 | scale_beta = 1. # no scaling 67 | timestep_respacing = '' # can be used for ddim sampling, we don't use it. 68 | learn_sigma = False 69 | rescale_timesteps = False 70 | 71 | betas = gd.get_named_beta_schedule(args.noise_schedule, steps, scale_beta) 72 | loss_type = gd.LossType.MSE 73 | 74 | if not timestep_respacing: 75 | timestep_respacing = [steps] 76 | 77 | return SpacedDiffusion( 78 | use_timesteps=space_timesteps(steps, timestep_respacing), 79 | betas=betas, 80 | model_mean_type=( 81 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 82 | ), 83 | model_var_type=( 84 | ( 85 | gd.ModelVarType.FIXED_LARGE 86 | if not args.sigma_small 87 | else gd.ModelVarType.FIXED_SMALL 88 | ) 89 | if not learn_sigma 90 | else gd.ModelVarType.LEARNED_RANGE 91 | ), 92 | loss_type=loss_type, 93 | rescale_timesteps=rescale_timesteps, 94 | lambda_vel=args.lambda_vel, 95 | lambda_rcxyz=args.lambda_rcxyz, 96 | lambda_fc=args.lambda_fc, 97 | ) -------------------------------------------------------------------------------- /teach/model/losses/actioncompute.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import hydra 20 | import torch 21 | from torch import Tensor 22 | from torch.distributions.distribution import Distribution 23 | from torchmetrics import Metric 24 | 25 | 26 | class ActionComputeLosses(Metric): 27 | def __init__(self, vae: bool, 28 | dist_sync_on_step=False, **kwargs): 29 | super().__init__(dist_sync_on_step=dist_sync_on_step) 30 | 31 | # Save parameters 32 | self.vae = vae 33 | losses = ["recons_text2jfeats", "recons_jfeats2jfeats"] 34 | losses.append("latent_manifold") 35 | 36 | if vae: 37 | kl_losses = [] 38 | kl_losses.extend(["kl_texts", "kl_motion"]) 39 | losses.extend(kl_losses) 40 | 41 | losses.append("total") 42 | 43 | for loss in losses: 44 | self.add_state(loss, default=torch.tensor(0.0), dist_reduce_fx="sum") 45 | self.add_state("count", default=torch.tensor(0), dist_reduce_fx="sum") 46 | self.losses = losses 47 | 48 | # Instantiate loss functions 49 | self._losses_func = {loss: hydra.utils.instantiate(kwargs[loss + "_func"]) 50 | for loss in losses if loss != "total"} 51 | # Save the lambda parameters 52 | self._params = {loss: kwargs[loss] for loss in losses if loss != "total"} 53 | 54 | def update(self, ds_text, ds_motion, ds_ref, 55 | lats_text, lat_motion, diss_text, dis_motion, dis_ref, diss_ref): 56 | total: float = 0.0 57 | 58 | total += self._update_loss("recons_text2jfeats", ds_text.jfeats, ds_ref.jfeats) 59 | total += self._update_loss("recons_jfeats2jfeats", ds_motion.jfeats, ds_ref.jfeats) 60 | 61 | if self.vae: 62 | total += self._update_loss("kl_texts", diss_text, diss_ref) 63 | total += self._update_loss("kl_motion", dis_motion, dis_ref) 64 | 65 | # Average over the actions 66 | lat_text = torch.stack([action_lat.mean(0) for action_lat in lats_text]) 67 | total += self._update_loss("latent_manifold", lat_text, lat_motion) 68 | 69 | self.total += total.detach() 70 | self.count += 1 71 | 72 | return total 73 | 74 | def compute(self, split): 75 | count = getattr(self, "count") 76 | return {loss: getattr(self, loss)/count for loss in self.losses} 77 | 78 | def _update_loss(self, loss: str, outputs, inputs): 79 | # Update the loss 80 | val = self._losses_func[loss](outputs, inputs) 81 | getattr(self, loss).__iadd__(val.detach()) 82 | # Return a weighted sum 83 | return self._params[loss] * val 84 | 85 | def loss2logname(self, loss: str, split: str): 86 | if loss == "total": 87 | log_name = f"{loss}/{split}" 88 | else: 89 | loss_type, name = loss.split("_") 90 | log_name = f"{loss_type}/{name}/{split}" 91 | return log_name 92 | -------------------------------------------------------------------------------- /teach/render/blender/scene.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import bpy 18 | from .materials import plane_mat # noqa 19 | 20 | 21 | def setup_cycles(cycle=True): 22 | bpy.context.scene.render.engine = 'CYCLES' 23 | bpy.data.scenes[0].render.engine = "CYCLES" 24 | bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" 25 | bpy.context.scene.cycles.device = "GPU" 26 | bpy.context.preferences.addons["cycles"].preferences.get_devices() 27 | print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type) 28 | 29 | if cycle: 30 | bpy.context.scene.cycles.use_denoising = True 31 | # added for python versions >3, I have 3.1.2(BLENDER) 32 | if bpy.app.version[0] == 3: 33 | if bpy.context.scene.cycles.device == "GPU": 34 | bpy.context.scene.cycles.tile_size = 256 35 | else: 36 | bpy.context.scene.cycles.tile_size = 32 37 | else: 38 | bpy.context.scene.render.tile_x = 256 39 | bpy.context.scene.render.tile_y = 256 40 | 41 | bpy.context.scene.cycles.samples = 64 42 | # bpy.context.scene.cycles.denoiser = 'OPTIX' 43 | 44 | 45 | # Setup scene 46 | def setup_scene(cycle=True, res='low'): 47 | scene = bpy.data.scenes['Scene'] 48 | assert res in ["ultra", "high", "med", "low"] 49 | if res == "high": 50 | scene.render.resolution_x = 1280 51 | scene.render.resolution_y = 1024 52 | elif res == "med": 53 | scene.render.resolution_x = 1280//2 54 | scene.render.resolution_y = 1024//2 55 | elif res == "low": 56 | scene.render.resolution_x = 1280//4 57 | scene.render.resolution_y = 1024//4 58 | elif res == "ultra": 59 | scene.render.resolution_x = 1280*2 60 | scene.render.resolution_y = 1024*2 61 | world = bpy.data.worlds['World'] 62 | world.use_nodes = True 63 | bg = world.node_tree.nodes['Background'] 64 | bg.inputs[0].default_value[:3] = (1.0, 1.0, 1.0) 65 | bg.inputs[1].default_value = 1.0 66 | 67 | # Remove default cube 68 | if 'Cube' in bpy.data.objects: 69 | bpy.data.objects['Cube'].select_set(True) 70 | bpy.ops.object.delete() 71 | 72 | bpy.ops.object.light_add(type='SUN', align='WORLD', 73 | location=(0, 0, 0), scale=(1, 1, 1)) 74 | bpy.data.objects["Sun"].data.energy = 1.5 75 | 76 | # rotate camera 77 | bpy.ops.object.empty_add(type='PLAIN_AXES', align='WORLD', location=(0, 0, 0), scale=(1, 1, 1)) 78 | bpy.ops.transform.resize(value=(10, 10, 10), orient_type='GLOBAL', orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), 79 | orient_matrix_type='GLOBAL', mirror=True, use_proportional_edit=False, 80 | proportional_edit_falloff='SMOOTH', proportional_size=1, 81 | use_proportional_connected=False, use_proportional_projected=False) 82 | bpy.ops.object.select_all(action='DESELECT') 83 | 84 | setup_cycles(cycle=cycle) 85 | return scene 86 | -------------------------------------------------------------------------------- /teach/logger/wandb_log.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Dict, Optional 18 | from pytorch_lightning.utilities import rank_zero_only 19 | from pytorch_lightning.loggers import WandbLogger as _pl_WandbLogger 20 | import os 21 | from pathlib import Path 22 | import time 23 | 24 | # Fix the step logging 25 | class WandbLogger(_pl_WandbLogger): 26 | @rank_zero_only 27 | def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None: 28 | 29 | assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0" 30 | metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR) 31 | 32 | # if 'epoch' not in metrics: 33 | # self.experiment.log({**metrics, "trainer/global_step": step}, 34 | # step=step) 35 | # else: 36 | wandb_step = int(metrics["epoch"]) 37 | 38 | if step is not None: 39 | self.experiment.log({**metrics, "trainer/global_step": step}, 40 | step=wandb_step) 41 | else: 42 | self.experiment.log(metrics, step=wandb_step) 43 | 44 | @property 45 | def name(self) -> Optional[str]: 46 | """ Override the method because model checkpointing define the path before 47 | the initialization, and in offline mode you can't get the good path 48 | """ 49 | # don't create an experiment if we don't have one 50 | # return self._experiment.project_name() if self._experiment else self._name 51 | return self._wandb_init["project"] 52 | 53 | def symlink_checkpoint(self, code_dir, project, run_id): 54 | # this is the hydra run dir!! see train.yaml 55 | local_project_dir = Path("wandb") / project 56 | local_project_dir.mkdir(parents=True, exist_ok=True) 57 | 58 | # ... but code_dir is the current dir see path.yaml 59 | Path(code_dir) / project / run_id 60 | os.symlink(Path(code_dir) / "wandb" / project / run_id, 61 | local_project_dir / f'{run_id}_{time.strftime("%Y%m%d%H%M%S")}') 62 | # # Creating a another symlink for easy access 63 | os.symlink(Path(code_dir) / "wandb" / project / run_id / "checkpoints", 64 | Path("checkpoints")) 65 | # if it exists an error is spawned which makes sense, but ... 66 | 67 | def symlink_run(self, checkpoint_folder: str): 68 | 69 | code_dir = checkpoint_folder.split("wandb/")[0] 70 | # # local run 71 | local_wandb = Path("wandb/wandb") 72 | local_wandb.mkdir(parents=True, exist_ok=True) 73 | offline_run = self.experiment.dir.split("wandb/wandb/")[1].split("/files")[0] 74 | # # Create the symlink 75 | os.symlink(Path(code_dir) / "wandb/wandb" / offline_run, local_wandb / offline_run) 76 | 77 | def begin(self, code_dir, project, run_id): 78 | self.symlink_checkpoint(code_dir, project, run_id) 79 | 80 | def end(self, checkpoint_folder): 81 | self.symlink_run(checkpoint_folder) 82 | -------------------------------------------------------------------------------- /transforms/joints2jfeats/tools.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn.functional as F 19 | 20 | from teach.info.joints import mmm_joints 21 | 22 | # Get the indexes of particular body part SMPLH case 23 | # Feet 24 | # LM, RM = smplh_joints.index("left_ankle"), smplh_joints.index("right_ankle") 25 | # LF, RF = smplh_joints.index("left_foot"), smplh_joints.index("right_foot") 26 | # # Shoulders 27 | # LS, RS = smplh_joints.index("left_shoulder"), smplh_joints.index("right_shoulder") 28 | # # Hips 29 | # LH, RH = smplh_joints.index("left_hip"), smplh_joints.index("right_hip") 30 | 31 | # Get the indexes of particular body part 32 | # Feet 33 | LM, RM = mmm_joints.index("LMrot"), mmm_joints.index("RMrot") 34 | LF, RF = mmm_joints.index("LF"), mmm_joints.index("RF") 35 | # Shoulders 36 | LS, RS = mmm_joints.index("LS"), mmm_joints.index("RS") 37 | # Hips 38 | LH, RH = mmm_joints.index("LH"), mmm_joints.index("RH") 39 | 40 | def get_forward_direction(poses, jointstype="mmm"): 41 | assert jointstype == 'mmm' 42 | across = poses[..., RH, :] - poses[..., LH, :] + poses[..., RS, :] - poses[..., LS, :] 43 | forward = torch.stack((-across[..., 2], across[..., 0]), axis=-1) 44 | forward = torch.nn.functional.normalize(forward, dim=-1) 45 | return forward 46 | 47 | 48 | def get_floor(poses, jointstype="mmm"): 49 | assert jointstype == 'mmm' 50 | ndim = len(poses.shape) 51 | foot_heights = poses[..., (LM, LF, RM, RF), 1].min(-1).values 52 | floor_height = softmin(foot_heights, softness=0.5, dim=-1) 53 | # changed this thing Mathis version 1.11 pytorch 54 | return floor_height[(ndim - 2) * [None]].transpose(0, -1) 55 | 56 | 57 | def softmax(x, softness=1.0, dim=None): 58 | maxi, mini = x.max(dim=dim).values, x.min(dim=dim).values 59 | return maxi + torch.log(softness + torch.exp(mini - maxi)) 60 | 61 | 62 | def softmin(x, softness=1.0, dim=0): 63 | return -softmax(-x, softness=softness, dim=dim) 64 | 65 | 66 | def gaussian_filter1d(_inputs, sigma, truncate=4.0): 67 | # Code adapted/mixed from scipy library into pytorch 68 | # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L211 69 | # and gaussian kernel 70 | # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L179 71 | # Correspond to mode="nearest" and order = 0 72 | # But works batched 73 | if len(_inputs.shape) == 2: 74 | inputs = _inputs[None] 75 | else: 76 | inputs = _inputs 77 | 78 | sd = float(sigma) 79 | radius = int(truncate * sd + 0.5) 80 | sigma2 = sigma * sigma 81 | x = torch.arange(-radius, radius + 1, device=inputs.device, dtype=inputs.dtype) 82 | phi_x = torch.exp(-0.5 / sigma2 * x ** 2) 83 | phi_x = phi_x / phi_x.sum() 84 | 85 | # Conv1d weights 86 | groups = inputs.shape[-1] 87 | weights = torch.tile(phi_x, (groups, 1, 1)) 88 | inputs = inputs.transpose(-1, -2) 89 | outputs = F.conv1d(inputs, weights, padding="same", groups=groups).transpose(-1, -2) 90 | 91 | return outputs.reshape(_inputs.shape) 92 | -------------------------------------------------------------------------------- /teach/transforms/joints2jfeats/tools.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn.functional as F 19 | 20 | from teach.info.joints import mmm_joints 21 | 22 | # Get the indexes of particular body part SMPLH case 23 | # Feet 24 | # LM, RM = smplh_joints.index("left_ankle"), smplh_joints.index("right_ankle") 25 | # LF, RF = smplh_joints.index("left_foot"), smplh_joints.index("right_foot") 26 | # # Shoulders 27 | # LS, RS = smplh_joints.index("left_shoulder"), smplh_joints.index("right_shoulder") 28 | # # Hips 29 | # LH, RH = smplh_joints.index("left_hip"), smplh_joints.index("right_hip") 30 | 31 | # Get the indexes of particular body part 32 | # Feet 33 | LM, RM = mmm_joints.index("LMrot"), mmm_joints.index("RMrot") 34 | LF, RF = mmm_joints.index("LF"), mmm_joints.index("RF") 35 | # Shoulders 36 | LS, RS = mmm_joints.index("LS"), mmm_joints.index("RS") 37 | # Hips 38 | LH, RH = mmm_joints.index("LH"), mmm_joints.index("RH") 39 | 40 | def get_forward_direction(poses, jointstype="mmm"): 41 | assert jointstype == 'mmm' 42 | across = poses[..., RH, :] - poses[..., LH, :] + poses[..., RS, :] - poses[..., LS, :] 43 | forward = torch.stack((-across[..., 2], across[..., 0]), axis=-1) 44 | forward = torch.nn.functional.normalize(forward, dim=-1) 45 | return forward 46 | 47 | 48 | def get_floor(poses, jointstype="mmm"): 49 | assert jointstype == 'mmm' 50 | ndim = len(poses.shape) 51 | foot_heights = poses[..., (LM, LF, RM, RF), 1].min(-1).values 52 | floor_height = softmin(foot_heights, softness=0.5, dim=-1) 53 | # changed this thing Mathis version 1.11 pytorch 54 | return floor_height[(ndim - 2) * [None]].transpose(0, -1) 55 | 56 | 57 | def softmax(x, softness=1.0, dim=None): 58 | maxi, mini = x.max(dim=dim).values, x.min(dim=dim).values 59 | return maxi + torch.log(softness + torch.exp(mini - maxi)) 60 | 61 | 62 | def softmin(x, softness=1.0, dim=0): 63 | return -softmax(-x, softness=softness, dim=dim) 64 | 65 | 66 | def gaussian_filter1d(_inputs, sigma, truncate=4.0): 67 | # Code adapted/mixed from scipy library into pytorch 68 | # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L211 69 | # and gaussian kernel 70 | # https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/ndimage/filters.py#L179 71 | # Correspond to mode="nearest" and order = 0 72 | # But works batched 73 | if len(_inputs.shape) == 2: 74 | inputs = _inputs[None] 75 | else: 76 | inputs = _inputs 77 | 78 | sd = float(sigma) 79 | radius = int(truncate * sd + 0.5) 80 | sigma2 = sigma * sigma 81 | x = torch.arange(-radius, radius + 1, device=inputs.device, dtype=inputs.dtype) 82 | phi_x = torch.exp(-0.5 / sigma2 * x ** 2) 83 | phi_x = phi_x / phi_x.sum() 84 | 85 | # Conv1d weights 86 | groups = inputs.shape[-1] 87 | weights = torch.tile(phi_x, (groups, 1, 1)) 88 | inputs = inputs.transpose(-1, -2) 89 | outputs = F.conv1d(inputs, weights, padding="same", groups=groups).transpose(-1, -2) 90 | 91 | return outputs.reshape(_inputs.shape) 92 | -------------------------------------------------------------------------------- /model/rotation2xyz.py: -------------------------------------------------------------------------------- 1 | # This code is based on https://github.com/Mathux/ACTOR.git 2 | import torch 3 | import utils.rotation_conversions as geometry 4 | 5 | 6 | from model.smpl import SMPL, JOINTSTYPE_ROOT 7 | # from .get_model import JOINTSTYPES 8 | JOINTSTYPES = ["a2m", "a2mpl", "smpl", "vibe", "vertices"] 9 | 10 | 11 | class Rotation2xyz: 12 | def __init__(self, device, dataset='amass'): 13 | self.device = device 14 | self.dataset = dataset 15 | self.smpl_model = SMPL().eval().to(device) 16 | 17 | def __call__(self, x, mask, pose_rep, translation, glob, 18 | jointstype, vertstrans, betas=None, beta=0, 19 | glob_rot=None, get_rotations_back=False, **kwargs): 20 | if pose_rep == "xyz": 21 | return x 22 | 23 | if mask is None: 24 | mask = torch.ones((x.shape[0], x.shape[-1]), dtype=bool, device=x.device) 25 | 26 | if not glob and glob_rot is None: 27 | raise TypeError("You must specify global rotation if glob is False") 28 | 29 | if jointstype not in JOINTSTYPES: 30 | raise NotImplementedError("This jointstype is not implemented.") 31 | 32 | if translation: 33 | x_translations = x[:, -1, :3] 34 | x_rotations = x[:, :-1] 35 | else: 36 | x_rotations = x 37 | 38 | x_rotations = x_rotations.permute(0, 3, 1, 2) 39 | nsamples, time, njoints, feats = x_rotations.shape 40 | 41 | # Compute rotations (convert only masked sequences output) 42 | if pose_rep == "rotvec": 43 | rotations = geometry.axis_angle_to_matrix(x_rotations[mask]) 44 | elif pose_rep == "rotmat": 45 | rotations = x_rotations[mask].view(-1, njoints, 3, 3) 46 | elif pose_rep == "rotquat": 47 | rotations = geometry.quaternion_to_matrix(x_rotations[mask]) 48 | elif pose_rep == "rot6d": 49 | rotations = geometry.rotation_6d_to_matrix(x_rotations[mask]) 50 | else: 51 | raise NotImplementedError("No geometry for this one.") 52 | 53 | if not glob: 54 | global_orient = torch.tensor(glob_rot, device=x.device) 55 | global_orient = geometry.axis_angle_to_matrix(global_orient).view(1, 1, 3, 3) 56 | global_orient = global_orient.repeat(len(rotations), 1, 1, 1) 57 | else: 58 | global_orient = rotations[:, 0] 59 | rotations = rotations[:, 1:] 60 | 61 | if betas is None: 62 | betas = torch.zeros([rotations.shape[0], self.smpl_model.num_betas], 63 | dtype=rotations.dtype, device=rotations.device) 64 | betas[:, 1] = beta 65 | # import ipdb; ipdb.set_trace() 66 | out = self.smpl_model(body_pose=rotations, global_orient=global_orient, betas=betas) 67 | 68 | # get the desirable joints 69 | joints = out[jointstype] 70 | 71 | x_xyz = torch.empty(nsamples, time, joints.shape[1], 3, device=x.device, dtype=x.dtype) 72 | x_xyz[~mask] = 0 73 | x_xyz[mask] = joints 74 | 75 | x_xyz = x_xyz.permute(0, 2, 3, 1).contiguous() 76 | 77 | # the first translation root at the origin on the prediction 78 | if jointstype != "vertices": 79 | rootindex = JOINTSTYPE_ROOT[jointstype] 80 | x_xyz = x_xyz - x_xyz[:, [rootindex], :, :] 81 | 82 | if translation and vertstrans: 83 | # the first translation root at the origin 84 | x_translations = x_translations - x_translations[:, :, [0]] 85 | 86 | # add the translation to all the joints 87 | x_xyz = x_xyz + x_translations[:, None, :, :] 88 | 89 | if get_rotations_back: 90 | return x_xyz, rotations, global_orient 91 | else: 92 | return x_xyz 93 | -------------------------------------------------------------------------------- /teach/data/sampling/frames.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from typing import Optional 18 | 19 | import numpy as np 20 | from numpy import ndarray as Array 21 | import random 22 | 23 | 24 | def get_frameix_from_data_index(num_frames: int, 25 | max_len: Optional[int], 26 | request_frames: Optional[int], 27 | sampling: str = "conseq", 28 | sampling_step: int = 1) -> Array: 29 | nframes = num_frames 30 | # do not pad small sequences sample from long ones 31 | if request_frames is None or request_frames > nframes: 32 | frame_ix = np.arange(nframes) 33 | else: 34 | # sampling goal: input: ----------- 11 nframes 35 | # o--o--o--o- 4 ninputs 36 | # 37 | # step number is computed like that: [(11-1)/(4-1)] = 3 38 | # [---][---][---][- 39 | # So step = 3, and we take 0 to step*ninputs+1 with steps 40 | # [o--][o--][o--][o-] 41 | # then we can randomly shift the vector 42 | # -[o--][o--][o--]o 43 | # If there are too much frames required 44 | # Nikos: It never gets here now. Should add a pad flag instead of this. 45 | if request_frames > nframes: 46 | fair = False # True 47 | if fair: 48 | # distills redundancy everywhere 49 | choices = np.random.choice(range(nframes), 50 | request_frames, 51 | replace=True) 52 | frame_ix = sorted(choices) 53 | else: 54 | # adding the last frame until done 55 | # Nikos: do not pad 56 | ntoadd = max(0, request_frames - nframes) 57 | lastframe = nframes - 1 58 | padding = lastframe * np.ones(ntoadd, dtype=int) 59 | frame_ix = np.concatenate((np.arange(0, nframes), 60 | padding)) 61 | 62 | elif sampling in ["conseq", "random_conseq"]: 63 | step_max = (nframes - 1) // (request_frames - 1) 64 | if sampling == "conseq": 65 | if sampling_step == -1 or sampling_step * (request_frames - 1) >= nframes: 66 | step = step_max 67 | else: 68 | step = sampling_step 69 | elif sampling == "random_conseq": 70 | step = random.randint(1, step_max) 71 | 72 | lastone = step * (request_frames - 1) 73 | shift_max = nframes - lastone - 1 74 | shift = random.randint(0, max(0, shift_max - 1)) 75 | frame_ix = shift + np.arange(0, lastone + 1, step) 76 | 77 | elif sampling == "random": 78 | choices = np.random.choice(range(nframes), 79 | request_frames, 80 | replace=False) 81 | frame_ix = sorted(choices) 82 | 83 | else: 84 | raise ValueError("Sampling not recognized.") 85 | 86 | return frame_ix 87 | -------------------------------------------------------------------------------- /teach/model/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | from pytorch_lightning import LightningModule 19 | from hydra.utils import instantiate 20 | 21 | from teach.model.metrics import ComputeMetrics 22 | from torchmetrics import MetricCollection 23 | import torch 24 | 25 | class BaseModel(LightningModule): 26 | def __init__(self, *args, **kwargs): 27 | super().__init__(*args, **kwargs) 28 | self.save_hyperparameters(logger=False) 29 | 30 | # Save visuals, one validation step per validation epoch 31 | self.store_examples = {"train": None, 32 | "val": None, 33 | "test": None} 34 | 35 | # Need to define: 36 | # forward 37 | # allsplit_step() 38 | # metrics() 39 | # losses() 40 | 41 | def __post_init__(self): 42 | trainable, nontrainable = 0, 0 43 | for p in self.parameters(): 44 | if p.requires_grad: 45 | trainable += np.prod(p.size()) 46 | else: 47 | nontrainable += np.prod(p.size()) 48 | self.hparams.n_params_trainable = trainable 49 | self.hparams.n_params_nontrainable = nontrainable 50 | 51 | 52 | def training_step(self, batch, batch_idx): 53 | return self.allsplit_step("train", batch, batch_idx) 54 | 55 | def validation_step(self, batch, batch_idx): 56 | return self.allsplit_step("val", batch, batch_idx) 57 | 58 | def test_step(self, batch, batch_idx): 59 | return self.allsplit_step("test", batch, batch_idx) 60 | 61 | def allsplit_epoch_end(self, split: str, outputs): 62 | losses = self.losses[split] 63 | loss_dict = losses.compute(split) 64 | dico = {losses.loss2logname(loss, split): value.item() 65 | for loss, value in loss_dict.items()} 66 | dico.update({"epoch": float(self.trainer.current_epoch), 67 | "step": float(self.trainer.current_epoch)}) 68 | # workaround for LR, assuming 1 optimizer, 1 scheduler, very weak 69 | curr_lr = self.trainer.optimizers[0].param_groups[0]['lr'] 70 | dico.update({'Learning Rate': curr_lr}) 71 | self.log_dict(dico) 72 | 73 | def training_epoch_end(self, outputs): 74 | return self.allsplit_epoch_end("train", outputs) 75 | 76 | def validation_epoch_end(self, outputs): 77 | return self.allsplit_epoch_end("val", outputs) 78 | 79 | def test_epoch_end(self, outputs): 80 | return self.allsplit_epoch_end("test", outputs) 81 | 82 | def configure_optimizers(self): 83 | optim_dict = {} 84 | optimizer = instantiate(self.hparams.optim, params=self.parameters()) 85 | optim_dict['optimizer'] = optimizer 86 | 87 | if self.hparams.lr_scheduler == 'reduceonplateau': 88 | optim_dict['lr_scheduler'] = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, threshold=1e-3) 89 | optim_dict['monitor'] = 'losses/total/train' 90 | elif self.hparams.lr_scheduler == 'steplr': 91 | optim_dict['lr_scheduler'] = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100) 92 | 93 | return optim_dict 94 | -------------------------------------------------------------------------------- /teach/render/blender/floor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import bpy 18 | from .materials import colored_material as get_mat 19 | 20 | 21 | def get_trajectory(data, is_mesh): 22 | if is_mesh: 23 | # mean of the vertices 24 | trajectory = data[:, :, [0, 1]].mean(1) 25 | else: 26 | # get the root joint 27 | trajectory = data[:, 0, [0, 1]] 28 | return trajectory 29 | 30 | 31 | def plot_floor(data, color_alpha=None): 32 | # Create a floor 33 | minx, miny, _ = data.min(axis=(0, 1)) 34 | maxx, maxy, _ = data.max(axis=(0, 1)) 35 | minz = 0 36 | 37 | location = ((maxx + minx)/2, (maxy + miny)/2, 0) 38 | scale = ((maxx - minx)/2, (maxy - miny)/2, 1) 39 | 40 | bpy.ops.mesh.primitive_plane_add(size=2, enter_editmode=False, align='WORLD', location=location, scale=(1, 1, 1)) 41 | 42 | bpy.ops.transform.resize(value=scale, orient_type='GLOBAL', orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), orient_matrix_type='GLOBAL', 43 | constraint_axis=(False, True, False), mirror=True, use_proportional_edit=False, 44 | proportional_edit_falloff='SMOOTH', proportional_size=1, use_proportional_connected=False, 45 | use_proportional_projected=False, release_confirm=True) 46 | obj = bpy.data.objects["Plane"] 47 | obj.name = "SmallPlane" 48 | obj.data.name = "SmallPlane" 49 | if color_alpha is not None: 50 | # obj.active_material = get_mat(0.1, 0.1, 0.1, 1*color_alpha) 51 | obj.active_material = get_mat(65/255, 105/255, 225/255, 1) 52 | else: 53 | obj.active_material = get_mat(0.1, 0.1, 0.1, 1) 54 | location = ((maxx + minx)/2, (maxy + miny)/2, -0.01) 55 | bpy.ops.mesh.primitive_plane_add(size=2, enter_editmode=False, align='WORLD', location=location, scale=(1, 1, 1)) 56 | 57 | bpy.ops.transform.resize(value=[2*x for x in scale], orient_type='GLOBAL', orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), orient_matrix_type='GLOBAL', 58 | constraint_axis=(False, True, False), mirror=True, use_proportional_edit=False, 59 | proportional_edit_falloff='SMOOTH', proportional_size=1, use_proportional_connected=False, 60 | use_proportional_projected=False, release_confirm=True) 61 | 62 | obj = bpy.data.objects["Plane"] 63 | obj.name = "BigPlane" 64 | obj.data.name = "BigPlane" 65 | if color_alpha is not None: 66 | obj.active_material = get_mat(70/255, 130/255, 180/255, 1) 67 | else: 68 | obj.active_material = get_mat(0.2, 0.2, 0.2, 1) 69 | 70 | 71 | def show_traj(coords): 72 | # create the Curve Datablock 73 | curveData = bpy.data.curves.new('myCurve', type='CURVE') 74 | curveData.dimensions = '3D' 75 | curveData.resolution_u = 2 76 | 77 | # map coords to spline 78 | polyline = curveData.splines.new('POLY') 79 | polyline.points.add(len(coords)-1) 80 | for i, coord in enumerate(coords): 81 | x, y = coord 82 | polyline.points[i].co = (x, y, 0.001, 1) 83 | 84 | # create Object 85 | curveOB = bpy.data.objects.new('myCurve', curveData) 86 | curveData.bevel_depth = 0.01 87 | 88 | bpy.context.collection.objects.link(curveOB) 89 | -------------------------------------------------------------------------------- /teach/data/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import pytorch_lightning as pl 18 | from torch.utils.data import DataLoader 19 | 20 | from teach.data.tools.collate import collate_pairs_and_text, collate_datastruct_and_text 21 | 22 | 23 | class BASEDataModule(pl.LightningDataModule): 24 | def __init__(self, batch_size: int, 25 | num_workers: int, 26 | datatype: None): 27 | super().__init__() 28 | 29 | if datatype == 'separate_pairs': 30 | collate_fn = collate_pairs_and_text 31 | else: 32 | collate_fn = collate_datastruct_and_text 33 | 34 | self.dataloader_options = {"batch_size": batch_size, "num_workers": num_workers, 35 | "collate_fn": collate_fn} 36 | # need to be overloaded: 37 | # - self.Dataset 38 | # - self._sample_set => load only a small subset 39 | # There is an helper below (get_sample_set) 40 | # - self.nfeats 41 | # - self.transforms 42 | 43 | self._train_dataset = None 44 | self._val_dataset = None 45 | self._test_dataset = None 46 | 47 | # Optional 48 | self._subset_dataset = None 49 | 50 | def get_sample_set(self, overrides={}): 51 | sample_params = self.hparams.copy() 52 | sample_params.update(overrides) 53 | return self.Dataset(**sample_params) 54 | 55 | @property 56 | def train_dataset(self): 57 | if self._train_dataset is None: 58 | self._train_dataset = self.Dataset(split="train", **self.hparams) 59 | return self._train_dataset 60 | 61 | @property 62 | def val_dataset(self): 63 | if self._val_dataset is None: 64 | self._val_dataset = self.Dataset(split="val", **self.hparams) 65 | return self._val_dataset 66 | 67 | @property 68 | def test_dataset(self): 69 | if self._test_dataset is None: 70 | self._test_dataset = self.Dataset(split="test", **self.hparams) 71 | return self._test_dataset 72 | 73 | @property 74 | def subset_dataset(self): 75 | if self._subset_dataset is None: 76 | self._subset_dataset = self.Dataset(split="subset", **self.hparams) 77 | return self._subset_dataset 78 | 79 | def setup(self, stage=None): 80 | # Use the getter the first time to load the data 81 | if stage in (None, "fit"): 82 | _ = self.train_dataset 83 | _ = self.val_dataset 84 | if stage in (None, "test"): 85 | _ = self.test_dataset 86 | 87 | def train_dataloader(self): 88 | return DataLoader(self.train_dataset, shuffle=True, **self.dataloader_options) 89 | 90 | def predict_dataloader(self): 91 | return DataLoader(self.train_dataset, shuffle=False, **self.dataloader_options) 92 | 93 | def val_dataloader(self): 94 | return DataLoader(self.val_dataset, shuffle=False, **self.dataloader_options) 95 | 96 | def test_dataloader(self): 97 | return DataLoader(self.test_dataset, shuffle=False, **self.dataloader_options) 98 | 99 | def subset_dataloader(self): 100 | return DataLoader(self.subset_dataset, shuffle=False, **self.dataloader_options) 101 | -------------------------------------------------------------------------------- /model/motion_encoder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from transformers import RobertaTokenizer, RobertaModel 3 | from typing import List 4 | import torch 5 | import numpy as np 6 | from torch.nn import TransformerEncoder, TransformerEncoderLayer 7 | import torch.nn.functional as F 8 | import math 9 | from teach.data.tools import lengths_to_mask 10 | 11 | def timestep_embedding(timesteps, dim, max_period=10000): 12 | """ 13 | Create sinusoidal timestep embeddings. 14 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 15 | These may be fractional. 16 | :param dim: the dimension of the output. 17 | :param max_period: controls the minimum frequency of the embeddings. 18 | :return: an [N x dim] Tensor of positional embeddings. 19 | """ 20 | half = dim // 2 21 | freqs = torch.exp( 22 | -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half 23 | ).to(device=timesteps.device) 24 | args = timesteps[:, None].float() * freqs[None] 25 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) 26 | if dim % 2: 27 | embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) 28 | return embedding 29 | 30 | class MotionEncoder(nn.Module): 31 | def __init__( 32 | self, 33 | n_layers: int = 8, 34 | n_heads: int = 4, 35 | input_dim: int = 135, 36 | embed_dim: int = 256, 37 | dim_feedforward: int = 1024, 38 | dropout: float = 0.1, 39 | max_motion_length: int = 512, 40 | ): 41 | super().__init__() 42 | 43 | self.n_layers = n_layers 44 | self.n_heads = n_heads 45 | self.input_dim = input_dim 46 | self.embed_dim = embed_dim 47 | self.dim_feedforward = dim_feedforward 48 | self.dropout = dropout 49 | self.max_motion_length = max_motion_length 50 | 51 | self.input_projection = nn.Linear(input_dim, embed_dim) 52 | # transformer blocks 53 | encoder_layers = TransformerEncoderLayer( 54 | d_model=embed_dim, 55 | nhead=n_heads, 56 | dim_feedforward=dim_feedforward, 57 | dropout=dropout, 58 | batch_first=True, 59 | ) 60 | 61 | # self.motion_length_emb = nn.Embedding(max_motion_length, embed_dim) 62 | self.motion_encoder = TransformerEncoder(encoder_layers, n_layers) 63 | self.CLS_TOKEN = nn.Parameter(torch.randn(1, 1, embed_dim)) 64 | # self.last_projection = nn.Linear(embed_dim, lm_hidden_size) 65 | 66 | def forward(self, motion: torch.Tensor, motion_length: List[int]): 67 | 68 | # motion = motion.permute(0, 2, 1) # (N,C,L) -> (N,T,C) 69 | device = motion.device 70 | mask = lengths_to_mask(motion_length, device) 71 | B, T, C = motion.shape 72 | 73 | cls_token = self.CLS_TOKEN.repeat((B, 1, 1)) 74 | 75 | # motion length embedding 76 | # ml_token = self.motion_length_emb(motion_length)[:, None, :] 77 | 78 | motion_input_proj = self.input_projection(motion) # (N,L,E) 79 | input_tokens = torch.cat( 80 | [cls_token, motion_input_proj], dim=1 81 | ) # [B, T + 2, E] 82 | 83 | pos_enc = timestep_embedding( 84 | torch.arange(T + 1, device=device), self.embed_dim 85 | ).repeat((B, 1, 1)) 86 | input_tokens_pe = input_tokens + pos_enc 87 | 88 | mask_cls_len = torch.ones( 89 | (B, 1), dtype=bool, device=device 90 | ) # extend mask for CLS token and length token 91 | mask_ext = torch.cat([mask_cls_len, mask], dim=1) 92 | 93 | out = self.motion_encoder( 94 | src=input_tokens_pe, src_key_padding_mask=~mask_ext 95 | ) # mask_ext: (N,T+1) 96 | # out = self.last_projection(encoder_out) 97 | 98 | output_dict = { 99 | "pooler_output": out[:, 0, :], 100 | "last_hidden_state": out, 101 | } 102 | return output_dict["pooler_output"] -------------------------------------------------------------------------------- /model/motion_clip.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from model.motion_encoder import MotionEncoder 4 | from model.text_encoder import TextEncoder 5 | import torch.nn.functional as F 6 | import pytorch_lightning as pl 7 | from torch import optim 8 | 9 | class MotionClip(pl.LightningModule): 10 | def __init__( 11 | self, 12 | temperature, 13 | motion_hidden_dim, 14 | text_hidden_dim, 15 | projection_dim, 16 | lr, 17 | weight_decay, 18 | ): 19 | super().__init__() 20 | self.save_hyperparameters() 21 | 22 | self.temperature = temperature 23 | 24 | self.motion_hidden_dim = motion_hidden_dim 25 | self.text_hidden_dim = text_hidden_dim 26 | 27 | self.projection_dim = projection_dim 28 | 29 | self.lr = lr 30 | self.weight_decay = weight_decay 31 | 32 | self.motion_encoder = MotionEncoder() 33 | self.text_encoder = TextEncoder() 34 | 35 | self.motion_proj_head = nn.Sequential( 36 | nn.Linear(self.motion_hidden_dim, self.motion_hidden_dim), 37 | nn.ReLU(inplace=True), 38 | nn.Linear(self.motion_hidden_dim, self.projection_dim) 39 | ) 40 | self.text_proj_head = nn.Sequential( 41 | nn.Linear(self.text_hidden_dim, self.text_hidden_dim), 42 | nn.ReLU(inplace=True), 43 | nn.Linear(self.text_hidden_dim, self.projection_dim) 44 | ) 45 | 46 | def forward(self, features_motion, features_text): 47 | batch_size = features_motion.size(0) 48 | 49 | # normalized features 50 | features_motion = F.normalize(features_motion, dim=-1) 51 | features_text = F.normalize(features_text, dim=-1) 52 | 53 | # cosine similarity as logits 54 | logits_per_motion = features_motion @ features_text.t() / self.temperature 55 | logits_per_text = logits_per_motion.t() 56 | 57 | labels = torch.arange(batch_size, dtype=torch.long, device=self.device) # 大小为B 58 | loss_motion = F.cross_entropy(logits_per_motion, labels) 59 | loss_text = F.cross_entropy(logits_per_text, labels) 60 | loss = (loss_motion + loss_text) / 2 61 | 62 | return logits_per_motion, logits_per_text, loss 63 | 64 | def configure_optimizers(self): 65 | # High lr because of small dataset and small model 66 | optimizer = optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) 67 | return optimizer 68 | 69 | def training_step(self, batch, batch_idx): 70 | 71 | # batch = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in batch.items()} 72 | motion_feats = batch["motion_feats"] 73 | length = batch["length"] 74 | text = batch["text"] 75 | 76 | motion_rep = self.motion_encoder(motion_feats, length) 77 | motion_rep = self.motion_proj_head(motion_rep) 78 | 79 | text_rep = self.text_encoder(text) 80 | text_rep = self.text_proj_head(text_rep) 81 | 82 | _, _, loss = self.forward(motion_rep, text_rep) 83 | 84 | self.log("train_loss", loss) 85 | 86 | return loss 87 | 88 | @staticmethod 89 | def add_model_specific_args(parent_parser): 90 | parser = parent_parser.add_argument_group("MotionClip") 91 | # train mode 92 | parser.add_argument('--temperature', type=float, default=0.1, help='the temperature of NT_XentLoss') 93 | 94 | parser.add_argument('--motion_hidden_dim', type=int, default=256, help='') 95 | 96 | parser.add_argument('--text_hidden_dim', type=int, default=512, help='') 97 | 98 | parser.add_argument('--projection_dim', type=int, default=256) 99 | 100 | # optimization 101 | parser.add_argument('--lr', type=float, default=0.0001, help='optimizer learning rate') 102 | parser.add_argument('--weight_decay', type=float, default=1e-5, help='optimizer weight decay') 103 | 104 | return parent_parser 105 | 106 | -------------------------------------------------------------------------------- /data_loaders/get_data.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import DataLoader 2 | from data_loaders.tensors import collate as all_collate 3 | from data_loaders.tensors import t2m_collate 4 | from teach.data.tools.collate import collate_pairs_and_text, collate_datastruct_and_text, collate_contrastive 5 | from tqdm import tqdm 6 | from teach.data.sampling.base import FrameSampler 7 | 8 | def get_dataset_class(name): 9 | if name == "amass": 10 | from .amass import AMASS 11 | return AMASS 12 | elif name == "uestc": 13 | from .a2m.uestc import UESTC 14 | return UESTC 15 | elif name == "humanact12": 16 | from .a2m.humanact12poses import HumanAct12Poses 17 | return HumanAct12Poses 18 | elif name == "humanml": 19 | from data_loaders.humanml.data.dataset import HumanML3D 20 | return HumanML3D 21 | elif name == "kit": 22 | from data_loaders.humanml.data.dataset import KIT 23 | return KIT 24 | elif name == "babel": 25 | from data_loaders.multi_motion.data.dataset import BABEL 26 | return BABEL 27 | else: 28 | raise ValueError(f'Unsupported dataset name [{name}]') 29 | 30 | def get_collate_fn(name, hml_mode='train'): 31 | if hml_mode == 'gt': 32 | from data_loaders.humanml.data.dataset import collate_fn as t2m_eval_collate 33 | return t2m_eval_collate 34 | if name in ["humanml", "kit"]: 35 | return t2m_collate 36 | else: 37 | return all_collate 38 | 39 | 40 | def get_dataset(name, num_frames, split='train', hml_mode='train'): 41 | DATA = get_dataset_class(name) 42 | if name in ["humanml", "kit"]: 43 | dataset = DATA(split=split, num_frames=num_frames, mode=hml_mode) 44 | else: 45 | dataset = DATA(split=split, num_frames=num_frames) 46 | return dataset 47 | 48 | def get_dataset_loader(args, name, batch_size, num_frames, split='train', hml_mode='train'): 49 | if name != 'babel': 50 | dataset = get_dataset(name, num_frames, split, hml_mode) 51 | collate = get_collate_fn(name, hml_mode) 52 | else: 53 | from data_loaders.multi_motion.data.dataset import BABEL 54 | datapath = './data/babel/babel-smplh-30fps-male' 55 | framerate = 30 56 | dtype = 'separate_pairs' 57 | from teach.transforms.smpl import SMPLTransform 58 | from teach.transforms.joints2jfeats import Rifke 59 | from teach.transforms.rots2joints import SMPLH 60 | from teach.transforms.rots2rfeats import Globalvelandy 61 | rifke = Rifke(jointstype='mmm', forward_filter=False, 62 | path='./deps/transforms/joints2jfeats/rifke/babel-amass', 63 | normalization=True 64 | ) 65 | bs = 128 if split == 'val' else 16 66 | smplh = SMPLH(path='./data/smpl_models/smplh', 67 | jointstype='mmm', input_pose_rep='matrix', batch_size=bs, gender='male') 68 | globalvelandy = Globalvelandy(canonicalize=True, pose_rep='rot6d', offset=True, 69 | path='./deps/transforms/rots2rfeats/globalvelandy/rot6d/babel-amass', 70 | # path='./deps/transforms/rots2rfeats/globalvelandy/rot6d/babel-amass/separate_pairs', 71 | normalization=True) 72 | transforms = SMPLTransform(rots2rfeats=globalvelandy, rots2joints=smplh, 73 | joints2jfeats=rifke) 74 | frame_sampler = FrameSampler() 75 | frame_sampler.max_len = args.max_len 76 | dataset = BABEL(datapath=datapath, framerate=framerate, dtype=dtype, 77 | sampler=frame_sampler, split=split, transforms=transforms, tiny=args.tiny) 78 | datatype = 'separate_pairs' 79 | if hml_mode == 'train': 80 | collate = collate_pairs_and_text 81 | else: 82 | collate = collate_contrastive 83 | 84 | # collate = get_collate_fn(name, hml_mode) 85 | 86 | # batch_size = 1 87 | loader = DataLoader( 88 | dataset, batch_size=batch_size, shuffle=True, 89 | num_workers=8, drop_last=True, collate_fn=collate 90 | ) 91 | # print(len(loader)) 92 | # for batch in tqdm(loader): 93 | # print(batch) 94 | return loader -------------------------------------------------------------------------------- /model/smpl.py: -------------------------------------------------------------------------------- 1 | # This code is based on https://github.com/Mathux/ACTOR.git 2 | import numpy as np 3 | import torch 4 | 5 | import contextlib 6 | 7 | from smplx import SMPLLayer as _SMPLLayer 8 | from smplx.lbs import vertices2joints 9 | 10 | 11 | # action2motion_joints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 21, 24, 38] 12 | # change 0 and 8 13 | action2motion_joints = [8, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 21, 24, 38] 14 | 15 | from utils.config import SMPL_MODEL_PATH, JOINT_REGRESSOR_TRAIN_EXTRA 16 | 17 | JOINTSTYPE_ROOT = {"a2m": 0, # action2motion 18 | "smpl": 0, 19 | "a2mpl": 0, # set(smpl, a2m) 20 | "vibe": 8} # 0 is the 8 position: OP MidHip below 21 | 22 | JOINT_MAP = { 23 | 'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17, 24 | 'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16, 25 | 'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0, 26 | 'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8, 27 | 'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7, 28 | 'OP REye': 25, 'OP LEye': 26, 'OP REar': 27, 29 | 'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30, 30 | 'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34, 31 | 'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45, 32 | 'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7, 33 | 'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17, 34 | 'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20, 35 | 'Neck (LSP)': 47, 'Top of Head (LSP)': 48, 36 | 'Pelvis (MPII)': 49, 'Thorax (MPII)': 50, 37 | 'Spine (H36M)': 51, 'Jaw (H36M)': 52, 38 | 'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26, 39 | 'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27 40 | } 41 | 42 | JOINT_NAMES = [ 43 | 'OP Nose', 'OP Neck', 'OP RShoulder', 44 | 'OP RElbow', 'OP RWrist', 'OP LShoulder', 45 | 'OP LElbow', 'OP LWrist', 'OP MidHip', 46 | 'OP RHip', 'OP RKnee', 'OP RAnkle', 47 | 'OP LHip', 'OP LKnee', 'OP LAnkle', 48 | 'OP REye', 'OP LEye', 'OP REar', 49 | 'OP LEar', 'OP LBigToe', 'OP LSmallToe', 50 | 'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel', 51 | 'Right Ankle', 'Right Knee', 'Right Hip', 52 | 'Left Hip', 'Left Knee', 'Left Ankle', 53 | 'Right Wrist', 'Right Elbow', 'Right Shoulder', 54 | 'Left Shoulder', 'Left Elbow', 'Left Wrist', 55 | 'Neck (LSP)', 'Top of Head (LSP)', 56 | 'Pelvis (MPII)', 'Thorax (MPII)', 57 | 'Spine (H36M)', 'Jaw (H36M)', 58 | 'Head (H36M)', 'Nose', 'Left Eye', 59 | 'Right Eye', 'Left Ear', 'Right Ear' 60 | ] 61 | 62 | 63 | # adapted from VIBE/SPIN to output smpl_joints, vibe joints and action2motion joints 64 | class SMPL(_SMPLLayer): 65 | """ Extension of the official SMPL implementation to support more joints """ 66 | 67 | def __init__(self, model_path=SMPL_MODEL_PATH, **kwargs): 68 | kwargs["model_path"] = model_path 69 | 70 | # remove the verbosity for the 10-shapes beta parameters 71 | with contextlib.redirect_stdout(None): 72 | super(SMPL, self).__init__(**kwargs) 73 | 74 | J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA) 75 | self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32)) 76 | vibe_indexes = np.array([JOINT_MAP[i] for i in JOINT_NAMES]) 77 | a2m_indexes = vibe_indexes[action2motion_joints] 78 | smpl_indexes = np.arange(24) 79 | a2mpl_indexes = np.unique(np.r_[smpl_indexes, a2m_indexes]) 80 | 81 | self.maps = {"vibe": vibe_indexes, 82 | "a2m": a2m_indexes, 83 | "smpl": smpl_indexes, 84 | "a2mpl": a2mpl_indexes} 85 | 86 | def forward(self, *args, **kwargs): 87 | smpl_output = super(SMPL, self).forward(*args, **kwargs) 88 | 89 | extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices) 90 | all_joints = torch.cat([smpl_output.joints, extra_joints], dim=1) 91 | 92 | output = {"vertices": smpl_output.vertices} 93 | 94 | for joinstype, indexes in self.maps.items(): 95 | output[joinstype] = all_joints[:, indexes] 96 | 97 | return output -------------------------------------------------------------------------------- /teach/logger/tools.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | # Taken from pytorch lighting / loggers / base 18 | # Mimic the log hyperparams of wandb 19 | from argparse import Namespace 20 | from typing import Dict, Any, MutableMapping, Union 21 | from omegaconf import DictConfig 22 | 23 | import numpy as np 24 | import torch 25 | 26 | 27 | def _convert_params(params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]: 28 | # in case converting from namespace 29 | if isinstance(params, Namespace): 30 | params = vars(params) 31 | 32 | if params is None: 33 | params = {} 34 | 35 | return params 36 | 37 | 38 | def _flatten_dict(params: Dict[Any, Any], delimiter: str = "/") -> Dict[str, Any]: 39 | """Flatten hierarchical dict, e.g. ``{'a': {'b': 'c'}} -> {'a/b': 'c'}``. 40 | 41 | Args: 42 | params: Dictionary containing the hyperparameters 43 | delimiter: Delimiter to express the hierarchy. Defaults to ``'/'``. 44 | 45 | Returns: 46 | Flattened dict. 47 | 48 | Examples: 49 | >>> _flatten_dict({'a': {'b': 'c'}}) 50 | {'a/b': 'c'} 51 | >>> _flatten_dict({'a': {'b': 123}}) 52 | {'a/b': 123} 53 | >>> _flatten_dict({5: {'a': 123}}) 54 | {'5/a': 123} 55 | """ 56 | 57 | def _dict_generator(input_dict, prefixes=None): 58 | prefixes = prefixes[:] if prefixes else [] 59 | if isinstance(input_dict, MutableMapping): 60 | for key, value in input_dict.items(): 61 | key = str(key) 62 | if isinstance(value, (MutableMapping, Namespace)): 63 | value = vars(value) if isinstance(value, Namespace) else value 64 | yield from _dict_generator(value, prefixes + [key]) 65 | else: 66 | yield prefixes + [key, value if value is not None else str(None)] 67 | else: 68 | yield prefixes + [input_dict if input_dict is None else str(input_dict)] 69 | 70 | return {delimiter.join(keys): val for *keys, val in _dict_generator(params)} 71 | 72 | 73 | def _sanitize_params(params: Dict[str, Any]) -> Dict[str, Any]: 74 | """Returns params with non-primitvies converted to strings for logging. 75 | 76 | >>> params = {"float": 0.3, 77 | ... "int": 1, 78 | ... "string": "abc", 79 | ... "bool": True, 80 | ... "list": [1, 2, 3], 81 | ... "namespace": Namespace(foo=3), 82 | ... "layer": torch.nn.BatchNorm1d} 83 | >>> import pprint 84 | >>> pprint.pprint(_sanitize_params(params)) # doctest: +NORMALIZE_WHITESPACE 85 | {'bool': True, 86 | 'float': 0.3, 87 | 'int': 1, 88 | 'layer': "", 89 | 'list': '[1, 2, 3]', 90 | 'namespace': 'Namespace(foo=3)', 91 | 'string': 'abc'} 92 | """ 93 | for k in params.keys(): 94 | # convert relevant np scalars to python types first (instead of str) 95 | if isinstance(params[k], (np.bool_, np.integer, np.floating)): 96 | params[k] = params[k].item() 97 | elif type(params[k]) not in [bool, int, float, str, torch.Tensor]: 98 | params[k] = str(params[k]) 99 | return params 100 | 101 | 102 | def cfg_to_flatten_config(cfg: DictConfig): 103 | params = _convert_params(cfg) 104 | params = _flatten_dict(params) 105 | params = _sanitize_params(params) 106 | return params 107 | -------------------------------------------------------------------------------- /teach/render/blender/meshes.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | 19 | from .materials import colored_material 20 | 21 | import matplotlib 22 | # green 23 | # GT_SMPL = colored_material(0.009, 0.214, 0.029) 24 | GT_SMPL = colored_material(0.035, 0.415, 0.122) 25 | COLORMAPS = [matplotlib.cm.get_cmap('Blues'), matplotlib.cm.get_cmap('Greys'), matplotlib.cm.get_cmap('Purples'), matplotlib.cm.get_cmap('Greens'), 26 | matplotlib.cm.get_cmap('Reds')] 27 | # blue 28 | # GEN_SMPL = colored_material(0.022, 0.129, 0.439) 29 | # Blues => cmap(0.87) 30 | GEN_SMPL = colored_material(0.035, 0.322, 0.615) 31 | 32 | def prepare_meshes(data, canonicalize=True, always_on_floor=False): 33 | if canonicalize: 34 | print("No canonicalization for now") 35 | 36 | # fix axis 37 | data[..., 1] = - data[..., 1] 38 | data[..., 0] = - data[..., 0] 39 | 40 | # Remove the floor 41 | data[..., 2] -= data[..., 2].min() 42 | 43 | # Put all the body on the floor 44 | if always_on_floor: 45 | data[..., 2] -= data[..., 2].min(1)[:, None] 46 | 47 | return data 48 | 49 | 50 | class Meshes: 51 | def __init__(self, data, *, gt, mode, faces_path, canonicalize, always_on_floor, action_id=0, lengths=None, **kwargs): 52 | # data = prepare_meshes(data, canonicalize=canonicalize, always_on_floor=always_on_floor) 53 | 54 | self.faces = np.load(faces_path) 55 | self.data = data 56 | self.mode = mode 57 | 58 | self.N = len(data) 59 | self.trajectory = data[:, :, [0, 1]].mean(1) 60 | self.lengths = lengths 61 | 62 | self.action_id = action_id 63 | if lengths is None: 64 | if gt: 65 | self.mat = GT_SMPL 66 | else: 67 | self.colormap = COLORMAPS[action_id%len(COLORMAPS)] 68 | # self.mat = colored_material(*matplotlib.cm.get_cmap('Dark2')(0)) 69 | else: 70 | if mode == 'sequence': 71 | self.colormap = COLORMAPS[action_id%len(COLORMAPS)] 72 | elif mode == 'video': 73 | self.mat = matplotlib.cm.get_cmap('Dark2')(action_id) 74 | elif mode == 'frame': 75 | self.mat = matplotlib.cm.get_cmap('Dark2')(action_id) 76 | self.temp_data = data.copy() 77 | self.last_idx = None 78 | 79 | def get_sequence_mat(self, frac): 80 | # cmap = matplotlib.cm.get_cmap('Blues') 81 | if self.mode == 'sequence': 82 | cmap = self.colormap 83 | begin = 0.70 84 | end = 0.90 85 | rgbcolor = cmap(begin + (end-begin)*frac) 86 | mat = colored_material(*rgbcolor) 87 | elif self.mode == 'video': 88 | rgbcolor = self.mat 89 | mat = colored_material(*rgbcolor) 90 | elif self.mode == 'frame': 91 | rgbcolor = self.mat 92 | mat = colored_material(*rgbcolor) 93 | return mat 94 | 95 | def get_root(self, index): 96 | return self.data[index].mean(0) 97 | 98 | def get_mean_root(self): 99 | return self.data.mean((0, 1)) 100 | 101 | def load_in_blender(self, index, mat): 102 | vertices = self.data[index] 103 | faces = self.faces 104 | name = f"{str(index).zfill(4)}_{self.action_id}" 105 | 106 | from .tools import load_numpy_vertices_into_blender 107 | load_numpy_vertices_into_blender(vertices, faces, name, mat) 108 | 109 | return name 110 | 111 | def __len__(self): 112 | return self.N 113 | -------------------------------------------------------------------------------- /teach/model/motiondecoder/metactor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn as nn 19 | import numpy as np 20 | import pytorch_lightning as pl 21 | 22 | from typing import List, Optional 23 | from torch import nn, Tensor 24 | 25 | from teach.model.utils import PositionalEncoding 26 | from teach.data.tools import lengths_to_mask 27 | 28 | 29 | class MetaActorDecoder(pl.LightningModule): 30 | def __init__(self, nfeats: int, 31 | latent_dim: int = 256, ff_size: int = 1024, 32 | num_layers: int = 4, num_heads: int = 4, 33 | dropout: float = 0.1, 34 | mode="posencoding", 35 | activation: str = "gelu", **kwargs) -> None: 36 | 37 | assert mode in ["memory", "posencoding"] 38 | super().__init__() 39 | self.save_hyperparameters(logger=False) 40 | 41 | output_feats = nfeats 42 | 43 | self.sequence_pos_encoding = PositionalEncoding(latent_dim, dropout) 44 | 45 | input_feats = nfeats 46 | self.skel_embedding = nn.Linear(input_feats, latent_dim) 47 | 48 | seq_trans_decoder_layer = nn.TransformerDecoderLayer(d_model=latent_dim, 49 | nhead=num_heads, 50 | dim_feedforward=ff_size, 51 | dropout=dropout, 52 | activation=activation) 53 | 54 | self.seqTransDecoder = nn.TransformerDecoder(seq_trans_decoder_layer, 55 | num_layers=num_layers) 56 | 57 | self.final_layer = nn.Linear(latent_dim, output_feats) 58 | 59 | def forward(self, z: Tensor, first_frames: List[Tensor], lengths: List[int]): 60 | mask = lengths_to_mask(lengths, z.device) 61 | latent_dim = z.shape[1] 62 | bs, nframes = mask.shape 63 | nfeats = self.hparams.nfeats 64 | 65 | fframes_emb = self.skel_embedding(first_frames) 66 | 67 | if self.hparams.mode == "memory": 68 | # two elements in the memory 69 | # the latent vector 70 | # the embedding of the first frame 71 | memory = torch.stack((z, fframes_emb)) 72 | else: 73 | memory = z[None] # sequence of 1 element for the memory 74 | 75 | # Construct time queries 76 | time_queries = torch.zeros(nframes, bs, latent_dim, device=memory.device) 77 | time_queries = self.sequence_pos_encoding(time_queries) 78 | 79 | if self.hparams.mode == "posencoding": 80 | queries = torch.cat((fframes_emb[None], time_queries), axis=0) 81 | # create a bigger mask 82 | fframes_mask = torch.ones((bs, 1), dtype=bool, device=z.device) 83 | mask = torch.cat((fframes_mask, mask), 1) 84 | else: 85 | queries = time_queries 86 | 87 | # Pass through the transformer decoder 88 | # with the latent vector for memory 89 | output = self.seqTransDecoder(tgt=queries, memory=memory, 90 | tgt_key_padding_mask=~mask) 91 | 92 | output = self.final_layer(output) 93 | # zero for padded area 94 | output[~mask.T] = 0 95 | 96 | if self.hparams.mode == "posencoding": 97 | # remove the first one, as it is only to give clues 98 | output = output[1:] 99 | 100 | # Pytorch Transformer: [Sequence, Batch size, ...] 101 | feats = output.permute(1, 0, 2) 102 | return feats 103 | -------------------------------------------------------------------------------- /teach/model/motionencoder/actor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2020 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import torch 18 | import torch.nn as nn 19 | import numpy as np 20 | import pytorch_lightning as pl 21 | 22 | from typing import List, Optional, Union 23 | from torch import nn, Tensor 24 | from torch.distributions.distribution import Distribution 25 | 26 | from teach.model.utils import PositionalEncoding 27 | from teach.data.tools import lengths_to_mask 28 | 29 | 30 | class ActorAgnosticEncoder(pl.LightningModule): 31 | def __init__(self, nfeats: int, vae: bool, 32 | latent_dim: int = 256, ff_size: int = 1024, 33 | num_layers: int = 4, num_heads: int = 4, 34 | dropout: float = 0.1, 35 | activation: str = "gelu", **kwargs) -> None: 36 | super().__init__() 37 | self.save_hyperparameters(logger=False) 38 | 39 | input_feats = nfeats 40 | self.skel_embedding = nn.Linear(input_feats, latent_dim) 41 | 42 | # Action agnostic: only one set of params 43 | if vae: 44 | self.mu_token = nn.Parameter(torch.randn(latent_dim)) 45 | self.logvar_token = nn.Parameter(torch.randn(latent_dim)) 46 | else: 47 | self.emb_token = nn.Parameter(torch.randn(latent_dim)) 48 | 49 | self.sequence_pos_encoding = PositionalEncoding(latent_dim, dropout) 50 | 51 | seq_trans_encoder_layer = nn.TransformerEncoderLayer(d_model=latent_dim, 52 | nhead=num_heads, 53 | dim_feedforward=ff_size, 54 | dropout=dropout, 55 | activation=activation) 56 | 57 | self.seqTransEncoder = nn.TransformerEncoder(seq_trans_encoder_layer, 58 | num_layers=num_layers) 59 | 60 | def forward(self, features: Tensor, lengths: Optional[List[int]] = None) -> Union[Tensor, Distribution]: 61 | if lengths is None: 62 | lengths = [len(feature) for feature in features] 63 | 64 | device = features.device 65 | 66 | bs, nframes, nfeats = features.shape 67 | mask = lengths_to_mask(lengths, device) 68 | 69 | x = features 70 | # Embed each human poses into latent vectors 71 | x = self.skel_embedding(x) 72 | 73 | # Switch sequence and batch_size because the input of 74 | # Pytorch Transformer is [Sequence, Batch size, ...] 75 | x = x.permute(1, 0, 2) # now it is [nframes, bs, latent_dim] 76 | 77 | # Each batch has its own set of tokens 78 | if self.hparams.vae: 79 | mu_token = torch.tile(self.mu_token, (bs,)).reshape(bs, -1) 80 | logvar_token = torch.tile(self.logvar_token, (bs,)).reshape(bs, -1) 81 | 82 | # adding the distribution tokens for all sequences 83 | xseq = torch.cat((mu_token[None], logvar_token[None], x), 0) 84 | 85 | # create a bigger mask, to allow attend to mu and logvar 86 | token_mask = torch.ones((bs, 2), dtype=bool, device=x.device) 87 | aug_mask = torch.cat((token_mask, mask), 1) 88 | else: 89 | emb_token = torch.tile(self.emb_token, (bs,)).reshape(bs, -1) 90 | 91 | # adding the embedding token for all sequences 92 | xseq = torch.cat((emb_token[None], x), 0) 93 | 94 | # create a bigger mask, to allow attend to emb 95 | token_mask = torch.ones((bs, 1), dtype=bool, device=x.device) 96 | aug_mask = torch.cat((token_mask, mask), 1) 97 | 98 | # add positional encoding 99 | xseq = self.sequence_pos_encoding(xseq) 100 | final = self.seqTransEncoder(xseq, src_key_padding_mask=~aug_mask) 101 | 102 | if self.hparams.vae: 103 | mu, logvar = final[0], final[1] 104 | std = logvar.exp().pow(0.5) 105 | # https://github.com/kampta/pytorch-distributions/blob/master/gaussian_vae.py 106 | dist = torch.distributions.Normal(mu, std) 107 | return dist 108 | else: 109 | return final[0] 110 | --------------------------------------------------------------------------------