diff --git a/VQ-Trans/.gitignore b/VQ-Trans/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..25e4de635b2ddfcad42dbb14ccb4a9d57a0f5882 --- /dev/null +++ b/VQ-Trans/.gitignore @@ -0,0 +1,147 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +.vscode +dataset/dataset_TM_train_cb1_temp.py +train_gpt_cnn_temp.py +train_gpt_cnn_mask.py +start.sh +start_eval.sh +config.json +output_GPT_Final +output_vqfinal +output_transformer +glove +checkpoints +dataset/HumanML3D +dataset/KIT-ML +output +matrix_multi.py +body_models \ No newline at end of file diff --git a/VQ-Trans/GPT_eval_multi.py b/VQ-Trans/GPT_eval_multi.py new file mode 100644 index 0000000000000000000000000000000000000000..b5e3ebcb1199e42cf16748e60863b554a0046f00 --- /dev/null +++ b/VQ-Trans/GPT_eval_multi.py @@ -0,0 +1,121 @@ +import os +import torch +import numpy as np +from torch.utils.tensorboard import SummaryWriter +import json +import clip + +import options.option_transformer as option_trans +import models.vqvae as vqvae +import utils.utils_model as utils_model +import utils.eval_trans as eval_trans +from dataset import dataset_TM_eval +import models.t2m_trans as trans +from options.get_eval_option import get_opt +from models.evaluator_wrapper import EvaluatorModelWrapper +import warnings +warnings.filterwarnings('ignore') + +##### ---- Exp dirs ---- ##### +args = option_trans.get_args_parser() +torch.manual_seed(args.seed) + +args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}') +os.makedirs(args.out_dir, exist_ok = True) + +##### ---- Logger ---- ##### +logger = utils_model.get_logger(args.out_dir) +writer = SummaryWriter(args.out_dir) +logger.info(json.dumps(vars(args), indent=4, sort_keys=True)) + +from utils.word_vectorizer import WordVectorizer +w_vectorizer = WordVectorizer('./glove', 'our_vab') +val_loader = dataset_TM_eval.DATALoader(args.dataname, True, 32, w_vectorizer) + +dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataname == 'kit' else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt' + +wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda')) +eval_wrapper = EvaluatorModelWrapper(wrapper_opt) + +##### ---- Network ---- ##### + +## load clip model and datasets +clip_model, clip_preprocess = clip.load("ViT-B/32", device=torch.device('cuda'), jit=False, download_root='/apdcephfs_cq2/share_1290939/maelyszhang/.cache/clip') # Must set jit=False for training +clip.model.convert_weights(clip_model) # Actually this line is unnecessary since clip by default already on float16 +clip_model.eval() +for p in clip_model.parameters(): + p.requires_grad = False + +net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers + args.nb_code, + args.code_dim, + args.output_emb_width, + args.down_t, + args.stride_t, + args.width, + args.depth, + args.dilation_growth_rate) + + +trans_encoder = trans.Text2Motion_Transformer(num_vq=args.nb_code, + embed_dim=args.embed_dim_gpt, + clip_dim=args.clip_dim, + block_size=args.block_size, + num_layers=args.num_layers, + n_head=args.n_head_gpt, + drop_out_rate=args.drop_out_rate, + fc_rate=args.ff_rate) + + +print ('loading checkpoint from {}'.format(args.resume_pth)) +ckpt = torch.load(args.resume_pth, map_location='cpu') +net.load_state_dict(ckpt['net'], strict=True) +net.eval() +net.cuda() + +if args.resume_trans is not None: + print ('loading transformer checkpoint from {}'.format(args.resume_trans)) + ckpt = torch.load(args.resume_trans, map_location='cpu') + trans_encoder.load_state_dict(ckpt['trans'], strict=True) +trans_encoder.train() +trans_encoder.cuda() + + +fid = [] +div = [] +top1 = [] +top2 = [] +top3 = [] +matching = [] +multi = [] +repeat_time = 20 + + +for i in range(repeat_time): + best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, best_multi, writer, logger = eval_trans.evaluation_transformer_test(args.out_dir, val_loader, net, trans_encoder, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, best_multi=0, clip_model=clip_model, eval_wrapper=eval_wrapper, draw=False, savegif=False, save=False, savenpy=(i==0)) + fid.append(best_fid) + div.append(best_div) + top1.append(best_top1) + top2.append(best_top2) + top3.append(best_top3) + matching.append(best_matching) + multi.append(best_multi) + +print('final result:') +print('fid: ', sum(fid)/repeat_time) +print('div: ', sum(div)/repeat_time) +print('top1: ', sum(top1)/repeat_time) +print('top2: ', sum(top2)/repeat_time) +print('top3: ', sum(top3)/repeat_time) +print('matching: ', sum(matching)/repeat_time) +print('multi: ', sum(multi)/repeat_time) + +fid = np.array(fid) +div = np.array(div) +top1 = np.array(top1) +top2 = np.array(top2) +top3 = np.array(top3) +matching = np.array(matching) +multi = np.array(multi) +msg_final = f"FID. {np.mean(fid):.3f}, conf. {np.std(fid)*1.96/np.sqrt(repeat_time):.3f}, Diversity. {np.mean(div):.3f}, conf. {np.std(div)*1.96/np.sqrt(repeat_time):.3f}, TOP1. {np.mean(top1):.3f}, conf. {np.std(top1)*1.96/np.sqrt(repeat_time):.3f}, TOP2. {np.mean(top2):.3f}, conf. {np.std(top2)*1.96/np.sqrt(repeat_time):.3f}, TOP3. {np.mean(top3):.3f}, conf. {np.std(top3)*1.96/np.sqrt(repeat_time):.3f}, Matching. {np.mean(matching):.3f}, conf. {np.std(matching)*1.96/np.sqrt(repeat_time):.3f}, Multi. {np.mean(multi):.3f}, conf. {np.std(multi)*1.96/np.sqrt(repeat_time):.3f}" +logger.info(msg_final) \ No newline at end of file diff --git a/VQ-Trans/README.md b/VQ-Trans/README.md new file mode 100644 index 0000000000000000000000000000000000000000..547a1d4b52a5c76f0f86c641557f99d0688c0ffd --- /dev/null +++ b/VQ-Trans/README.md @@ -0,0 +1,400 @@ +# Motion VQ-Trans +Pytorch implementation of paper "Generating Human Motion from Textual Descriptions with High Quality Discrete Representation" + + +[[Notebook Demo]](https://colab.research.google.com/drive/1tAHlmcpKcjg_zZrqKku7AfpqdVAIFrF8?usp=sharing) + + +![teaser](img/Teaser.png) + +If our project is helpful for your research, please consider citing : (todo) +``` +@inproceedings{shen2020ransac, + title={RANSAC-Flow: generic two-stage image alignment}, + author={Shen, Xi and Darmon, Fran{\c{c}}ois and Efros, Alexei A and Aubry, Mathieu}, + booktitle={16th European Conference on Computer Vision} + year={2020} + } +``` + + +## Table of Content +* [1. Visual Results](#1-visual-results) +* [2. Installation](#2-installation) +* [3. Quick Start](#3-quick-start) +* [4. Train](#4-train) +* [5. Evaluation](#5-evaluation) +* [6. Motion Render](#6-motion-render) +* [7. Acknowledgement](#7-acknowledgement) +* [8. ChangLog](#8-changlog) + + + + +## 1. Visual Results (More results can be found in our project page (todo)) + +![visualization](img/ALLvis.png) + + +## 2. Installation + +### 2.1. Environment + + + +Our model can be learnt in a **single GPU V100-32G** + +```bash +conda env create -f environment.yml +conda activate VQTrans +``` + +The code was tested on Python 3.8 and PyTorch 1.8.1. + + +### 2.2. Dependencies + +```bash +bash dataset/prepare/download_glove.sh +``` + + +### 2.3. Datasets + + +We are using two 3D human motion-language dataset: HumanML3D and KIT-ML. For both datasets, you could find the details as well as download link [[here]](https://github.com/EricGuo5513/HumanML3D). + +Take HumanML3D for an example, the file directory should look like this: +``` +./dataset/HumanML3D/ +├── new_joint_vecs/ +├── texts/ +├── Mean.npy # same as in [HumanML3D](https://github.com/EricGuo5513/HumanML3D) +├── Std.npy # same as in [HumanML3D](https://github.com/EricGuo5513/HumanML3D) +├── train.txt +├── val.txt +├── test.txt +├── train_val.txt +└──all.txt +``` + + +### 2.4. Motion & text feature extractors: + +We use the same extractors provided by [t2m](https://github.com/EricGuo5513/text-to-motion) to evaluate our generated motions. Please download the extractors. + +```bash +bash dataset/prepare/download_extractor.sh +``` + +### 2.5. Pre-trained models + +The pretrained model files will be stored in the 'pretrained' folder: +```bash +bash dataset/prepare/download_model.sh +``` + + + +### 2.6. Render motion (optional) + +If you want to render the generated motion, you need to install: + +```bash +sudo sh dataset/prepare/download_smpl.sh +conda install -c menpo osmesa +conda install h5py +conda install -c conda-forge shapely pyrender trimesh mapbox_earcut +``` + + + +## 3. Quick Start + +A quick start guide of how to use our code is available in [demo.ipynb](https://colab.research.google.com/drive/1tAHlmcpKcjg_zZrqKku7AfpqdVAIFrF8?usp=sharing) + +

+demo +

+ + +## 4. Train + +Note that, for kit dataset, just need to set '--dataname kit'. + +### 4.1. VQ-VAE + +The results are saved in the folder output_vqfinal. + +
+ +VQ training + + +```bash +python3 train_vq.py \ +--batch-size 256 \ +--lr 2e-4 \ +--total-iter 300000 \ +--lr-scheduler 200000 \ +--nb-code 512 \ +--down-t 2 \ +--depth 3 \ +--dilation-growth-rate 3 \ +--out-dir output \ +--dataname t2m \ +--vq-act relu \ +--quantizer ema_reset \ +--loss-vel 0.5 \ +--recons-loss l1_smooth \ +--exp-name VQVAE +``` + +
+ +### 4.2. Motion-Transformer + +The results are saved in the folder output_transformer. + +
+ +MoTrans training + + +```bash +python3 train_t2m_trans.py \ +--exp-name VQTransformer \ +--batch-size 128 \ +--num-layers 9 \ +--embed-dim-gpt 1024 \ +--nb-code 512 \ +--n-head-gpt 16 \ +--block-size 51 \ +--ff-rate 4 \ +--drop-out-rate 0.1 \ +--resume-pth output/VQVAE/net_last.pth \ +--vq-name VQVAE \ +--out-dir output \ +--total-iter 300000 \ +--lr-scheduler 150000 \ +--lr 0.0001 \ +--dataname t2m \ +--down-t 2 \ +--depth 3 \ +--quantizer ema_reset \ +--eval-iter 10000 \ +--pkeep 0.5 \ +--dilation-growth-rate 3 \ +--vq-act relu +``` + +
+ +## 5. Evaluation + +### 5.1. VQ-VAE +
+ +VQ eval + + +```bash +python3 VQ_eval.py \ +--batch-size 256 \ +--lr 2e-4 \ +--total-iter 300000 \ +--lr-scheduler 200000 \ +--nb-code 512 \ +--down-t 2 \ +--depth 3 \ +--dilation-growth-rate 3 \ +--out-dir output \ +--dataname t2m \ +--vq-act relu \ +--quantizer ema_reset \ +--loss-vel 0.5 \ +--recons-loss l1_smooth \ +--exp-name TEST_VQVAE \ +--resume-pth output/VQVAE/net_last.pth +``` + +
+ +### 5.2. Motion-Transformer + +
+ +MoTrans eval + + +```bash +python3 GPT_eval_multi.py \ +--exp-name TEST_VQTransformer \ +--batch-size 128 \ +--num-layers 9 \ +--embed-dim-gpt 1024 \ +--nb-code 512 \ +--n-head-gpt 16 \ +--block-size 51 \ +--ff-rate 4 \ +--drop-out-rate 0.1 \ +--resume-pth output/VQVAE/net_last.pth \ +--vq-name VQVAE \ +--out-dir output \ +--total-iter 300000 \ +--lr-scheduler 150000 \ +--lr 0.0001 \ +--dataname t2m \ +--down-t 2 \ +--depth 3 \ +--quantizer ema_reset \ +--eval-iter 10000 \ +--pkeep 0.5 \ +--dilation-growth-rate 3 \ +--vq-act relu \ +--resume-gpt output/VQTransformer/net_best_fid.pth +``` + +
+ + +## 6. Motion Render + +
+ +Motion Render + + +You should input the npy folder address and the motion names. Here is an example: + +```bash +python3 render_final.py --filedir output/TEST_VQTransformer/ --motion-list 000019 005485 +``` + +
+ +### 7. Acknowledgement + +We appreciate helps from : + +* Public code like [text-to-motion](https://github.com/EricGuo5513/text-to-motion), [TM2T](https://github.com/EricGuo5513/TM2T) etc. + +### 8. ChangLog + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/VQ-Trans/VQ_eval.py b/VQ-Trans/VQ_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..f1b7f269e344f730797eba13a45c9672f323b9f5 --- /dev/null +++ b/VQ-Trans/VQ_eval.py @@ -0,0 +1,95 @@ +import os +import json + +import torch +from torch.utils.tensorboard import SummaryWriter +import numpy as np +import models.vqvae as vqvae +import options.option_vq as option_vq +import utils.utils_model as utils_model +from dataset import dataset_TM_eval +import utils.eval_trans as eval_trans +from options.get_eval_option import get_opt +from models.evaluator_wrapper import EvaluatorModelWrapper +import warnings +warnings.filterwarnings('ignore') +import numpy as np +##### ---- Exp dirs ---- ##### +args = option_vq.get_args_parser() +torch.manual_seed(args.seed) + +args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}') +os.makedirs(args.out_dir, exist_ok = True) + +##### ---- Logger ---- ##### +logger = utils_model.get_logger(args.out_dir) +writer = SummaryWriter(args.out_dir) +logger.info(json.dumps(vars(args), indent=4, sort_keys=True)) + + +from utils.word_vectorizer import WordVectorizer +w_vectorizer = WordVectorizer('./glove', 'our_vab') + + +dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataname == 'kit' else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt' + +wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda')) +eval_wrapper = EvaluatorModelWrapper(wrapper_opt) + + +##### ---- Dataloader ---- ##### +args.nb_joints = 21 if args.dataname == 'kit' else 22 + +val_loader = dataset_TM_eval.DATALoader(args.dataname, True, 32, w_vectorizer, unit_length=2**args.down_t) + +##### ---- Network ---- ##### +net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers + args.nb_code, + args.code_dim, + args.output_emb_width, + args.down_t, + args.stride_t, + args.width, + args.depth, + args.dilation_growth_rate, + args.vq_act, + args.vq_norm) + +if args.resume_pth : + logger.info('loading checkpoint from {}'.format(args.resume_pth)) + ckpt = torch.load(args.resume_pth, map_location='cpu') + net.load_state_dict(ckpt['net'], strict=True) +net.train() +net.cuda() + +fid = [] +div = [] +top1 = [] +top2 = [] +top3 = [] +matching = [] +repeat_time = 20 +for i in range(repeat_time): + best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, eval_wrapper=eval_wrapper, draw=False, save=False, savenpy=(i==0)) + fid.append(best_fid) + div.append(best_div) + top1.append(best_top1) + top2.append(best_top2) + top3.append(best_top3) + matching.append(best_matching) +print('final result:') +print('fid: ', sum(fid)/repeat_time) +print('div: ', sum(div)/repeat_time) +print('top1: ', sum(top1)/repeat_time) +print('top2: ', sum(top2)/repeat_time) +print('top3: ', sum(top3)/repeat_time) +print('matching: ', sum(matching)/repeat_time) + +fid = np.array(fid) +div = np.array(div) +top1 = np.array(top1) +top2 = np.array(top2) +top3 = np.array(top3) +matching = np.array(matching) +msg_final = f"FID. {np.mean(fid):.3f}, conf. {np.std(fid)*1.96/np.sqrt(repeat_time):.3f}, Diversity. {np.mean(div):.3f}, conf. {np.std(div)*1.96/np.sqrt(repeat_time):.3f}, TOP1. {np.mean(top1):.3f}, conf. {np.std(top1)*1.96/np.sqrt(repeat_time):.3f}, TOP2. {np.mean(top2):.3f}, conf. {np.std(top2)*1.96/np.sqrt(repeat_time):.3f}, TOP3. {np.mean(top3):.3f}, conf. {np.std(top3)*1.96/np.sqrt(repeat_time):.3f}, Matching. {np.mean(matching):.3f}, conf. {np.std(matching)*1.96/np.sqrt(repeat_time):.3f}" +logger.info(msg_final) \ No newline at end of file diff --git a/VQ-Trans/ViT-B-32.pt b/VQ-Trans/ViT-B-32.pt new file mode 100644 index 0000000000000000000000000000000000000000..06a4dea8587eb4948a3221b1e1b2e2475e0e203b --- /dev/null +++ b/VQ-Trans/ViT-B-32.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af +size 353976522 diff --git a/VQ-Trans/dataset/dataset_TM_eval.py b/VQ-Trans/dataset/dataset_TM_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..576a53b7dabd8095bed59dcc86199e30f2798838 --- /dev/null +++ b/VQ-Trans/dataset/dataset_TM_eval.py @@ -0,0 +1,217 @@ +import torch +from torch.utils import data +import numpy as np +from os.path import join as pjoin +import random +import codecs as cs +from tqdm import tqdm + +import utils.paramUtil as paramUtil +from torch.utils.data._utils.collate import default_collate + + +def collate_fn(batch): + batch.sort(key=lambda x: x[3], reverse=True) + return default_collate(batch) + + +'''For use of training text-2-motion generative model''' +class Text2MotionDataset(data.Dataset): + def __init__(self, dataset_name, is_test, w_vectorizer, feat_bias = 5, max_text_len = 20, unit_length = 4): + + self.max_length = 20 + self.pointer = 0 + self.dataset_name = dataset_name + self.is_test = is_test + self.max_text_len = max_text_len + self.unit_length = unit_length + self.w_vectorizer = w_vectorizer + if dataset_name == 't2m': + self.data_root = './dataset/HumanML3D' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 22 + radius = 4 + fps = 20 + self.max_motion_length = 196 + dim_pose = 263 + kinematic_chain = paramUtil.t2m_kinematic_chain + self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + elif dataset_name == 'kit': + self.data_root = './dataset/KIT-ML' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 21 + radius = 240 * 8 + fps = 12.5 + dim_pose = 251 + self.max_motion_length = 196 + kinematic_chain = paramUtil.kit_kinematic_chain + self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + + mean = np.load(pjoin(self.meta_dir, 'mean.npy')) + std = np.load(pjoin(self.meta_dir, 'std.npy')) + + if is_test: + split_file = pjoin(self.data_root, 'test.txt') + else: + split_file = pjoin(self.data_root, 'val.txt') + + min_motion_len = 40 if self.dataset_name =='t2m' else 24 + # min_motion_len = 64 + + joints_num = self.joints_num + + data_dict = {} + id_list = [] + with cs.open(split_file, 'r') as f: + for line in f.readlines(): + id_list.append(line.strip()) + + new_name_list = [] + length_list = [] + for name in tqdm(id_list): + try: + motion = np.load(pjoin(self.motion_dir, name + '.npy')) + if (len(motion)) < min_motion_len or (len(motion) >= 200): + continue + text_data = [] + flag = False + with cs.open(pjoin(self.text_dir, name + '.txt')) as f: + for line in f.readlines(): + text_dict = {} + line_split = line.strip().split('#') + caption = line_split[0] + tokens = line_split[1].split(' ') + f_tag = float(line_split[2]) + to_tag = float(line_split[3]) + f_tag = 0.0 if np.isnan(f_tag) else f_tag + to_tag = 0.0 if np.isnan(to_tag) else to_tag + + text_dict['caption'] = caption + text_dict['tokens'] = tokens + if f_tag == 0.0 and to_tag == 0.0: + flag = True + text_data.append(text_dict) + else: + try: + n_motion = motion[int(f_tag*fps) : int(to_tag*fps)] + if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200): + continue + new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name + while new_name in data_dict: + new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name + data_dict[new_name] = {'motion': n_motion, + 'length': len(n_motion), + 'text':[text_dict]} + new_name_list.append(new_name) + length_list.append(len(n_motion)) + except: + print(line_split) + print(line_split[2], line_split[3], f_tag, to_tag, name) + # break + + if flag: + data_dict[name] = {'motion': motion, + 'length': len(motion), + 'text': text_data} + new_name_list.append(name) + length_list.append(len(motion)) + except Exception as e: + # print(e) + pass + + name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1])) + self.mean = mean + self.std = std + self.length_arr = np.array(length_list) + self.data_dict = data_dict + self.name_list = name_list + self.reset_max_len(self.max_length) + + def reset_max_len(self, length): + assert length <= self.max_motion_length + self.pointer = np.searchsorted(self.length_arr, length) + print("Pointer Pointing at %d"%self.pointer) + self.max_length = length + + def inv_transform(self, data): + return data * self.std + self.mean + + def forward_transform(self, data): + return (data - self.mean) / self.std + + def __len__(self): + return len(self.data_dict) - self.pointer + + def __getitem__(self, item): + idx = self.pointer + item + name = self.name_list[idx] + data = self.data_dict[name] + # data = self.data_dict[self.name_list[idx]] + motion, m_length, text_list = data['motion'], data['length'], data['text'] + # Randomly select a caption + text_data = random.choice(text_list) + caption, tokens = text_data['caption'], text_data['tokens'] + + if len(tokens) < self.max_text_len: + # pad with "unk" + tokens = ['sos/OTHER'] + tokens + ['eos/OTHER'] + sent_len = len(tokens) + tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - sent_len) + else: + # crop + tokens = tokens[:self.max_text_len] + tokens = ['sos/OTHER'] + tokens + ['eos/OTHER'] + sent_len = len(tokens) + pos_one_hots = [] + word_embeddings = [] + for token in tokens: + word_emb, pos_oh = self.w_vectorizer[token] + pos_one_hots.append(pos_oh[None, :]) + word_embeddings.append(word_emb[None, :]) + pos_one_hots = np.concatenate(pos_one_hots, axis=0) + word_embeddings = np.concatenate(word_embeddings, axis=0) + + if self.unit_length < 10: + coin2 = np.random.choice(['single', 'single', 'double']) + else: + coin2 = 'single' + + if coin2 == 'double': + m_length = (m_length // self.unit_length - 1) * self.unit_length + elif coin2 == 'single': + m_length = (m_length // self.unit_length) * self.unit_length + idx = random.randint(0, len(motion) - m_length) + motion = motion[idx:idx+m_length] + + "Z Normalization" + motion = (motion - self.mean) / self.std + + if m_length < self.max_motion_length: + motion = np.concatenate([motion, + np.zeros((self.max_motion_length - m_length, motion.shape[1])) + ], axis=0) + + return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens), name + + + + +def DATALoader(dataset_name, is_test, + batch_size, w_vectorizer, + num_workers = 8, unit_length = 4) : + + val_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, is_test, w_vectorizer, unit_length=unit_length), + batch_size, + shuffle = True, + num_workers=num_workers, + collate_fn=collate_fn, + drop_last = True) + return val_loader + + +def cycle(iterable): + while True: + for x in iterable: + yield x diff --git a/VQ-Trans/dataset/dataset_TM_train.py b/VQ-Trans/dataset/dataset_TM_train.py new file mode 100644 index 0000000000000000000000000000000000000000..0b0223effb01c1cf57fa6b2b6fb8d9d01b83f84a --- /dev/null +++ b/VQ-Trans/dataset/dataset_TM_train.py @@ -0,0 +1,161 @@ +import torch +from torch.utils import data +import numpy as np +from os.path import join as pjoin +import random +import codecs as cs +from tqdm import tqdm +import utils.paramUtil as paramUtil +from torch.utils.data._utils.collate import default_collate + + +def collate_fn(batch): + batch.sort(key=lambda x: x[3], reverse=True) + return default_collate(batch) + + +'''For use of training text-2-motion generative model''' +class Text2MotionDataset(data.Dataset): + def __init__(self, dataset_name, feat_bias = 5, unit_length = 4, codebook_size = 1024, tokenizer_name=None): + + self.max_length = 64 + self.pointer = 0 + self.dataset_name = dataset_name + + self.unit_length = unit_length + # self.mot_start_idx = codebook_size + self.mot_end_idx = codebook_size + self.mot_pad_idx = codebook_size + 1 + if dataset_name == 't2m': + self.data_root = './dataset/HumanML3D' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 22 + radius = 4 + fps = 20 + self.max_motion_length = 26 if unit_length == 8 else 51 + dim_pose = 263 + kinematic_chain = paramUtil.t2m_kinematic_chain + elif dataset_name == 'kit': + self.data_root = './dataset/KIT-ML' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 21 + radius = 240 * 8 + fps = 12.5 + dim_pose = 251 + self.max_motion_length = 26 if unit_length == 8 else 51 + kinematic_chain = paramUtil.kit_kinematic_chain + + split_file = pjoin(self.data_root, 'train.txt') + + + id_list = [] + with cs.open(split_file, 'r') as f: + for line in f.readlines(): + id_list.append(line.strip()) + + new_name_list = [] + data_dict = {} + for name in tqdm(id_list): + try: + m_token_list = np.load(pjoin(self.data_root, tokenizer_name, '%s.npy'%name)) + + # Read text + with cs.open(pjoin(self.text_dir, name + '.txt')) as f: + text_data = [] + flag = False + lines = f.readlines() + + for line in lines: + try: + text_dict = {} + line_split = line.strip().split('#') + caption = line_split[0] + t_tokens = line_split[1].split(' ') + f_tag = float(line_split[2]) + to_tag = float(line_split[3]) + f_tag = 0.0 if np.isnan(f_tag) else f_tag + to_tag = 0.0 if np.isnan(to_tag) else to_tag + + text_dict['caption'] = caption + text_dict['tokens'] = t_tokens + if f_tag == 0.0 and to_tag == 0.0: + flag = True + text_data.append(text_dict) + else: + m_token_list_new = [tokens[int(f_tag*fps/unit_length) : int(to_tag*fps/unit_length)] for tokens in m_token_list if int(f_tag*fps/unit_length) < int(to_tag*fps/unit_length)] + + if len(m_token_list_new) == 0: + continue + new_name = '%s_%f_%f'%(name, f_tag, to_tag) + + data_dict[new_name] = {'m_token_list': m_token_list_new, + 'text':[text_dict]} + new_name_list.append(new_name) + except: + pass + + if flag: + data_dict[name] = {'m_token_list': m_token_list, + 'text':text_data} + new_name_list.append(name) + except: + pass + self.data_dict = data_dict + self.name_list = new_name_list + + def __len__(self): + return len(self.data_dict) + + def __getitem__(self, item): + data = self.data_dict[self.name_list[item]] + m_token_list, text_list = data['m_token_list'], data['text'] + m_tokens = random.choice(m_token_list) + + text_data = random.choice(text_list) + caption= text_data['caption'] + + + coin = np.random.choice([False, False, True]) + # print(len(m_tokens)) + if coin: + # drop one token at the head or tail + coin2 = np.random.choice([True, False]) + if coin2: + m_tokens = m_tokens[:-1] + else: + m_tokens = m_tokens[1:] + m_tokens_len = m_tokens.shape[0] + + if m_tokens_len+1 < self.max_motion_length: + m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx, np.ones((self.max_motion_length-1-m_tokens_len), dtype=int) * self.mot_pad_idx], axis=0) + else: + m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx], axis=0) + + return caption, m_tokens.reshape(-1), m_tokens_len + + + + +def DATALoader(dataset_name, + batch_size, codebook_size, tokenizer_name, unit_length=4, + num_workers = 8) : + + train_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, codebook_size = codebook_size, tokenizer_name = tokenizer_name, unit_length=unit_length), + batch_size, + shuffle=True, + num_workers=num_workers, + #collate_fn=collate_fn, + drop_last = True) + + + return train_loader + + +def cycle(iterable): + while True: + for x in iterable: + yield x + + diff --git a/VQ-Trans/dataset/dataset_VQ.py b/VQ-Trans/dataset/dataset_VQ.py new file mode 100644 index 0000000000000000000000000000000000000000..2342de946f2cbdf64729a5145168df1bdda54fa0 --- /dev/null +++ b/VQ-Trans/dataset/dataset_VQ.py @@ -0,0 +1,109 @@ +import torch +from torch.utils import data +import numpy as np +from os.path import join as pjoin +import random +import codecs as cs +from tqdm import tqdm + + + +class VQMotionDataset(data.Dataset): + def __init__(self, dataset_name, window_size = 64, unit_length = 4): + self.window_size = window_size + self.unit_length = unit_length + self.dataset_name = dataset_name + + if dataset_name == 't2m': + self.data_root = './dataset/HumanML3D' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 22 + self.max_motion_length = 196 + self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + + elif dataset_name == 'kit': + self.data_root = './dataset/KIT-ML' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 21 + + self.max_motion_length = 196 + self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + + joints_num = self.joints_num + + mean = np.load(pjoin(self.meta_dir, 'mean.npy')) + std = np.load(pjoin(self.meta_dir, 'std.npy')) + + split_file = pjoin(self.data_root, 'train.txt') + + self.data = [] + self.lengths = [] + id_list = [] + with cs.open(split_file, 'r') as f: + for line in f.readlines(): + id_list.append(line.strip()) + + for name in tqdm(id_list): + try: + motion = np.load(pjoin(self.motion_dir, name + '.npy')) + if motion.shape[0] < self.window_size: + continue + self.lengths.append(motion.shape[0] - self.window_size) + self.data.append(motion) + except: + # Some motion may not exist in KIT dataset + pass + + + self.mean = mean + self.std = std + print("Total number of motions {}".format(len(self.data))) + + def inv_transform(self, data): + return data * self.std + self.mean + + def compute_sampling_prob(self) : + + prob = np.array(self.lengths, dtype=np.float32) + prob /= np.sum(prob) + return prob + + def __len__(self): + return len(self.data) + + def __getitem__(self, item): + motion = self.data[item] + + idx = random.randint(0, len(motion) - self.window_size) + + motion = motion[idx:idx+self.window_size] + "Z Normalization" + motion = (motion - self.mean) / self.std + + return motion + +def DATALoader(dataset_name, + batch_size, + num_workers = 8, + window_size = 64, + unit_length = 4): + + trainSet = VQMotionDataset(dataset_name, window_size=window_size, unit_length=unit_length) + prob = trainSet.compute_sampling_prob() + sampler = torch.utils.data.WeightedRandomSampler(prob, num_samples = len(trainSet) * 1000, replacement=True) + train_loader = torch.utils.data.DataLoader(trainSet, + batch_size, + shuffle=True, + #sampler=sampler, + num_workers=num_workers, + #collate_fn=collate_fn, + drop_last = True) + + return train_loader + +def cycle(iterable): + while True: + for x in iterable: + yield x diff --git a/VQ-Trans/dataset/dataset_tokenize.py b/VQ-Trans/dataset/dataset_tokenize.py new file mode 100644 index 0000000000000000000000000000000000000000..641a02a75f2cfaadea45851cad2a95b39bfa1eae --- /dev/null +++ b/VQ-Trans/dataset/dataset_tokenize.py @@ -0,0 +1,117 @@ +import torch +from torch.utils import data +import numpy as np +from os.path import join as pjoin +import random +import codecs as cs +from tqdm import tqdm + + + +class VQMotionDataset(data.Dataset): + def __init__(self, dataset_name, feat_bias = 5, window_size = 64, unit_length = 8): + self.window_size = window_size + self.unit_length = unit_length + self.feat_bias = feat_bias + + self.dataset_name = dataset_name + min_motion_len = 40 if dataset_name =='t2m' else 24 + + if dataset_name == 't2m': + self.data_root = './dataset/HumanML3D' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 22 + radius = 4 + fps = 20 + self.max_motion_length = 196 + dim_pose = 263 + self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + #kinematic_chain = paramUtil.t2m_kinematic_chain + elif dataset_name == 'kit': + self.data_root = './dataset/KIT-ML' + self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') + self.text_dir = pjoin(self.data_root, 'texts') + self.joints_num = 21 + radius = 240 * 8 + fps = 12.5 + dim_pose = 251 + self.max_motion_length = 196 + self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' + #kinematic_chain = paramUtil.kit_kinematic_chain + + joints_num = self.joints_num + + mean = np.load(pjoin(self.meta_dir, 'mean.npy')) + std = np.load(pjoin(self.meta_dir, 'std.npy')) + + split_file = pjoin(self.data_root, 'train.txt') + + data_dict = {} + id_list = [] + with cs.open(split_file, 'r') as f: + for line in f.readlines(): + id_list.append(line.strip()) + + new_name_list = [] + length_list = [] + for name in tqdm(id_list): + try: + motion = np.load(pjoin(self.motion_dir, name + '.npy')) + if (len(motion)) < min_motion_len or (len(motion) >= 200): + continue + + data_dict[name] = {'motion': motion, + 'length': len(motion), + 'name': name} + new_name_list.append(name) + length_list.append(len(motion)) + except: + # Some motion may not exist in KIT dataset + pass + + + self.mean = mean + self.std = std + self.length_arr = np.array(length_list) + self.data_dict = data_dict + self.name_list = new_name_list + + def inv_transform(self, data): + return data * self.std + self.mean + + def __len__(self): + return len(self.data_dict) + + def __getitem__(self, item): + name = self.name_list[item] + data = self.data_dict[name] + motion, m_length = data['motion'], data['length'] + + m_length = (m_length // self.unit_length) * self.unit_length + + idx = random.randint(0, len(motion) - m_length) + motion = motion[idx:idx+m_length] + + "Z Normalization" + motion = (motion - self.mean) / self.std + + return motion, name + +def DATALoader(dataset_name, + batch_size = 1, + num_workers = 8, unit_length = 4) : + + train_loader = torch.utils.data.DataLoader(VQMotionDataset(dataset_name, unit_length=unit_length), + batch_size, + shuffle=True, + num_workers=num_workers, + #collate_fn=collate_fn, + drop_last = True) + + return train_loader + +def cycle(iterable): + while True: + for x in iterable: + yield x diff --git a/VQ-Trans/dataset/prepare/download_extractor.sh b/VQ-Trans/dataset/prepare/download_extractor.sh new file mode 100644 index 0000000000000000000000000000000000000000..b1c456e8311a59a1c8d86e85da5ddd3aa7e1f9a4 --- /dev/null +++ b/VQ-Trans/dataset/prepare/download_extractor.sh @@ -0,0 +1,15 @@ +rm -rf checkpoints +mkdir checkpoints +cd checkpoints +echo -e "Downloading extractors" +gdown --fuzzy https://drive.google.com/file/d/1o7RTDQcToJjTm9_mNWTyzvZvjTWpZfug/view +gdown --fuzzy https://drive.google.com/file/d/1tX79xk0fflp07EZ660Xz1RAFE33iEyJR/view + + +unzip t2m.zip +unzip kit.zip + +echo -e "Cleaning\n" +rm t2m.zip +rm kit.zip +echo -e "Downloading done!" \ No newline at end of file diff --git a/VQ-Trans/dataset/prepare/download_glove.sh b/VQ-Trans/dataset/prepare/download_glove.sh new file mode 100644 index 0000000000000000000000000000000000000000..058599aa32c9c97e0e3fc0a9658822e9c904955a --- /dev/null +++ b/VQ-Trans/dataset/prepare/download_glove.sh @@ -0,0 +1,9 @@ +echo -e "Downloading glove (in use by the evaluators)" +gdown --fuzzy https://drive.google.com/file/d/1bCeS6Sh_mLVTebxIgiUHgdPrroW06mb6/view?usp=sharing +rm -rf glove + +unzip glove.zip +echo -e "Cleaning\n" +rm glove.zip + +echo -e "Downloading done!" \ No newline at end of file diff --git a/VQ-Trans/dataset/prepare/download_model.sh b/VQ-Trans/dataset/prepare/download_model.sh new file mode 100644 index 0000000000000000000000000000000000000000..da32436f6efa93e0c14e1dd52f97068bd75956ab --- /dev/null +++ b/VQ-Trans/dataset/prepare/download_model.sh @@ -0,0 +1,12 @@ + +mkdir -p pretrained +cd pretrained/ + +echo -e "The pretrained model files will be stored in the 'pretrained' folder\n" +gdown 1LaOvwypF-jM2Axnq5dc-Iuvv3w_G-WDE + +unzip VQTrans_pretrained.zip +echo -e "Cleaning\n" +rm VQTrans_pretrained.zip + +echo -e "Downloading done!" \ No newline at end of file diff --git a/VQ-Trans/dataset/prepare/download_smpl.sh b/VQ-Trans/dataset/prepare/download_smpl.sh new file mode 100644 index 0000000000000000000000000000000000000000..411325b509e891d96b859bf28f7b983005ca360a --- /dev/null +++ b/VQ-Trans/dataset/prepare/download_smpl.sh @@ -0,0 +1,13 @@ + +mkdir -p body_models +cd body_models/ + +echo -e "The smpl files will be stored in the 'body_models/smpl/' folder\n" +gdown 1INYlGA76ak_cKGzvpOV2Pe6RkYTlXTW2 +rm -rf smpl + +unzip smpl.zip +echo -e "Cleaning\n" +rm smpl.zip + +echo -e "Downloading done!" \ No newline at end of file diff --git a/VQ-Trans/environment.yml b/VQ-Trans/environment.yml new file mode 100644 index 0000000000000000000000000000000000000000..cb0abb7f5c278d1eaee782c02abb3a46da736f90 --- /dev/null +++ b/VQ-Trans/environment.yml @@ -0,0 +1,121 @@ +name: VQTrans +channels: + - pytorch + - defaults +dependencies: + - _libgcc_mutex=0.1=main + - _openmp_mutex=4.5=1_gnu + - blas=1.0=mkl + - bzip2=1.0.8=h7b6447c_0 + - ca-certificates=2021.7.5=h06a4308_1 + - certifi=2021.5.30=py38h06a4308_0 + - cudatoolkit=10.1.243=h6bb024c_0 + - ffmpeg=4.3=hf484d3e_0 + - freetype=2.10.4=h5ab3b9f_0 + - gmp=6.2.1=h2531618_2 + - gnutls=3.6.15=he1e5248_0 + - intel-openmp=2021.3.0=h06a4308_3350 + - jpeg=9b=h024ee3a_2 + - lame=3.100=h7b6447c_0 + - lcms2=2.12=h3be6417_0 + - ld_impl_linux-64=2.35.1=h7274673_9 + - libffi=3.3=he6710b0_2 + - libgcc-ng=9.3.0=h5101ec6_17 + - libgomp=9.3.0=h5101ec6_17 + - libiconv=1.15=h63c8f33_5 + - libidn2=2.3.2=h7f8727e_0 + - libpng=1.6.37=hbc83047_0 + - libstdcxx-ng=9.3.0=hd4cf53a_17 + - libtasn1=4.16.0=h27cfd23_0 + - libtiff=4.2.0=h85742a9_0 + - libunistring=0.9.10=h27cfd23_0 + - libuv=1.40.0=h7b6447c_0 + - libwebp-base=1.2.0=h27cfd23_0 + - lz4-c=1.9.3=h295c915_1 + - mkl=2021.3.0=h06a4308_520 + - mkl-service=2.4.0=py38h7f8727e_0 + - mkl_fft=1.3.0=py38h42c9631_2 + - mkl_random=1.2.2=py38h51133e4_0 + - ncurses=6.2=he6710b0_1 + - nettle=3.7.3=hbbd107a_1 + - ninja=1.10.2=hff7bd54_1 + - numpy=1.20.3=py38hf144106_0 + - numpy-base=1.20.3=py38h74d4b33_0 + - olefile=0.46=py_0 + - openh264=2.1.0=hd408876_0 + - openjpeg=2.3.0=h05c96fa_1 + - openssl=1.1.1k=h27cfd23_0 + - pillow=8.3.1=py38h2c7a002_0 + - pip=21.0.1=py38h06a4308_0 + - python=3.8.11=h12debd9_0_cpython + - pytorch=1.8.1=py3.8_cuda10.1_cudnn7.6.3_0 + - readline=8.1=h27cfd23_0 + - setuptools=52.0.0=py38h06a4308_0 + - six=1.16.0=pyhd3eb1b0_0 + - sqlite=3.36.0=hc218d9a_0 + - tk=8.6.10=hbc83047_0 + - torchaudio=0.8.1=py38 + - torchvision=0.9.1=py38_cu101 + - typing_extensions=3.10.0.0=pyh06a4308_0 + - wheel=0.37.0=pyhd3eb1b0_0 + - xz=5.2.5=h7b6447c_0 + - zlib=1.2.11=h7b6447c_3 + - zstd=1.4.9=haebb681_0 + - pip: + - absl-py==0.13.0 + - backcall==0.2.0 + - cachetools==4.2.2 + - charset-normalizer==2.0.4 + - chumpy==0.70 + - cycler==0.10.0 + - decorator==5.0.9 + - google-auth==1.35.0 + - google-auth-oauthlib==0.4.5 + - grpcio==1.39.0 + - idna==3.2 + - imageio==2.9.0 + - ipdb==0.13.9 + - ipython==7.26.0 + - ipython-genutils==0.2.0 + - jedi==0.18.0 + - joblib==1.0.1 + - kiwisolver==1.3.1 + - markdown==3.3.4 + - matplotlib==3.4.3 + - matplotlib-inline==0.1.2 + - oauthlib==3.1.1 + - pandas==1.3.2 + - parso==0.8.2 + - pexpect==4.8.0 + - pickleshare==0.7.5 + - prompt-toolkit==3.0.20 + - protobuf==3.17.3 + - ptyprocess==0.7.0 + - pyasn1==0.4.8 + - pyasn1-modules==0.2.8 + - pygments==2.10.0 + - pyparsing==2.4.7 + - python-dateutil==2.8.2 + - pytz==2021.1 + - pyyaml==5.4.1 + - requests==2.26.0 + - requests-oauthlib==1.3.0 + - rsa==4.7.2 + - scikit-learn==0.24.2 + - scipy==1.7.1 + - sklearn==0.0 + - smplx==0.1.28 + - tensorboard==2.6.0 + - tensorboard-data-server==0.6.1 + - tensorboard-plugin-wit==1.8.0 + - threadpoolctl==2.2.0 + - toml==0.10.2 + - tqdm==4.62.2 + - traitlets==5.0.5 + - urllib3==1.26.6 + - wcwidth==0.2.5 + - werkzeug==2.0.1 + - git+https://github.com/openai/CLIP.git + - git+https://github.com/nghorbani/human_body_prior + - gdown + - moviepy \ No newline at end of file diff --git a/VQ-Trans/models/encdec.py b/VQ-Trans/models/encdec.py new file mode 100644 index 0000000000000000000000000000000000000000..ae72afaa5aa59ad67cadb38e0d83e420fc6edb09 --- /dev/null +++ b/VQ-Trans/models/encdec.py @@ -0,0 +1,67 @@ +import torch.nn as nn +from models.resnet import Resnet1D + +class Encoder(nn.Module): + def __init__(self, + input_emb_width = 3, + output_emb_width = 512, + down_t = 3, + stride_t = 2, + width = 512, + depth = 3, + dilation_growth_rate = 3, + activation='relu', + norm=None): + super().__init__() + + blocks = [] + filter_t, pad_t = stride_t * 2, stride_t // 2 + blocks.append(nn.Conv1d(input_emb_width, width, 3, 1, 1)) + blocks.append(nn.ReLU()) + + for i in range(down_t): + input_dim = width + block = nn.Sequential( + nn.Conv1d(input_dim, width, filter_t, stride_t, pad_t), + Resnet1D(width, depth, dilation_growth_rate, activation=activation, norm=norm), + ) + blocks.append(block) + blocks.append(nn.Conv1d(width, output_emb_width, 3, 1, 1)) + self.model = nn.Sequential(*blocks) + + def forward(self, x): + return self.model(x) + +class Decoder(nn.Module): + def __init__(self, + input_emb_width = 3, + output_emb_width = 512, + down_t = 3, + stride_t = 2, + width = 512, + depth = 3, + dilation_growth_rate = 3, + activation='relu', + norm=None): + super().__init__() + blocks = [] + + filter_t, pad_t = stride_t * 2, stride_t // 2 + blocks.append(nn.Conv1d(output_emb_width, width, 3, 1, 1)) + blocks.append(nn.ReLU()) + for i in range(down_t): + out_dim = width + block = nn.Sequential( + Resnet1D(width, depth, dilation_growth_rate, reverse_dilation=True, activation=activation, norm=norm), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv1d(width, out_dim, 3, 1, 1) + ) + blocks.append(block) + blocks.append(nn.Conv1d(width, width, 3, 1, 1)) + blocks.append(nn.ReLU()) + blocks.append(nn.Conv1d(width, input_emb_width, 3, 1, 1)) + self.model = nn.Sequential(*blocks) + + def forward(self, x): + return self.model(x) + diff --git a/VQ-Trans/models/evaluator_wrapper.py b/VQ-Trans/models/evaluator_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..fe4558a13ccc2ce0579540b8b77f958096e9984c --- /dev/null +++ b/VQ-Trans/models/evaluator_wrapper.py @@ -0,0 +1,92 @@ + +import torch +from os.path import join as pjoin +import numpy as np +from models.modules import MovementConvEncoder, TextEncoderBiGRUCo, MotionEncoderBiGRUCo +from utils.word_vectorizer import POS_enumerator + +def build_models(opt): + movement_enc = MovementConvEncoder(opt.dim_pose-4, opt.dim_movement_enc_hidden, opt.dim_movement_latent) + text_enc = TextEncoderBiGRUCo(word_size=opt.dim_word, + pos_size=opt.dim_pos_ohot, + hidden_size=opt.dim_text_hidden, + output_size=opt.dim_coemb_hidden, + device=opt.device) + + motion_enc = MotionEncoderBiGRUCo(input_size=opt.dim_movement_latent, + hidden_size=opt.dim_motion_hidden, + output_size=opt.dim_coemb_hidden, + device=opt.device) + + checkpoint = torch.load(pjoin(opt.checkpoints_dir, opt.dataset_name, 'text_mot_match', 'model', 'finest.tar'), + map_location=opt.device) + movement_enc.load_state_dict(checkpoint['movement_encoder']) + text_enc.load_state_dict(checkpoint['text_encoder']) + motion_enc.load_state_dict(checkpoint['motion_encoder']) + print('Loading Evaluation Model Wrapper (Epoch %d) Completed!!' % (checkpoint['epoch'])) + return text_enc, motion_enc, movement_enc + + +class EvaluatorModelWrapper(object): + + def __init__(self, opt): + + if opt.dataset_name == 't2m': + opt.dim_pose = 263 + elif opt.dataset_name == 'kit': + opt.dim_pose = 251 + else: + raise KeyError('Dataset not Recognized!!!') + + opt.dim_word = 300 + opt.max_motion_length = 196 + opt.dim_pos_ohot = len(POS_enumerator) + opt.dim_motion_hidden = 1024 + opt.max_text_len = 20 + opt.dim_text_hidden = 512 + opt.dim_coemb_hidden = 512 + + # print(opt) + + self.text_encoder, self.motion_encoder, self.movement_encoder = build_models(opt) + self.opt = opt + self.device = opt.device + + self.text_encoder.to(opt.device) + self.motion_encoder.to(opt.device) + self.movement_encoder.to(opt.device) + + self.text_encoder.eval() + self.motion_encoder.eval() + self.movement_encoder.eval() + + # Please note that the results does not following the order of inputs + def get_co_embeddings(self, word_embs, pos_ohot, cap_lens, motions, m_lens): + with torch.no_grad(): + word_embs = word_embs.detach().to(self.device).float() + pos_ohot = pos_ohot.detach().to(self.device).float() + motions = motions.detach().to(self.device).float() + + '''Movement Encoding''' + movements = self.movement_encoder(motions[..., :-4]).detach() + m_lens = m_lens // self.opt.unit_length + motion_embedding = self.motion_encoder(movements, m_lens) + + '''Text Encoding''' + text_embedding = self.text_encoder(word_embs, pos_ohot, cap_lens) + return text_embedding, motion_embedding + + # Please note that the results does not following the order of inputs + def get_motion_embeddings(self, motions, m_lens): + with torch.no_grad(): + motions = motions.detach().to(self.device).float() + + align_idx = np.argsort(m_lens.data.tolist())[::-1].copy() + motions = motions[align_idx] + m_lens = m_lens[align_idx] + + '''Movement Encoding''' + movements = self.movement_encoder(motions[..., :-4]).detach() + m_lens = m_lens // self.opt.unit_length + motion_embedding = self.motion_encoder(movements, m_lens) + return motion_embedding diff --git a/VQ-Trans/models/modules.py b/VQ-Trans/models/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..4f06cd98d4f6029bd3df073095cf50498483d54a --- /dev/null +++ b/VQ-Trans/models/modules.py @@ -0,0 +1,109 @@ +import torch +import torch.nn as nn +from torch.nn.utils.rnn import pack_padded_sequence + +def init_weight(m): + if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear) or isinstance(m, nn.ConvTranspose1d): + nn.init.xavier_normal_(m.weight) + # m.bias.data.fill_(0.01) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + +class MovementConvEncoder(nn.Module): + def __init__(self, input_size, hidden_size, output_size): + super(MovementConvEncoder, self).__init__() + self.main = nn.Sequential( + nn.Conv1d(input_size, hidden_size, 4, 2, 1), + nn.Dropout(0.2, inplace=True), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv1d(hidden_size, output_size, 4, 2, 1), + nn.Dropout(0.2, inplace=True), + nn.LeakyReLU(0.2, inplace=True), + ) + self.out_net = nn.Linear(output_size, output_size) + self.main.apply(init_weight) + self.out_net.apply(init_weight) + + def forward(self, inputs): + inputs = inputs.permute(0, 2, 1) + outputs = self.main(inputs).permute(0, 2, 1) + # print(outputs.shape) + return self.out_net(outputs) + + + +class TextEncoderBiGRUCo(nn.Module): + def __init__(self, word_size, pos_size, hidden_size, output_size, device): + super(TextEncoderBiGRUCo, self).__init__() + self.device = device + + self.pos_emb = nn.Linear(pos_size, word_size) + self.input_emb = nn.Linear(word_size, hidden_size) + self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True) + self.output_net = nn.Sequential( + nn.Linear(hidden_size * 2, hidden_size), + nn.LayerNorm(hidden_size), + nn.LeakyReLU(0.2, inplace=True), + nn.Linear(hidden_size, output_size) + ) + + self.input_emb.apply(init_weight) + self.pos_emb.apply(init_weight) + self.output_net.apply(init_weight) + self.hidden_size = hidden_size + self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True)) + + # input(batch_size, seq_len, dim) + def forward(self, word_embs, pos_onehot, cap_lens): + num_samples = word_embs.shape[0] + + pos_embs = self.pos_emb(pos_onehot) + inputs = word_embs + pos_embs + input_embs = self.input_emb(inputs) + hidden = self.hidden.repeat(1, num_samples, 1) + + cap_lens = cap_lens.data.tolist() + emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True) + + gru_seq, gru_last = self.gru(emb, hidden) + + gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1) + + return self.output_net(gru_last) + + +class MotionEncoderBiGRUCo(nn.Module): + def __init__(self, input_size, hidden_size, output_size, device): + super(MotionEncoderBiGRUCo, self).__init__() + self.device = device + + self.input_emb = nn.Linear(input_size, hidden_size) + self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True) + self.output_net = nn.Sequential( + nn.Linear(hidden_size*2, hidden_size), + nn.LayerNorm(hidden_size), + nn.LeakyReLU(0.2, inplace=True), + nn.Linear(hidden_size, output_size) + ) + + self.input_emb.apply(init_weight) + self.output_net.apply(init_weight) + self.hidden_size = hidden_size + self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True)) + + # input(batch_size, seq_len, dim) + def forward(self, inputs, m_lens): + num_samples = inputs.shape[0] + + input_embs = self.input_emb(inputs) + hidden = self.hidden.repeat(1, num_samples, 1) + + cap_lens = m_lens.data.tolist() + emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True, enforce_sorted=False) + + gru_seq, gru_last = self.gru(emb, hidden) + + gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1) + + return self.output_net(gru_last) diff --git a/VQ-Trans/models/pos_encoding.py b/VQ-Trans/models/pos_encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..066be3e1f8a1636f7eaabd1c534b9c618ee3e9f8 --- /dev/null +++ b/VQ-Trans/models/pos_encoding.py @@ -0,0 +1,43 @@ +""" +Various positional encodings for the transformer. +""" +import math +import torch +from torch import nn + +def PE1d_sincos(seq_length, dim): + """ + :param d_model: dimension of the model + :param length: length of positions + :return: length*d_model position matrix + """ + if dim % 2 != 0: + raise ValueError("Cannot use sin/cos positional encoding with " + "odd dim (got dim={:d})".format(dim)) + pe = torch.zeros(seq_length, dim) + position = torch.arange(0, seq_length).unsqueeze(1) + div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) * + -(math.log(10000.0) / dim))) + pe[:, 0::2] = torch.sin(position.float() * div_term) + pe[:, 1::2] = torch.cos(position.float() * div_term) + + return pe.unsqueeze(1) + + +class PositionEmbedding(nn.Module): + """ + Absolute pos embedding (standard), learned. + """ + def __init__(self, seq_length, dim, dropout, grad=False): + super().__init__() + self.embed = nn.Parameter(data=PE1d_sincos(seq_length, dim), requires_grad=grad) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x): + # x.shape: bs, seq_len, feat_dim + l = x.shape[1] + x = x.permute(1, 0, 2) + self.embed[:l].expand(x.permute(1, 0, 2).shape) + x = self.dropout(x.permute(1, 0, 2)) + return x + + \ No newline at end of file diff --git a/VQ-Trans/models/quantize_cnn.py b/VQ-Trans/models/quantize_cnn.py new file mode 100644 index 0000000000000000000000000000000000000000..b796772749efda9a225bdcb0e7262791a972a710 --- /dev/null +++ b/VQ-Trans/models/quantize_cnn.py @@ -0,0 +1,415 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +class QuantizeEMAReset(nn.Module): + def __init__(self, nb_code, code_dim, args): + super().__init__() + self.nb_code = nb_code + self.code_dim = code_dim + self.mu = args.mu + self.reset_codebook() + + def reset_codebook(self): + self.init = False + self.code_sum = None + self.code_count = None + if torch.cuda.is_available(): + self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim).cuda()) + else: + self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim)) + + def _tile(self, x): + nb_code_x, code_dim = x.shape + if nb_code_x < self.nb_code: + n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x + std = 0.01 / np.sqrt(code_dim) + out = x.repeat(n_repeats, 1) + out = out + torch.randn_like(out) * std + else : + out = x + return out + + def init_codebook(self, x): + out = self._tile(x) + self.codebook = out[:self.nb_code] + self.code_sum = self.codebook.clone() + self.code_count = torch.ones(self.nb_code, device=self.codebook.device) + self.init = True + + @torch.no_grad() + def compute_perplexity(self, code_idx) : + # Calculate new centres + code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1) + + code_count = code_onehot.sum(dim=-1) # nb_code + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + return perplexity + + @torch.no_grad() + def update_codebook(self, x, code_idx): + + code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1) + + code_sum = torch.matmul(code_onehot, x) # nb_code, w + code_count = code_onehot.sum(dim=-1) # nb_code + + out = self._tile(x) + code_rand = out[:self.nb_code] + + # Update centres + self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum # w, nb_code + self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count # nb_code + + usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float() + code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1) + + self.codebook = usage * code_update + (1 - usage) * code_rand + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + + + return perplexity + + def preprocess(self, x): + # NCT -> NTC -> [NT, C] + x = x.permute(0, 2, 1).contiguous() + x = x.view(-1, x.shape[-1]) + return x + + def quantize(self, x): + # Calculate latent code x_l + k_w = self.codebook.t() + distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0, + keepdim=True) # (N * L, b) + _, code_idx = torch.min(distance, dim=-1) + return code_idx + + def dequantize(self, code_idx): + x = F.embedding(code_idx, self.codebook) + return x + + + def forward(self, x): + N, width, T = x.shape + + # Preprocess + x = self.preprocess(x) + + # Init codebook if not inited + if self.training and not self.init: + self.init_codebook(x) + + # quantize and dequantize through bottleneck + code_idx = self.quantize(x) + x_d = self.dequantize(code_idx) + + # Update embeddings + if self.training: + perplexity = self.update_codebook(x, code_idx) + else : + perplexity = self.compute_perplexity(code_idx) + + # Loss + commit_loss = F.mse_loss(x, x_d.detach()) + + # Passthrough + x_d = x + (x_d - x).detach() + + # Postprocess + x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous() #(N, DIM, T) + + return x_d, commit_loss, perplexity + + + +class Quantizer(nn.Module): + def __init__(self, n_e, e_dim, beta): + super(Quantizer, self).__init__() + + self.e_dim = e_dim + self.n_e = n_e + self.beta = beta + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + def forward(self, z): + + N, width, T = z.shape + z = self.preprocess(z) + assert z.shape[-1] == self.e_dim + z_flattened = z.contiguous().view(-1, self.e_dim) + + # B x V + d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight**2, dim=1) - 2 * \ + torch.matmul(z_flattened, self.embedding.weight.t()) + # B x 1 + min_encoding_indices = torch.argmin(d, dim=1) + z_q = self.embedding(min_encoding_indices).view(z.shape) + + # compute loss for embedding + loss = torch.mean((z_q - z.detach())**2) + self.beta * \ + torch.mean((z_q.detach() - z)**2) + + # preserve gradients + z_q = z + (z_q - z).detach() + z_q = z_q.view(N, T, -1).permute(0, 2, 1).contiguous() #(N, DIM, T) + + min_encodings = F.one_hot(min_encoding_indices, self.n_e).type(z.dtype) + e_mean = torch.mean(min_encodings, dim=0) + perplexity = torch.exp(-torch.sum(e_mean*torch.log(e_mean + 1e-10))) + return z_q, loss, perplexity + + def quantize(self, z): + + assert z.shape[-1] == self.e_dim + + # B x V + d = torch.sum(z ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \ + torch.matmul(z, self.embedding.weight.t()) + # B x 1 + min_encoding_indices = torch.argmin(d, dim=1) + return min_encoding_indices + + def dequantize(self, indices): + + index_flattened = indices.view(-1) + z_q = self.embedding(index_flattened) + z_q = z_q.view(indices.shape + (self.e_dim, )).contiguous() + return z_q + + def preprocess(self, x): + # NCT -> NTC -> [NT, C] + x = x.permute(0, 2, 1).contiguous() + x = x.view(-1, x.shape[-1]) + return x + + + +class QuantizeReset(nn.Module): + def __init__(self, nb_code, code_dim, args): + super().__init__() + self.nb_code = nb_code + self.code_dim = code_dim + self.reset_codebook() + self.codebook = nn.Parameter(torch.randn(nb_code, code_dim)) + + def reset_codebook(self): + self.init = False + self.code_count = None + + def _tile(self, x): + nb_code_x, code_dim = x.shape + if nb_code_x < self.nb_code: + n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x + std = 0.01 / np.sqrt(code_dim) + out = x.repeat(n_repeats, 1) + out = out + torch.randn_like(out) * std + else : + out = x + return out + + def init_codebook(self, x): + out = self._tile(x) + self.codebook = nn.Parameter(out[:self.nb_code]) + self.code_count = torch.ones(self.nb_code, device=self.codebook.device) + self.init = True + + @torch.no_grad() + def compute_perplexity(self, code_idx) : + # Calculate new centres + code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1) + + code_count = code_onehot.sum(dim=-1) # nb_code + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + return perplexity + + def update_codebook(self, x, code_idx): + + code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1) + + code_count = code_onehot.sum(dim=-1) # nb_code + + out = self._tile(x) + code_rand = out[:self.nb_code] + + # Update centres + self.code_count = code_count # nb_code + usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float() + + self.codebook.data = usage * self.codebook.data + (1 - usage) * code_rand + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + + + return perplexity + + def preprocess(self, x): + # NCT -> NTC -> [NT, C] + x = x.permute(0, 2, 1).contiguous() + x = x.view(-1, x.shape[-1]) + return x + + def quantize(self, x): + # Calculate latent code x_l + k_w = self.codebook.t() + distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0, + keepdim=True) # (N * L, b) + _, code_idx = torch.min(distance, dim=-1) + return code_idx + + def dequantize(self, code_idx): + x = F.embedding(code_idx, self.codebook) + return x + + + def forward(self, x): + N, width, T = x.shape + # Preprocess + x = self.preprocess(x) + # Init codebook if not inited + if self.training and not self.init: + self.init_codebook(x) + # quantize and dequantize through bottleneck + code_idx = self.quantize(x) + x_d = self.dequantize(code_idx) + # Update embeddings + if self.training: + perplexity = self.update_codebook(x, code_idx) + else : + perplexity = self.compute_perplexity(code_idx) + + # Loss + commit_loss = F.mse_loss(x, x_d.detach()) + + # Passthrough + x_d = x + (x_d - x).detach() + + # Postprocess + x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous() #(N, DIM, T) + + return x_d, commit_loss, perplexity + +class QuantizeEMA(nn.Module): + def __init__(self, nb_code, code_dim, args): + super().__init__() + self.nb_code = nb_code + self.code_dim = code_dim + self.mu = 0.99 + self.reset_codebook() + + def reset_codebook(self): + self.init = False + self.code_sum = None + self.code_count = None + self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim).cuda()) + + def _tile(self, x): + nb_code_x, code_dim = x.shape + if nb_code_x < self.nb_code: + n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x + std = 0.01 / np.sqrt(code_dim) + out = x.repeat(n_repeats, 1) + out = out + torch.randn_like(out) * std + else : + out = x + return out + + def init_codebook(self, x): + out = self._tile(x) + self.codebook = out[:self.nb_code] + self.code_sum = self.codebook.clone() + self.code_count = torch.ones(self.nb_code, device=self.codebook.device) + self.init = True + + @torch.no_grad() + def compute_perplexity(self, code_idx) : + # Calculate new centres + code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1) + + code_count = code_onehot.sum(dim=-1) # nb_code + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + return perplexity + + @torch.no_grad() + def update_codebook(self, x, code_idx): + + code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device) # nb_code, N * L + code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1) + + code_sum = torch.matmul(code_onehot, x) # nb_code, w + code_count = code_onehot.sum(dim=-1) # nb_code + + # Update centres + self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum # w, nb_code + self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count # nb_code + + code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1) + + self.codebook = code_update + prob = code_count / torch.sum(code_count) + perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7))) + + return perplexity + + def preprocess(self, x): + # NCT -> NTC -> [NT, C] + x = x.permute(0, 2, 1).contiguous() + x = x.view(-1, x.shape[-1]) + return x + + def quantize(self, x): + # Calculate latent code x_l + k_w = self.codebook.t() + distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0, + keepdim=True) # (N * L, b) + _, code_idx = torch.min(distance, dim=-1) + return code_idx + + def dequantize(self, code_idx): + x = F.embedding(code_idx, self.codebook) + return x + + + def forward(self, x): + N, width, T = x.shape + + # Preprocess + x = self.preprocess(x) + + # Init codebook if not inited + if self.training and not self.init: + self.init_codebook(x) + + # quantize and dequantize through bottleneck + code_idx = self.quantize(x) + x_d = self.dequantize(code_idx) + + # Update embeddings + if self.training: + perplexity = self.update_codebook(x, code_idx) + else : + perplexity = self.compute_perplexity(code_idx) + + # Loss + commit_loss = F.mse_loss(x, x_d.detach()) + + # Passthrough + x_d = x + (x_d - x).detach() + + # Postprocess + x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous() #(N, DIM, T) + + return x_d, commit_loss, perplexity \ No newline at end of file diff --git a/VQ-Trans/models/resnet.py b/VQ-Trans/models/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..062346e3ba2fc4d6ae5636f228c5b7565bdb62b7 --- /dev/null +++ b/VQ-Trans/models/resnet.py @@ -0,0 +1,82 @@ +import torch.nn as nn +import torch + +class nonlinearity(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + # swish + return x * torch.sigmoid(x) + +class ResConv1DBlock(nn.Module): + def __init__(self, n_in, n_state, dilation=1, activation='silu', norm=None, dropout=None): + super().__init__() + padding = dilation + self.norm = norm + if norm == "LN": + self.norm1 = nn.LayerNorm(n_in) + self.norm2 = nn.LayerNorm(n_in) + elif norm == "GN": + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True) + elif norm == "BN": + self.norm1 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True) + self.norm2 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True) + + else: + self.norm1 = nn.Identity() + self.norm2 = nn.Identity() + + if activation == "relu": + self.activation1 = nn.ReLU() + self.activation2 = nn.ReLU() + + elif activation == "silu": + self.activation1 = nonlinearity() + self.activation2 = nonlinearity() + + elif activation == "gelu": + self.activation1 = nn.GELU() + self.activation2 = nn.GELU() + + + + self.conv1 = nn.Conv1d(n_in, n_state, 3, 1, padding, dilation) + self.conv2 = nn.Conv1d(n_state, n_in, 1, 1, 0,) + + + def forward(self, x): + x_orig = x + if self.norm == "LN": + x = self.norm1(x.transpose(-2, -1)) + x = self.activation1(x.transpose(-2, -1)) + else: + x = self.norm1(x) + x = self.activation1(x) + + x = self.conv1(x) + + if self.norm == "LN": + x = self.norm2(x.transpose(-2, -1)) + x = self.activation2(x.transpose(-2, -1)) + else: + x = self.norm2(x) + x = self.activation2(x) + + x = self.conv2(x) + x = x + x_orig + return x + +class Resnet1D(nn.Module): + def __init__(self, n_in, n_depth, dilation_growth_rate=1, reverse_dilation=True, activation='relu', norm=None): + super().__init__() + + blocks = [ResConv1DBlock(n_in, n_in, dilation=dilation_growth_rate ** depth, activation=activation, norm=norm) for depth in range(n_depth)] + if reverse_dilation: + blocks = blocks[::-1] + + self.model = nn.Sequential(*blocks) + + def forward(self, x): + return self.model(x) \ No newline at end of file diff --git a/VQ-Trans/models/rotation2xyz.py b/VQ-Trans/models/rotation2xyz.py new file mode 100644 index 0000000000000000000000000000000000000000..44f6cb6c3fd0fd263bd6256803b908e9e2b4184b --- /dev/null +++ b/VQ-Trans/models/rotation2xyz.py @@ -0,0 +1,92 @@ +# This code is based on https://github.com/Mathux/ACTOR.git +import torch +import utils.rotation_conversions as geometry + + +from models.smpl import SMPL, JOINTSTYPE_ROOT +# from .get_model import JOINTSTYPES +JOINTSTYPES = ["a2m", "a2mpl", "smpl", "vibe", "vertices"] + + +class Rotation2xyz: + def __init__(self, device, dataset='amass'): + self.device = device + self.dataset = dataset + self.smpl_model = SMPL().eval().to(device) + + def __call__(self, x, mask, pose_rep, translation, glob, + jointstype, vertstrans, betas=None, beta=0, + glob_rot=None, get_rotations_back=False, **kwargs): + if pose_rep == "xyz": + return x + + if mask is None: + mask = torch.ones((x.shape[0], x.shape[-1]), dtype=bool, device=x.device) + + if not glob and glob_rot is None: + raise TypeError("You must specify global rotation if glob is False") + + if jointstype not in JOINTSTYPES: + raise NotImplementedError("This jointstype is not implemented.") + + if translation: + x_translations = x[:, -1, :3] + x_rotations = x[:, :-1] + else: + x_rotations = x + + x_rotations = x_rotations.permute(0, 3, 1, 2) + nsamples, time, njoints, feats = x_rotations.shape + + # Compute rotations (convert only masked sequences output) + if pose_rep == "rotvec": + rotations = geometry.axis_angle_to_matrix(x_rotations[mask]) + elif pose_rep == "rotmat": + rotations = x_rotations[mask].view(-1, njoints, 3, 3) + elif pose_rep == "rotquat": + rotations = geometry.quaternion_to_matrix(x_rotations[mask]) + elif pose_rep == "rot6d": + rotations = geometry.rotation_6d_to_matrix(x_rotations[mask]) + else: + raise NotImplementedError("No geometry for this one.") + + if not glob: + global_orient = torch.tensor(glob_rot, device=x.device) + global_orient = geometry.axis_angle_to_matrix(global_orient).view(1, 1, 3, 3) + global_orient = global_orient.repeat(len(rotations), 1, 1, 1) + else: + global_orient = rotations[:, 0] + rotations = rotations[:, 1:] + + if betas is None: + betas = torch.zeros([rotations.shape[0], self.smpl_model.num_betas], + dtype=rotations.dtype, device=rotations.device) + betas[:, 1] = beta + # import ipdb; ipdb.set_trace() + out = self.smpl_model(body_pose=rotations, global_orient=global_orient, betas=betas) + + # get the desirable joints + joints = out[jointstype] + + x_xyz = torch.empty(nsamples, time, joints.shape[1], 3, device=x.device, dtype=x.dtype) + x_xyz[~mask] = 0 + x_xyz[mask] = joints + + x_xyz = x_xyz.permute(0, 2, 3, 1).contiguous() + + # the first translation root at the origin on the prediction + if jointstype != "vertices": + rootindex = JOINTSTYPE_ROOT[jointstype] + x_xyz = x_xyz - x_xyz[:, [rootindex], :, :] + + if translation and vertstrans: + # the first translation root at the origin + x_translations = x_translations - x_translations[:, :, [0]] + + # add the translation to all the joints + x_xyz = x_xyz + x_translations[:, None, :, :] + + if get_rotations_back: + return x_xyz, rotations, global_orient + else: + return x_xyz diff --git a/VQ-Trans/models/smpl.py b/VQ-Trans/models/smpl.py new file mode 100644 index 0000000000000000000000000000000000000000..587f5419601a74df92c1e37263b28d4aa6a7c0a9 --- /dev/null +++ b/VQ-Trans/models/smpl.py @@ -0,0 +1,97 @@ +# This code is based on https://github.com/Mathux/ACTOR.git +import numpy as np +import torch + +import contextlib + +from smplx import SMPLLayer as _SMPLLayer +from smplx.lbs import vertices2joints + + +# action2motion_joints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 21, 24, 38] +# change 0 and 8 +action2motion_joints = [8, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 21, 24, 38] + +from utils.config import SMPL_MODEL_PATH, JOINT_REGRESSOR_TRAIN_EXTRA + +JOINTSTYPE_ROOT = {"a2m": 0, # action2motion + "smpl": 0, + "a2mpl": 0, # set(smpl, a2m) + "vibe": 8} # 0 is the 8 position: OP MidHip below + +JOINT_MAP = { + 'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17, + 'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16, + 'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0, + 'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8, + 'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7, + 'OP REye': 25, 'OP LEye': 26, 'OP REar': 27, + 'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30, + 'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34, + 'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45, + 'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7, + 'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17, + 'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20, + 'Neck (LSP)': 47, 'Top of Head (LSP)': 48, + 'Pelvis (MPII)': 49, 'Thorax (MPII)': 50, + 'Spine (H36M)': 51, 'Jaw (H36M)': 52, + 'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26, + 'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27 +} + +JOINT_NAMES = [ + 'OP Nose', 'OP Neck', 'OP RShoulder', + 'OP RElbow', 'OP RWrist', 'OP LShoulder', + 'OP LElbow', 'OP LWrist', 'OP MidHip', + 'OP RHip', 'OP RKnee', 'OP RAnkle', + 'OP LHip', 'OP LKnee', 'OP LAnkle', + 'OP REye', 'OP LEye', 'OP REar', + 'OP LEar', 'OP LBigToe', 'OP LSmallToe', + 'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel', + 'Right Ankle', 'Right Knee', 'Right Hip', + 'Left Hip', 'Left Knee', 'Left Ankle', + 'Right Wrist', 'Right Elbow', 'Right Shoulder', + 'Left Shoulder', 'Left Elbow', 'Left Wrist', + 'Neck (LSP)', 'Top of Head (LSP)', + 'Pelvis (MPII)', 'Thorax (MPII)', + 'Spine (H36M)', 'Jaw (H36M)', + 'Head (H36M)', 'Nose', 'Left Eye', + 'Right Eye', 'Left Ear', 'Right Ear' +] + + +# adapted from VIBE/SPIN to output smpl_joints, vibe joints and action2motion joints +class SMPL(_SMPLLayer): + """ Extension of the official SMPL implementation to support more joints """ + + def __init__(self, model_path=SMPL_MODEL_PATH, **kwargs): + kwargs["model_path"] = model_path + + # remove the verbosity for the 10-shapes beta parameters + with contextlib.redirect_stdout(None): + super(SMPL, self).__init__(**kwargs) + + J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA) + self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32)) + vibe_indexes = np.array([JOINT_MAP[i] for i in JOINT_NAMES]) + a2m_indexes = vibe_indexes[action2motion_joints] + smpl_indexes = np.arange(24) + a2mpl_indexes = np.unique(np.r_[smpl_indexes, a2m_indexes]) + + self.maps = {"vibe": vibe_indexes, + "a2m": a2m_indexes, + "smpl": smpl_indexes, + "a2mpl": a2mpl_indexes} + + def forward(self, *args, **kwargs): + smpl_output = super(SMPL, self).forward(*args, **kwargs) + + extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices) + all_joints = torch.cat([smpl_output.joints, extra_joints], dim=1) + + output = {"vertices": smpl_output.vertices} + + for joinstype, indexes in self.maps.items(): + output[joinstype] = all_joints[:, indexes] + + return output \ No newline at end of file diff --git a/VQ-Trans/models/t2m_trans.py b/VQ-Trans/models/t2m_trans.py new file mode 100644 index 0000000000000000000000000000000000000000..54bd0a485d7e8dbeaaac91d049f63ebd136cb074 --- /dev/null +++ b/VQ-Trans/models/t2m_trans.py @@ -0,0 +1,211 @@ +import math +import torch +import torch.nn as nn +from torch.nn import functional as F +from torch.distributions import Categorical +import models.pos_encoding as pos_encoding + +class Text2Motion_Transformer(nn.Module): + + def __init__(self, + num_vq=1024, + embed_dim=512, + clip_dim=512, + block_size=16, + num_layers=2, + n_head=8, + drop_out_rate=0.1, + fc_rate=4): + super().__init__() + self.trans_base = CrossCondTransBase(num_vq, embed_dim, clip_dim, block_size, num_layers, n_head, drop_out_rate, fc_rate) + self.trans_head = CrossCondTransHead(num_vq, embed_dim, block_size, num_layers, n_head, drop_out_rate, fc_rate) + self.block_size = block_size + self.num_vq = num_vq + + def get_block_size(self): + return self.block_size + + def forward(self, idxs, clip_feature): + feat = self.trans_base(idxs, clip_feature) + logits = self.trans_head(feat) + return logits + + def sample(self, clip_feature, if_categorial=False): + for k in range(self.block_size): + if k == 0: + x = [] + else: + x = xs + logits = self.forward(x, clip_feature) + logits = logits[:, -1, :] + probs = F.softmax(logits, dim=-1) + if if_categorial: + dist = Categorical(probs) + idx = dist.sample() + if idx == self.num_vq: + break + idx = idx.unsqueeze(-1) + else: + _, idx = torch.topk(probs, k=1, dim=-1) + if idx[0] == self.num_vq: + break + # append to the sequence and continue + if k == 0: + xs = idx + else: + xs = torch.cat((xs, idx), dim=1) + + if k == self.block_size - 1: + return xs[:, :-1] + return xs + +class CausalCrossConditionalSelfAttention(nn.Module): + + def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1): + super().__init__() + assert embed_dim % 8 == 0 + # key, query, value projections for all heads + self.key = nn.Linear(embed_dim, embed_dim) + self.query = nn.Linear(embed_dim, embed_dim) + self.value = nn.Linear(embed_dim, embed_dim) + + self.attn_drop = nn.Dropout(drop_out_rate) + self.resid_drop = nn.Dropout(drop_out_rate) + + self.proj = nn.Linear(embed_dim, embed_dim) + # causal mask to ensure that attention is only applied to the left in the input sequence + self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)) + self.n_head = n_head + + def forward(self, x): + B, T, C = x.size() + + # calculate query, key, values for all heads in batch and move head forward to be the batch dim + k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf')) + att = F.softmax(att, dim=-1) + att = self.attn_drop(att) + y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) + y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side + + # output projection + y = self.resid_drop(self.proj(y)) + return y + +class Block(nn.Module): + + def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1, fc_rate=4): + super().__init__() + self.ln1 = nn.LayerNorm(embed_dim) + self.ln2 = nn.LayerNorm(embed_dim) + self.attn = CausalCrossConditionalSelfAttention(embed_dim, block_size, n_head, drop_out_rate) + self.mlp = nn.Sequential( + nn.Linear(embed_dim, fc_rate * embed_dim), + nn.GELU(), + nn.Linear(fc_rate * embed_dim, embed_dim), + nn.Dropout(drop_out_rate), + ) + + def forward(self, x): + x = x + self.attn(self.ln1(x)) + x = x + self.mlp(self.ln2(x)) + return x + +class CrossCondTransBase(nn.Module): + + def __init__(self, + num_vq=1024, + embed_dim=512, + clip_dim=512, + block_size=16, + num_layers=2, + n_head=8, + drop_out_rate=0.1, + fc_rate=4): + super().__init__() + self.tok_emb = nn.Embedding(num_vq + 2, embed_dim) + self.cond_emb = nn.Linear(clip_dim, embed_dim) + self.pos_embedding = nn.Embedding(block_size, embed_dim) + self.drop = nn.Dropout(drop_out_rate) + # transformer block + self.blocks = nn.Sequential(*[Block(embed_dim, block_size, n_head, drop_out_rate, fc_rate) for _ in range(num_layers)]) + self.pos_embed = pos_encoding.PositionEmbedding(block_size, embed_dim, 0.0, False) + + self.block_size = block_size + + self.apply(self._init_weights) + + def get_block_size(self): + return self.block_size + + def _init_weights(self, module): + if isinstance(module, (nn.Linear, nn.Embedding)): + module.weight.data.normal_(mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + def forward(self, idx, clip_feature): + if len(idx) == 0: + token_embeddings = self.cond_emb(clip_feature).unsqueeze(1) + else: + b, t = idx.size() + assert t <= self.block_size, "Cannot forward, model block size is exhausted." + # forward the Trans model + token_embeddings = self.tok_emb(idx) + token_embeddings = torch.cat([self.cond_emb(clip_feature).unsqueeze(1), token_embeddings], dim=1) + + x = self.pos_embed(token_embeddings) + x = self.blocks(x) + + return x + + +class CrossCondTransHead(nn.Module): + + def __init__(self, + num_vq=1024, + embed_dim=512, + block_size=16, + num_layers=2, + n_head=8, + drop_out_rate=0.1, + fc_rate=4): + super().__init__() + + self.blocks = nn.Sequential(*[Block(embed_dim, block_size, n_head, drop_out_rate, fc_rate) for _ in range(num_layers)]) + self.ln_f = nn.LayerNorm(embed_dim) + self.head = nn.Linear(embed_dim, num_vq + 1, bias=False) + self.block_size = block_size + + self.apply(self._init_weights) + + def get_block_size(self): + return self.block_size + + def _init_weights(self, module): + if isinstance(module, (nn.Linear, nn.Embedding)): + module.weight.data.normal_(mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + def forward(self, x): + x = self.blocks(x) + x = self.ln_f(x) + logits = self.head(x) + return logits + + + + + + diff --git a/VQ-Trans/models/vqvae.py b/VQ-Trans/models/vqvae.py new file mode 100644 index 0000000000000000000000000000000000000000..7e6c940674d460853e8418514bf2306f774689fd --- /dev/null +++ b/VQ-Trans/models/vqvae.py @@ -0,0 +1,118 @@ +import torch.nn as nn +from models.encdec import Encoder, Decoder +from models.quantize_cnn import QuantizeEMAReset, Quantizer, QuantizeEMA, QuantizeReset + + +class VQVAE_251(nn.Module): + def __init__(self, + args, + nb_code=1024, + code_dim=512, + output_emb_width=512, + down_t=3, + stride_t=2, + width=512, + depth=3, + dilation_growth_rate=3, + activation='relu', + norm=None): + + super().__init__() + self.code_dim = code_dim + self.num_code = nb_code + self.quant = args.quantizer + self.encoder = Encoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm) + self.decoder = Decoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm) + if args.quantizer == "ema_reset": + self.quantizer = QuantizeEMAReset(nb_code, code_dim, args) + elif args.quantizer == "orig": + self.quantizer = Quantizer(nb_code, code_dim, 1.0) + elif args.quantizer == "ema": + self.quantizer = QuantizeEMA(nb_code, code_dim, args) + elif args.quantizer == "reset": + self.quantizer = QuantizeReset(nb_code, code_dim, args) + + + def preprocess(self, x): + # (bs, T, Jx3) -> (bs, Jx3, T) + x = x.permute(0,2,1).float() + return x + + + def postprocess(self, x): + # (bs, Jx3, T) -> (bs, T, Jx3) + x = x.permute(0,2,1) + return x + + + def encode(self, x): + N, T, _ = x.shape + x_in = self.preprocess(x) + x_encoder = self.encoder(x_in) + x_encoder = self.postprocess(x_encoder) + x_encoder = x_encoder.contiguous().view(-1, x_encoder.shape[-1]) # (NT, C) + code_idx = self.quantizer.quantize(x_encoder) + code_idx = code_idx.view(N, -1) + return code_idx + + + def forward(self, x): + + x_in = self.preprocess(x) + # Encode + x_encoder = self.encoder(x_in) + + ## quantization + x_quantized, loss, perplexity = self.quantizer(x_encoder) + + ## decoder + x_decoder = self.decoder(x_quantized) + x_out = self.postprocess(x_decoder) + return x_out, loss, perplexity + + + def forward_decoder(self, x): + x_d = self.quantizer.dequantize(x) + x_d = x_d.view(1, -1, self.code_dim).permute(0, 2, 1).contiguous() + + # decoder + x_decoder = self.decoder(x_d) + x_out = self.postprocess(x_decoder) + return x_out + + + +class HumanVQVAE(nn.Module): + def __init__(self, + args, + nb_code=512, + code_dim=512, + output_emb_width=512, + down_t=3, + stride_t=2, + width=512, + depth=3, + dilation_growth_rate=3, + activation='relu', + norm=None): + + super().__init__() + + self.nb_joints = 21 if args.dataname == 'kit' else 22 + self.vqvae = VQVAE_251(args, nb_code, code_dim, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm) + + def encode(self, x): + b, t, c = x.size() + quants = self.vqvae.encode(x) # (N, T) + return quants + + def forward(self, x): + + x_out, loss, perplexity = self.vqvae(x) + + return x_out, loss, perplexity + + def forward_decoder(self, x): + x_out = self.vqvae.forward_decoder(x) + return x_out + \ No newline at end of file diff --git a/VQ-Trans/options/get_eval_option.py b/VQ-Trans/options/get_eval_option.py new file mode 100644 index 0000000000000000000000000000000000000000..d0989ba1a8116068753ada2cb1806744e4512447 --- /dev/null +++ b/VQ-Trans/options/get_eval_option.py @@ -0,0 +1,83 @@ +from argparse import Namespace +import re +from os.path import join as pjoin + + +def is_float(numStr): + flag = False + numStr = str(numStr).strip().lstrip('-').lstrip('+') + try: + reg = re.compile(r'^[-+]?[0-9]+\.[0-9]+$') + res = reg.match(str(numStr)) + if res: + flag = True + except Exception as ex: + print("is_float() - error: " + str(ex)) + return flag + + +def is_number(numStr): + flag = False + numStr = str(numStr).strip().lstrip('-').lstrip('+') + if str(numStr).isdigit(): + flag = True + return flag + + +def get_opt(opt_path, device): + opt = Namespace() + opt_dict = vars(opt) + + skip = ('-------------- End ----------------', + '------------ Options -------------', + '\n') + print('Reading', opt_path) + with open(opt_path) as f: + for line in f: + if line.strip() not in skip: + # print(line.strip()) + key, value = line.strip().split(': ') + if value in ('True', 'False'): + opt_dict[key] = (value == 'True') + # print(key, value) + elif is_float(value): + opt_dict[key] = float(value) + elif is_number(value): + opt_dict[key] = int(value) + else: + opt_dict[key] = str(value) + + # print(opt) + opt_dict['which_epoch'] = 'finest' + opt.save_root = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name) + opt.model_dir = pjoin(opt.save_root, 'model') + opt.meta_dir = pjoin(opt.save_root, 'meta') + + if opt.dataset_name == 't2m': + opt.data_root = './dataset/HumanML3D/' + opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs') + opt.text_dir = pjoin(opt.data_root, 'texts') + opt.joints_num = 22 + opt.dim_pose = 263 + opt.max_motion_length = 196 + opt.max_motion_frame = 196 + opt.max_motion_token = 55 + elif opt.dataset_name == 'kit': + opt.data_root = './dataset/KIT-ML/' + opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs') + opt.text_dir = pjoin(opt.data_root, 'texts') + opt.joints_num = 21 + opt.dim_pose = 251 + opt.max_motion_length = 196 + opt.max_motion_frame = 196 + opt.max_motion_token = 55 + else: + raise KeyError('Dataset not recognized') + + opt.dim_word = 300 + opt.num_classes = 200 // opt.unit_length + opt.is_train = False + opt.is_continue = False + opt.device = device + + return opt \ No newline at end of file diff --git a/VQ-Trans/options/option_transformer.py b/VQ-Trans/options/option_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..cf48ce1fdac663ec44419d67721ac268806f8127 --- /dev/null +++ b/VQ-Trans/options/option_transformer.py @@ -0,0 +1,68 @@ +import argparse + +def get_args_parser(): + parser = argparse.ArgumentParser(description='Optimal Transport AutoEncoder training for Amass', + add_help=True, + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + ## dataloader + + parser.add_argument('--dataname', type=str, default='kit', help='dataset directory') + parser.add_argument('--batch-size', default=128, type=int, help='batch size') + parser.add_argument('--fps', default=[20], nargs="+", type=int, help='frames per second') + parser.add_argument('--seq-len', type=int, default=64, help='training motion length') + + ## optimization + parser.add_argument('--total-iter', default=100000, type=int, help='number of total iterations to run') + parser.add_argument('--warm-up-iter', default=1000, type=int, help='number of total iterations for warmup') + parser.add_argument('--lr', default=2e-4, type=float, help='max learning rate') + parser.add_argument('--lr-scheduler', default=[60000], nargs="+", type=int, help="learning rate schedule (iterations)") + parser.add_argument('--gamma', default=0.05, type=float, help="learning rate decay") + + parser.add_argument('--weight-decay', default=1e-6, type=float, help='weight decay') + parser.add_argument('--decay-option',default='all', type=str, choices=['all', 'noVQ'], help='disable weight decay on codebook') + parser.add_argument('--optimizer',default='adamw', type=str, choices=['adam', 'adamw'], help='disable weight decay on codebook') + + ## vqvae arch + parser.add_argument("--code-dim", type=int, default=512, help="embedding dimension") + parser.add_argument("--nb-code", type=int, default=512, help="nb of embedding") + parser.add_argument("--mu", type=float, default=0.99, help="exponential moving average to update the codebook") + parser.add_argument("--down-t", type=int, default=3, help="downsampling rate") + parser.add_argument("--stride-t", type=int, default=2, help="stride size") + parser.add_argument("--width", type=int, default=512, help="width of the network") + parser.add_argument("--depth", type=int, default=3, help="depth of the network") + parser.add_argument("--dilation-growth-rate", type=int, default=3, help="dilation growth rate") + parser.add_argument("--output-emb-width", type=int, default=512, help="output embedding width") + parser.add_argument('--vq-act', type=str, default='relu', choices = ['relu', 'silu', 'gelu'], help='dataset directory') + + ## gpt arch + parser.add_argument("--block-size", type=int, default=25, help="seq len") + parser.add_argument("--embed-dim-gpt", type=int, default=512, help="embedding dimension") + parser.add_argument("--clip-dim", type=int, default=512, help="latent dimension in the clip feature") + parser.add_argument("--num-layers", type=int, default=2, help="nb of transformer layers") + parser.add_argument("--n-head-gpt", type=int, default=8, help="nb of heads") + parser.add_argument("--ff-rate", type=int, default=4, help="feedforward size") + parser.add_argument("--drop-out-rate", type=float, default=0.1, help="dropout ratio in the pos encoding") + + ## quantizer + parser.add_argument("--quantizer", type=str, default='ema_reset', choices = ['ema', 'orig', 'ema_reset', 'reset'], help="eps for optimal transport") + parser.add_argument('--quantbeta', type=float, default=1.0, help='dataset directory') + + ## resume + parser.add_argument("--resume-pth", type=str, default=None, help='resume vq pth') + parser.add_argument("--resume-trans", type=str, default=None, help='resume gpt pth') + + + ## output directory + parser.add_argument('--out-dir', type=str, default='output_GPT_Final/', help='output directory') + parser.add_argument('--exp-name', type=str, default='exp_debug', help='name of the experiment, will create a file inside out-dir') + parser.add_argument('--vq-name', type=str, default='exp_debug', help='name of the generated dataset .npy, will create a file inside out-dir') + ## other + parser.add_argument('--print-iter', default=200, type=int, help='print frequency') + parser.add_argument('--eval-iter', default=5000, type=int, help='evaluation frequency') + parser.add_argument('--seed', default=123, type=int, help='seed for initializing training. ') + parser.add_argument("--if-maxtest", action='store_true', help="test in max") + parser.add_argument('--pkeep', type=float, default=1.0, help='keep rate for gpt training') + + + return parser.parse_args() \ No newline at end of file diff --git a/VQ-Trans/options/option_vq.py b/VQ-Trans/options/option_vq.py new file mode 100644 index 0000000000000000000000000000000000000000..08a53ff1270facc10ab44ec0647e673ed1336d0d --- /dev/null +++ b/VQ-Trans/options/option_vq.py @@ -0,0 +1,61 @@ +import argparse + +def get_args_parser(): + parser = argparse.ArgumentParser(description='Optimal Transport AutoEncoder training for AIST', + add_help=True, + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + ## dataloader + parser.add_argument('--dataname', type=str, default='kit', help='dataset directory') + parser.add_argument('--batch-size', default=128, type=int, help='batch size') + parser.add_argument('--window-size', type=int, default=64, help='training motion length') + + ## optimization + parser.add_argument('--total-iter', default=200000, type=int, help='number of total iterations to run') + parser.add_argument('--warm-up-iter', default=1000, type=int, help='number of total iterations for warmup') + parser.add_argument('--lr', default=2e-4, type=float, help='max learning rate') + parser.add_argument('--lr-scheduler', default=[50000, 400000], nargs="+", type=int, help="learning rate schedule (iterations)") + parser.add_argument('--gamma', default=0.05, type=float, help="learning rate decay") + + parser.add_argument('--weight-decay', default=0.0, type=float, help='weight decay') + parser.add_argument("--commit", type=float, default=0.02, help="hyper-parameter for the commitment loss") + parser.add_argument('--loss-vel', type=float, default=0.1, help='hyper-parameter for the velocity loss') + parser.add_argument('--recons-loss', type=str, default='l2', help='reconstruction loss') + + ## vqvae arch + parser.add_argument("--code-dim", type=int, default=512, help="embedding dimension") + parser.add_argument("--nb-code", type=int, default=512, help="nb of embedding") + parser.add_argument("--mu", type=float, default=0.99, help="exponential moving average to update the codebook") + parser.add_argument("--down-t", type=int, default=2, help="downsampling rate") + parser.add_argument("--stride-t", type=int, default=2, help="stride size") + parser.add_argument("--width", type=int, default=512, help="width of the network") + parser.add_argument("--depth", type=int, default=3, help="depth of the network") + parser.add_argument("--dilation-growth-rate", type=int, default=3, help="dilation growth rate") + parser.add_argument("--output-emb-width", type=int, default=512, help="output embedding width") + parser.add_argument('--vq-act', type=str, default='relu', choices = ['relu', 'silu', 'gelu'], help='dataset directory') + parser.add_argument('--vq-norm', type=str, default=None, help='dataset directory') + + ## quantizer + parser.add_argument("--quantizer", type=str, default='ema_reset', choices = ['ema', 'orig', 'ema_reset', 'reset'], help="eps for optimal transport") + parser.add_argument('--beta', type=float, default=1.0, help='commitment loss in standard VQ') + + ## resume + parser.add_argument("--resume-pth", type=str, default=None, help='resume pth for VQ') + parser.add_argument("--resume-gpt", type=str, default=None, help='resume pth for GPT') + + + ## output directory + parser.add_argument('--out-dir', type=str, default='output_vqfinal/', help='output directory') + parser.add_argument('--results-dir', type=str, default='visual_results/', help='output directory') + parser.add_argument('--visual-name', type=str, default='baseline', help='output directory') + parser.add_argument('--exp-name', type=str, default='exp_debug', help='name of the experiment, will create a file inside out-dir') + ## other + parser.add_argument('--print-iter', default=200, type=int, help='print frequency') + parser.add_argument('--eval-iter', default=1000, type=int, help='evaluation frequency') + parser.add_argument('--seed', default=123, type=int, help='seed for initializing training.') + + parser.add_argument('--vis-gt', action='store_true', help='whether visualize GT motions') + parser.add_argument('--nb-vis', default=20, type=int, help='nb of visualizations') + + + return parser.parse_args() \ No newline at end of file diff --git a/VQ-Trans/render_final.py b/VQ-Trans/render_final.py new file mode 100644 index 0000000000000000000000000000000000000000..41b3bfdb2e6bff74aeaceb8f1a7ebac9dc1acaba --- /dev/null +++ b/VQ-Trans/render_final.py @@ -0,0 +1,194 @@ +from models.rotation2xyz import Rotation2xyz +import numpy as np +from trimesh import Trimesh +import os +os.environ['PYOPENGL_PLATFORM'] = "osmesa" + +import torch +from visualize.simplify_loc2rot import joints2smpl +import pyrender +import matplotlib.pyplot as plt + +import io +import imageio +from shapely import geometry +import trimesh +from pyrender.constants import RenderFlags +import math +# import ffmpeg +from PIL import Image + +class WeakPerspectiveCamera(pyrender.Camera): + def __init__(self, + scale, + translation, + znear=pyrender.camera.DEFAULT_Z_NEAR, + zfar=None, + name=None): + super(WeakPerspectiveCamera, self).__init__( + znear=znear, + zfar=zfar, + name=name, + ) + self.scale = scale + self.translation = translation + + def get_projection_matrix(self, width=None, height=None): + P = np.eye(4) + P[0, 0] = self.scale[0] + P[1, 1] = self.scale[1] + P[0, 3] = self.translation[0] * self.scale[0] + P[1, 3] = -self.translation[1] * self.scale[1] + P[2, 2] = -1 + return P + +def render(motions, outdir='test_vis', device_id=0, name=None, pred=True): + frames, njoints, nfeats = motions.shape + MINS = motions.min(axis=0).min(axis=0) + MAXS = motions.max(axis=0).max(axis=0) + + height_offset = MINS[1] + motions[:, :, 1] -= height_offset + trajec = motions[:, 0, [0, 2]] + + j2s = joints2smpl(num_frames=frames, device_id=0, cuda=True) + rot2xyz = Rotation2xyz(device=torch.device("cuda:0")) + faces = rot2xyz.smpl_model.faces + + if (not os.path.exists(outdir + name+'_pred.pt') and pred) or (not os.path.exists(outdir + name+'_gt.pt') and not pred): + print(f'Running SMPLify, it may take a few minutes.') + motion_tensor, opt_dict = j2s.joint2smpl(motions) # [nframes, njoints, 3] + + vertices = rot2xyz(torch.tensor(motion_tensor).clone(), mask=None, + pose_rep='rot6d', translation=True, glob=True, + jointstype='vertices', + vertstrans=True) + + if pred: + torch.save(vertices, outdir + name+'_pred.pt') + else: + torch.save(vertices, outdir + name+'_gt.pt') + else: + if pred: + vertices = torch.load(outdir + name+'_pred.pt') + else: + vertices = torch.load(outdir + name+'_gt.pt') + frames = vertices.shape[3] # shape: 1, nb_frames, 3, nb_joints + print (vertices.shape) + MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0] + MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0] + # vertices[:,:,1,:] -= MINS[1] + 1e-5 + + + out_list = [] + + minx = MINS[0] - 0.5 + maxx = MAXS[0] + 0.5 + minz = MINS[2] - 0.5 + maxz = MAXS[2] + 0.5 + polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]]) + polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5) + + vid = [] + for i in range(frames): + if i % 10 == 0: + print(i) + + mesh = Trimesh(vertices=vertices[0, :, :, i].squeeze().tolist(), faces=faces) + + base_color = (0.11, 0.53, 0.8, 0.5) + ## OPAQUE rendering without alpha + ## BLEND rendering consider alpha + material = pyrender.MetallicRoughnessMaterial( + metallicFactor=0.7, + alphaMode='OPAQUE', + baseColorFactor=base_color + ) + + + mesh = pyrender.Mesh.from_trimesh(mesh, material=material) + + polygon_mesh.visual.face_colors = [0, 0, 0, 0.21] + polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False) + + bg_color = [1, 1, 1, 0.8] + scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4)) + + sx, sy, tx, ty = [0.75, 0.75, 0, 0.10] + + camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0)) + + light = pyrender.DirectionalLight(color=[1,1,1], intensity=300) + + scene.add(mesh) + + c = np.pi / 2 + + scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0], + + [ 0, np.cos(c), -np.sin(c), MINS[1].cpu().numpy()], + + [ 0, np.sin(c), np.cos(c), 0], + + [ 0, 0, 0, 1]])) + + light_pose = np.eye(4) + light_pose[:3, 3] = [0, -1, 1] + scene.add(light, pose=light_pose.copy()) + + light_pose[:3, 3] = [0, 1, 1] + scene.add(light, pose=light_pose.copy()) + + light_pose[:3, 3] = [1, 1, 2] + scene.add(light, pose=light_pose.copy()) + + + c = -np.pi / 6 + + scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx).cpu().numpy()/2], + + [ 0, np.cos(c), -np.sin(c), 1.5], + + [ 0, np.sin(c), np.cos(c), max(4, minz.cpu().numpy()+(1.5-MINS[1].cpu().numpy())*2, (maxx-minx).cpu().numpy())], + + [ 0, 0, 0, 1] + ]) + + # render scene + r = pyrender.OffscreenRenderer(960, 960) + + color, _ = r.render(scene, flags=RenderFlags.RGBA) + # Image.fromarray(color).save(outdir+'/'+name+'_'+str(i)+'.png') + + vid.append(color) + + r.delete() + + out = np.stack(vid, axis=0) + if pred: + imageio.mimsave(outdir + name+'_pred.gif', out, fps=20) + else: + imageio.mimsave(outdir + name+'_gt.gif', out, fps=20) + + + + + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--filedir", type=str, default=None, help='motion npy file dir') + parser.add_argument('--motion-list', default=None, nargs="+", type=str, help="motion name list") + args = parser.parse_args() + + filename_list = args.motion_list + filedir = args.filedir + + for filename in filename_list: + motions = np.load(filedir + filename+'_pred.npy') + print('pred', motions.shape, filename) + render(motions[0], outdir=filedir, device_id=0, name=filename, pred=True) + + motions = np.load(filedir + filename+'_gt.npy') + print('gt', motions.shape, filename) + render(motions[0], outdir=filedir, device_id=0, name=filename, pred=False) diff --git a/VQ-Trans/train_t2m_trans.py b/VQ-Trans/train_t2m_trans.py new file mode 100644 index 0000000000000000000000000000000000000000..8da444f87aa7ca71cd8bc3604868cf30a6c70e02 --- /dev/null +++ b/VQ-Trans/train_t2m_trans.py @@ -0,0 +1,191 @@ +import os +import torch +import numpy as np + +from torch.utils.tensorboard import SummaryWriter +from os.path import join as pjoin +from torch.distributions import Categorical +import json +import clip + +import options.option_transformer as option_trans +import models.vqvae as vqvae +import utils.utils_model as utils_model +import utils.eval_trans as eval_trans +from dataset import dataset_TM_train +from dataset import dataset_TM_eval +from dataset import dataset_tokenize +import models.t2m_trans as trans +from options.get_eval_option import get_opt +from models.evaluator_wrapper import EvaluatorModelWrapper +import warnings +warnings.filterwarnings('ignore') + +##### ---- Exp dirs ---- ##### +args = option_trans.get_args_parser() +torch.manual_seed(args.seed) + +args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}') +args.vq_dir= os.path.join("./dataset/KIT-ML" if args.dataname == 'kit' else "./dataset/HumanML3D", f'{args.vq_name}') +os.makedirs(args.out_dir, exist_ok = True) +os.makedirs(args.vq_dir, exist_ok = True) + +##### ---- Logger ---- ##### +logger = utils_model.get_logger(args.out_dir) +writer = SummaryWriter(args.out_dir) +logger.info(json.dumps(vars(args), indent=4, sort_keys=True)) + +##### ---- Dataloader ---- ##### +train_loader_token = dataset_tokenize.DATALoader(args.dataname, 1, unit_length=2**args.down_t) + +from utils.word_vectorizer import WordVectorizer +w_vectorizer = WordVectorizer('./glove', 'our_vab') +val_loader = dataset_TM_eval.DATALoader(args.dataname, False, 32, w_vectorizer) + +dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataname == 'kit' else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt' + +wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda')) +eval_wrapper = EvaluatorModelWrapper(wrapper_opt) + +##### ---- Network ---- ##### +clip_model, clip_preprocess = clip.load("ViT-B/32", device=torch.device('cuda'), jit=False, download_root='/apdcephfs_cq2/share_1290939/maelyszhang/.cache/clip') # Must set jit=False for training +clip.model.convert_weights(clip_model) # Actually this line is unnecessary since clip by default already on float16 +clip_model.eval() +for p in clip_model.parameters(): + p.requires_grad = False + +net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers + args.nb_code, + args.code_dim, + args.output_emb_width, + args.down_t, + args.stride_t, + args.width, + args.depth, + args.dilation_growth_rate) + + +trans_encoder = trans.Text2Motion_Transformer(num_vq=args.nb_code, + embed_dim=args.embed_dim_gpt, + clip_dim=args.clip_dim, + block_size=args.block_size, + num_layers=args.num_layers, + n_head=args.n_head_gpt, + drop_out_rate=args.drop_out_rate, + fc_rate=args.ff_rate) + + +print ('loading checkpoint from {}'.format(args.resume_pth)) +ckpt = torch.load(args.resume_pth, map_location='cpu') +net.load_state_dict(ckpt['net'], strict=True) +net.eval() +net.cuda() + +if args.resume_trans is not None: + print ('loading transformer checkpoint from {}'.format(args.resume_trans)) + ckpt = torch.load(args.resume_trans, map_location='cpu') + trans_encoder.load_state_dict(ckpt['trans'], strict=True) +trans_encoder.train() +trans_encoder.cuda() + +##### ---- Optimizer & Scheduler ---- ##### +optimizer = utils_model.initial_optim(args.decay_option, args.lr, args.weight_decay, trans_encoder, args.optimizer) +scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_scheduler, gamma=args.gamma) + +##### ---- Optimization goals ---- ##### +loss_ce = torch.nn.CrossEntropyLoss() + +nb_iter, avg_loss_cls, avg_acc = 0, 0., 0. +right_num = 0 +nb_sample_train = 0 + +##### ---- get code ---- ##### +for batch in train_loader_token: + pose, name = batch + bs, seq = pose.shape[0], pose.shape[1] + + pose = pose.cuda().float() # bs, nb_joints, joints_dim, seq_len + target = net.encode(pose) + target = target.cpu().numpy() + np.save(pjoin(args.vq_dir, name[0] +'.npy'), target) + + +train_loader = dataset_TM_train.DATALoader(args.dataname, args.batch_size, args.nb_code, args.vq_name, unit_length=2**args.down_t) +train_loader_iter = dataset_TM_train.cycle(train_loader) + + +##### ---- Training ---- ##### +best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_transformer(args.out_dir, val_loader, net, trans_encoder, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, clip_model=clip_model, eval_wrapper=eval_wrapper) +while nb_iter <= args.total_iter: + + batch = next(train_loader_iter) + clip_text, m_tokens, m_tokens_len = batch + m_tokens, m_tokens_len = m_tokens.cuda(), m_tokens_len.cuda() + bs = m_tokens.shape[0] + target = m_tokens # (bs, 26) + target = target.cuda() + + text = clip.tokenize(clip_text, truncate=True).cuda() + + feat_clip_text = clip_model.encode_text(text).float() + + input_index = target[:,:-1] + + if args.pkeep == -1: + proba = np.random.rand(1)[0] + mask = torch.bernoulli(proba * torch.ones(input_index.shape, + device=input_index.device)) + else: + mask = torch.bernoulli(args.pkeep * torch.ones(input_index.shape, + device=input_index.device)) + mask = mask.round().to(dtype=torch.int64) + r_indices = torch.randint_like(input_index, args.nb_code) + a_indices = mask*input_index+(1-mask)*r_indices + + cls_pred = trans_encoder(a_indices, feat_clip_text) + cls_pred = cls_pred.contiguous() + + loss_cls = 0.0 + for i in range(bs): + # loss function (26), (26, 513) + loss_cls += loss_ce(cls_pred[i][:m_tokens_len[i] + 1], target[i][:m_tokens_len[i] + 1]) / bs + + # Accuracy + probs = torch.softmax(cls_pred[i][:m_tokens_len[i] + 1], dim=-1) + + if args.if_maxtest: + _, cls_pred_index = torch.max(probs, dim=-1) + + else: + dist = Categorical(probs) + cls_pred_index = dist.sample() + right_num += (cls_pred_index.flatten(0) == target[i][:m_tokens_len[i] + 1].flatten(0)).sum().item() + + ## global loss + optimizer.zero_grad() + loss_cls.backward() + optimizer.step() + scheduler.step() + + avg_loss_cls = avg_loss_cls + loss_cls.item() + nb_sample_train = nb_sample_train + (m_tokens_len + 1).sum().item() + + nb_iter += 1 + if nb_iter % args.print_iter == 0 : + avg_loss_cls = avg_loss_cls / args.print_iter + avg_acc = right_num * 100 / nb_sample_train + writer.add_scalar('./Loss/train', avg_loss_cls, nb_iter) + writer.add_scalar('./ACC/train', avg_acc, nb_iter) + msg = f"Train. Iter {nb_iter} : Loss. {avg_loss_cls:.5f}, ACC. {avg_acc:.4f}" + logger.info(msg) + avg_loss_cls = 0. + right_num = 0 + nb_sample_train = 0 + + if nb_iter % args.eval_iter == 0: + best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_transformer(args.out_dir, val_loader, net, trans_encoder, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, clip_model=clip_model, eval_wrapper=eval_wrapper) + + if nb_iter == args.total_iter: + msg_final = f"Train. Iter {best_iter} : FID. {best_fid:.5f}, Diversity. {best_div:.4f}, TOP1. {best_top1:.4f}, TOP2. {best_top2:.4f}, TOP3. {best_top3:.4f}" + logger.info(msg_final) + break \ No newline at end of file diff --git a/VQ-Trans/train_vq.py b/VQ-Trans/train_vq.py new file mode 100644 index 0000000000000000000000000000000000000000..d89b9930ba1262747542df3d5b2f03f8fab1b04a --- /dev/null +++ b/VQ-Trans/train_vq.py @@ -0,0 +1,171 @@ +import os +import json + +import torch +import torch.optim as optim +from torch.utils.tensorboard import SummaryWriter + +import models.vqvae as vqvae +import utils.losses as losses +import options.option_vq as option_vq +import utils.utils_model as utils_model +from dataset import dataset_VQ, dataset_TM_eval +import utils.eval_trans as eval_trans +from options.get_eval_option import get_opt +from models.evaluator_wrapper import EvaluatorModelWrapper +import warnings +warnings.filterwarnings('ignore') +from utils.word_vectorizer import WordVectorizer + +def update_lr_warm_up(optimizer, nb_iter, warm_up_iter, lr): + + current_lr = lr * (nb_iter + 1) / (warm_up_iter + 1) + for param_group in optimizer.param_groups: + param_group["lr"] = current_lr + + return optimizer, current_lr + +##### ---- Exp dirs ---- ##### +args = option_vq.get_args_parser() +torch.manual_seed(args.seed) + +args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}') +os.makedirs(args.out_dir, exist_ok = True) + +##### ---- Logger ---- ##### +logger = utils_model.get_logger(args.out_dir) +writer = SummaryWriter(args.out_dir) +logger.info(json.dumps(vars(args), indent=4, sort_keys=True)) + + + +w_vectorizer = WordVectorizer('./glove', 'our_vab') + +if args.dataname == 'kit' : + dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' + args.nb_joints = 21 + +else : + dataset_opt_path = 'checkpoints/t2m/Comp_v6_KLD005/opt.txt' + args.nb_joints = 22 + +logger.info(f'Training on {args.dataname}, motions are with {args.nb_joints} joints') + +wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda')) +eval_wrapper = EvaluatorModelWrapper(wrapper_opt) + + +##### ---- Dataloader ---- ##### +train_loader = dataset_VQ.DATALoader(args.dataname, + args.batch_size, + window_size=args.window_size, + unit_length=2**args.down_t) + +train_loader_iter = dataset_VQ.cycle(train_loader) + +val_loader = dataset_TM_eval.DATALoader(args.dataname, False, + 32, + w_vectorizer, + unit_length=2**args.down_t) + +##### ---- Network ---- ##### +net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers + args.nb_code, + args.code_dim, + args.output_emb_width, + args.down_t, + args.stride_t, + args.width, + args.depth, + args.dilation_growth_rate, + args.vq_act, + args.vq_norm) + + +if args.resume_pth : + logger.info('loading checkpoint from {}'.format(args.resume_pth)) + ckpt = torch.load(args.resume_pth, map_location='cpu') + net.load_state_dict(ckpt['net'], strict=True) +net.train() +net.cuda() + +##### ---- Optimizer & Scheduler ---- ##### +optimizer = optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.weight_decay) +scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_scheduler, gamma=args.gamma) + + +Loss = losses.ReConsLoss(args.recons_loss, args.nb_joints) + +##### ------ warm-up ------- ##### +avg_recons, avg_perplexity, avg_commit = 0., 0., 0. + +for nb_iter in range(1, args.warm_up_iter): + + optimizer, current_lr = update_lr_warm_up(optimizer, nb_iter, args.warm_up_iter, args.lr) + + gt_motion = next(train_loader_iter) + gt_motion = gt_motion.cuda().float() # (bs, 64, dim) + + pred_motion, loss_commit, perplexity = net(gt_motion) + loss_motion = Loss(pred_motion, gt_motion) + loss_vel = Loss.forward_vel(pred_motion, gt_motion) + + loss = loss_motion + args.commit * loss_commit + args.loss_vel * loss_vel + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + avg_recons += loss_motion.item() + avg_perplexity += perplexity.item() + avg_commit += loss_commit.item() + + if nb_iter % args.print_iter == 0 : + avg_recons /= args.print_iter + avg_perplexity /= args.print_iter + avg_commit /= args.print_iter + + logger.info(f"Warmup. Iter {nb_iter} : lr {current_lr:.5f} \t Commit. {avg_commit:.5f} \t PPL. {avg_perplexity:.2f} \t Recons. {avg_recons:.5f}") + + avg_recons, avg_perplexity, avg_commit = 0., 0., 0. + +##### ---- Training ---- ##### +avg_recons, avg_perplexity, avg_commit = 0., 0., 0. +best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, eval_wrapper=eval_wrapper) + +for nb_iter in range(1, args.total_iter + 1): + + gt_motion = next(train_loader_iter) + gt_motion = gt_motion.cuda().float() # bs, nb_joints, joints_dim, seq_len + + pred_motion, loss_commit, perplexity = net(gt_motion) + loss_motion = Loss(pred_motion, gt_motion) + loss_vel = Loss.forward_vel(pred_motion, gt_motion) + + loss = loss_motion + args.commit * loss_commit + args.loss_vel * loss_vel + + optimizer.zero_grad() + loss.backward() + optimizer.step() + scheduler.step() + + avg_recons += loss_motion.item() + avg_perplexity += perplexity.item() + avg_commit += loss_commit.item() + + if nb_iter % args.print_iter == 0 : + avg_recons /= args.print_iter + avg_perplexity /= args.print_iter + avg_commit /= args.print_iter + + writer.add_scalar('./Train/L1', avg_recons, nb_iter) + writer.add_scalar('./Train/PPL', avg_perplexity, nb_iter) + writer.add_scalar('./Train/Commit', avg_commit, nb_iter) + + logger.info(f"Train. Iter {nb_iter} : \t Commit. {avg_commit:.5f} \t PPL. {avg_perplexity:.2f} \t Recons. {avg_recons:.5f}") + + avg_recons, avg_perplexity, avg_commit = 0., 0., 0., + + if nb_iter % args.eval_iter==0 : + best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, eval_wrapper=eval_wrapper) + \ No newline at end of file diff --git a/VQ-Trans/utils/config.py b/VQ-Trans/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..091d790e963959c326917688ee267e6a4ec136d1 --- /dev/null +++ b/VQ-Trans/utils/config.py @@ -0,0 +1,17 @@ +import os + +SMPL_DATA_PATH = "./body_models/smpl" + +SMPL_KINTREE_PATH = os.path.join(SMPL_DATA_PATH, "kintree_table.pkl") +SMPL_MODEL_PATH = os.path.join(SMPL_DATA_PATH, "SMPL_NEUTRAL.pkl") +JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(SMPL_DATA_PATH, 'J_regressor_extra.npy') + +ROT_CONVENTION_TO_ROT_NUMBER = { + 'legacy': 23, + 'no_hands': 21, + 'full_hands': 51, + 'mitten_hands': 33, +} + +GENDERS = ['neutral', 'male', 'female'] +NUM_BETAS = 10 \ No newline at end of file diff --git a/VQ-Trans/utils/eval_trans.py b/VQ-Trans/utils/eval_trans.py new file mode 100644 index 0000000000000000000000000000000000000000..8778bb8cb7e7a320e5f7f2f3b43c7ba0b4c285ab --- /dev/null +++ b/VQ-Trans/utils/eval_trans.py @@ -0,0 +1,580 @@ +import os + +import clip +import numpy as np +import torch +from scipy import linalg + +import visualization.plot_3d_global as plot_3d +from utils.motion_process import recover_from_ric + + +def tensorborad_add_video_xyz(writer, xyz, nb_iter, tag, nb_vis=4, title_batch=None, outname=None): + xyz = xyz[:1] + bs, seq = xyz.shape[:2] + xyz = xyz.reshape(bs, seq, -1, 3) + plot_xyz = plot_3d.draw_to_batch(xyz.cpu().numpy(),title_batch, outname) + plot_xyz =np.transpose(plot_xyz, (0, 1, 4, 2, 3)) + writer.add_video(tag, plot_xyz, nb_iter, fps = 20) + +@torch.no_grad() +def evaluation_vqvae(out_dir, val_loader, net, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, eval_wrapper, draw = True, save = True, savegif=False, savenpy=False) : + net.eval() + nb_sample = 0 + + draw_org = [] + draw_pred = [] + draw_text = [] + + + motion_annotation_list = [] + motion_pred_list = [] + + R_precision_real = 0 + R_precision = 0 + + nb_sample = 0 + matching_score_real = 0 + matching_score_pred = 0 + for batch in val_loader: + word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, token, name = batch + + motion = motion.cuda() + et, em = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, motion, m_length) + bs, seq = motion.shape[0], motion.shape[1] + + num_joints = 21 if motion.shape[-1] == 251 else 22 + + pred_pose_eval = torch.zeros((bs, seq, motion.shape[-1])).cuda() + + for i in range(bs): + pose = val_loader.dataset.inv_transform(motion[i:i+1, :m_length[i], :].detach().cpu().numpy()) + pose_xyz = recover_from_ric(torch.from_numpy(pose).float().cuda(), num_joints) + + + pred_pose, loss_commit, perplexity = net(motion[i:i+1, :m_length[i]]) + pred_denorm = val_loader.dataset.inv_transform(pred_pose.detach().cpu().numpy()) + pred_xyz = recover_from_ric(torch.from_numpy(pred_denorm).float().cuda(), num_joints) + + if savenpy: + np.save(os.path.join(out_dir, name[i]+'_gt.npy'), pose_xyz[:, :m_length[i]].cpu().numpy()) + np.save(os.path.join(out_dir, name[i]+'_pred.npy'), pred_xyz.detach().cpu().numpy()) + + pred_pose_eval[i:i+1,:m_length[i],:] = pred_pose + + if i < min(4, bs): + draw_org.append(pose_xyz) + draw_pred.append(pred_xyz) + draw_text.append(caption[i]) + + et_pred, em_pred = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, pred_pose_eval, m_length) + + motion_pred_list.append(em_pred) + motion_annotation_list.append(em) + + temp_R, temp_match = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True) + R_precision_real += temp_R + matching_score_real += temp_match + temp_R, temp_match = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True) + R_precision += temp_R + matching_score_pred += temp_match + + nb_sample += bs + + motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy() + motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy() + gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np) + mu, cov= calculate_activation_statistics(motion_pred_np) + + diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100) + diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100) + + R_precision_real = R_precision_real / nb_sample + R_precision = R_precision / nb_sample + + matching_score_real = matching_score_real / nb_sample + matching_score_pred = matching_score_pred / nb_sample + + fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov) + + msg = f"--> \t Eva. Iter {nb_iter} :, FID. {fid:.4f}, Diversity Real. {diversity_real:.4f}, Diversity. {diversity:.4f}, R_precision_real. {R_precision_real}, R_precision. {R_precision}, matching_score_real. {matching_score_real}, matching_score_pred. {matching_score_pred}" + logger.info(msg) + + if draw: + writer.add_scalar('./Test/FID', fid, nb_iter) + writer.add_scalar('./Test/Diversity', diversity, nb_iter) + writer.add_scalar('./Test/top1', R_precision[0], nb_iter) + writer.add_scalar('./Test/top2', R_precision[1], nb_iter) + writer.add_scalar('./Test/top3', R_precision[2], nb_iter) + writer.add_scalar('./Test/matching_score', matching_score_pred, nb_iter) + + + if nb_iter % 5000 == 0 : + for ii in range(4): + tensorborad_add_video_xyz(writer, draw_org[ii], nb_iter, tag='./Vis/org_eval'+str(ii), nb_vis=1, title_batch=[draw_text[ii]], outname=[os.path.join(out_dir, 'gt'+str(ii)+'.gif')] if savegif else None) + + if nb_iter % 5000 == 0 : + for ii in range(4): + tensorborad_add_video_xyz(writer, draw_pred[ii], nb_iter, tag='./Vis/pred_eval'+str(ii), nb_vis=1, title_batch=[draw_text[ii]], outname=[os.path.join(out_dir, 'pred'+str(ii)+'.gif')] if savegif else None) + + + if fid < best_fid : + msg = f"--> --> \t FID Improved from {best_fid:.5f} to {fid:.5f} !!!" + logger.info(msg) + best_fid, best_iter = fid, nb_iter + if save: + torch.save({'net' : net.state_dict()}, os.path.join(out_dir, 'net_best_fid.pth')) + + if abs(diversity_real - diversity) < abs(diversity_real - best_div) : + msg = f"--> --> \t Diversity Improved from {best_div:.5f} to {diversity:.5f} !!!" + logger.info(msg) + best_div = diversity + if save: + torch.save({'net' : net.state_dict()}, os.path.join(out_dir, 'net_best_div.pth')) + + if R_precision[0] > best_top1 : + msg = f"--> --> \t Top1 Improved from {best_top1:.4f} to {R_precision[0]:.4f} !!!" + logger.info(msg) + best_top1 = R_precision[0] + if save: + torch.save({'net' : net.state_dict()}, os.path.join(out_dir, 'net_best_top1.pth')) + + if R_precision[1] > best_top2 : + msg = f"--> --> \t Top2 Improved from {best_top2:.4f} to {R_precision[1]:.4f} !!!" + logger.info(msg) + best_top2 = R_precision[1] + + if R_precision[2] > best_top3 : + msg = f"--> --> \t Top3 Improved from {best_top3:.4f} to {R_precision[2]:.4f} !!!" + logger.info(msg) + best_top3 = R_precision[2] + + if matching_score_pred < best_matching : + msg = f"--> --> \t matching_score Improved from {best_matching:.5f} to {matching_score_pred:.5f} !!!" + logger.info(msg) + best_matching = matching_score_pred + if save: + torch.save({'net' : net.state_dict()}, os.path.join(out_dir, 'net_best_matching.pth')) + + if save: + torch.save({'net' : net.state_dict()}, os.path.join(out_dir, 'net_last.pth')) + + net.train() + return best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger + + +@torch.no_grad() +def evaluation_transformer(out_dir, val_loader, net, trans, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, clip_model, eval_wrapper, draw = True, save = True, savegif=False) : + + trans.eval() + nb_sample = 0 + + draw_org = [] + draw_pred = [] + draw_text = [] + draw_text_pred = [] + + motion_annotation_list = [] + motion_pred_list = [] + R_precision_real = 0 + R_precision = 0 + matching_score_real = 0 + matching_score_pred = 0 + + nb_sample = 0 + for i in range(1): + for batch in val_loader: + word_embeddings, pos_one_hots, clip_text, sent_len, pose, m_length, token, name = batch + + bs, seq = pose.shape[:2] + num_joints = 21 if pose.shape[-1] == 251 else 22 + + text = clip.tokenize(clip_text, truncate=True).cuda() + + feat_clip_text = clip_model.encode_text(text).float() + pred_pose_eval = torch.zeros((bs, seq, pose.shape[-1])).cuda() + pred_len = torch.ones(bs).long() + + for k in range(bs): + try: + index_motion = trans.sample(feat_clip_text[k:k+1], False) + except: + index_motion = torch.ones(1,1).cuda().long() + + pred_pose = net.forward_decoder(index_motion) + cur_len = pred_pose.shape[1] + + pred_len[k] = min(cur_len, seq) + pred_pose_eval[k:k+1, :cur_len] = pred_pose[:, :seq] + + if draw: + pred_denorm = val_loader.dataset.inv_transform(pred_pose.detach().cpu().numpy()) + pred_xyz = recover_from_ric(torch.from_numpy(pred_denorm).float().cuda(), num_joints) + + if i == 0 and k < 4: + draw_pred.append(pred_xyz) + draw_text_pred.append(clip_text[k]) + + et_pred, em_pred = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, pred_pose_eval, pred_len) + + if i == 0: + pose = pose.cuda().float() + + et, em = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, pose, m_length) + motion_annotation_list.append(em) + motion_pred_list.append(em_pred) + + if draw: + pose = val_loader.dataset.inv_transform(pose.detach().cpu().numpy()) + pose_xyz = recover_from_ric(torch.from_numpy(pose).float().cuda(), num_joints) + + + for j in range(min(4, bs)): + draw_org.append(pose_xyz[j][:m_length[j]].unsqueeze(0)) + draw_text.append(clip_text[j]) + + temp_R, temp_match = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True) + R_precision_real += temp_R + matching_score_real += temp_match + temp_R, temp_match = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True) + R_precision += temp_R + matching_score_pred += temp_match + + nb_sample += bs + + motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy() + motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy() + gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np) + mu, cov= calculate_activation_statistics(motion_pred_np) + + diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100) + diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100) + + R_precision_real = R_precision_real / nb_sample + R_precision = R_precision / nb_sample + + matching_score_real = matching_score_real / nb_sample + matching_score_pred = matching_score_pred / nb_sample + + + fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov) + + msg = f"--> \t Eva. Iter {nb_iter} :, FID. {fid:.4f}, Diversity Real. {diversity_real:.4f}, Diversity. {diversity:.4f}, R_precision_real. {R_precision_real}, R_precision. {R_precision}, matching_score_real. {matching_score_real}, matching_score_pred. {matching_score_pred}" + logger.info(msg) + + + if draw: + writer.add_scalar('./Test/FID', fid, nb_iter) + writer.add_scalar('./Test/Diversity', diversity, nb_iter) + writer.add_scalar('./Test/top1', R_precision[0], nb_iter) + writer.add_scalar('./Test/top2', R_precision[1], nb_iter) + writer.add_scalar('./Test/top3', R_precision[2], nb_iter) + writer.add_scalar('./Test/matching_score', matching_score_pred, nb_iter) + + + if nb_iter % 10000 == 0 : + for ii in range(4): + tensorborad_add_video_xyz(writer, draw_org[ii], nb_iter, tag='./Vis/org_eval'+str(ii), nb_vis=1, title_batch=[draw_text[ii]], outname=[os.path.join(out_dir, 'gt'+str(ii)+'.gif')] if savegif else None) + + if nb_iter % 10000 == 0 : + for ii in range(4): + tensorborad_add_video_xyz(writer, draw_pred[ii], nb_iter, tag='./Vis/pred_eval'+str(ii), nb_vis=1, title_batch=[draw_text_pred[ii]], outname=[os.path.join(out_dir, 'pred'+str(ii)+'.gif')] if savegif else None) + + + if fid < best_fid : + msg = f"--> --> \t FID Improved from {best_fid:.5f} to {fid:.5f} !!!" + logger.info(msg) + best_fid, best_iter = fid, nb_iter + if save: + torch.save({'trans' : trans.state_dict()}, os.path.join(out_dir, 'net_best_fid.pth')) + + if matching_score_pred < best_matching : + msg = f"--> --> \t matching_score Improved from {best_matching:.5f} to {matching_score_pred:.5f} !!!" + logger.info(msg) + best_matching = matching_score_pred + + if abs(diversity_real - diversity) < abs(diversity_real - best_div) : + msg = f"--> --> \t Diversity Improved from {best_div:.5f} to {diversity:.5f} !!!" + logger.info(msg) + best_div = diversity + + if R_precision[0] > best_top1 : + msg = f"--> --> \t Top1 Improved from {best_top1:.4f} to {R_precision[0]:.4f} !!!" + logger.info(msg) + best_top1 = R_precision[0] + + if R_precision[1] > best_top2 : + msg = f"--> --> \t Top2 Improved from {best_top2:.4f} to {R_precision[1]:.4f} !!!" + logger.info(msg) + best_top2 = R_precision[1] + + if R_precision[2] > best_top3 : + msg = f"--> --> \t Top3 Improved from {best_top3:.4f} to {R_precision[2]:.4f} !!!" + logger.info(msg) + best_top3 = R_precision[2] + + if save: + torch.save({'trans' : trans.state_dict()}, os.path.join(out_dir, 'net_last.pth')) + + trans.train() + return best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger + + +@torch.no_grad() +def evaluation_transformer_test(out_dir, val_loader, net, trans, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, best_multi, clip_model, eval_wrapper, draw = True, save = True, savegif=False, savenpy=False) : + + trans.eval() + nb_sample = 0 + + draw_org = [] + draw_pred = [] + draw_text = [] + draw_text_pred = [] + draw_name = [] + + motion_annotation_list = [] + motion_pred_list = [] + motion_multimodality = [] + R_precision_real = 0 + R_precision = 0 + matching_score_real = 0 + matching_score_pred = 0 + + nb_sample = 0 + + for batch in val_loader: + + word_embeddings, pos_one_hots, clip_text, sent_len, pose, m_length, token, name = batch + bs, seq = pose.shape[:2] + num_joints = 21 if pose.shape[-1] == 251 else 22 + + text = clip.tokenize(clip_text, truncate=True).cuda() + + feat_clip_text = clip_model.encode_text(text).float() + motion_multimodality_batch = [] + for i in range(30): + pred_pose_eval = torch.zeros((bs, seq, pose.shape[-1])).cuda() + pred_len = torch.ones(bs).long() + + for k in range(bs): + try: + index_motion = trans.sample(feat_clip_text[k:k+1], True) + except: + index_motion = torch.ones(1,1).cuda().long() + + pred_pose = net.forward_decoder(index_motion) + cur_len = pred_pose.shape[1] + + pred_len[k] = min(cur_len, seq) + pred_pose_eval[k:k+1, :cur_len] = pred_pose[:, :seq] + + if i == 0 and (draw or savenpy): + pred_denorm = val_loader.dataset.inv_transform(pred_pose.detach().cpu().numpy()) + pred_xyz = recover_from_ric(torch.from_numpy(pred_denorm).float().cuda(), num_joints) + + if savenpy: + np.save(os.path.join(out_dir, name[k]+'_pred.npy'), pred_xyz.detach().cpu().numpy()) + + if draw: + if i == 0: + draw_pred.append(pred_xyz) + draw_text_pred.append(clip_text[k]) + draw_name.append(name[k]) + + et_pred, em_pred = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, pred_pose_eval, pred_len) + + motion_multimodality_batch.append(em_pred.reshape(bs, 1, -1)) + + if i == 0: + pose = pose.cuda().float() + + et, em = eval_wrapper.get_co_embeddings(word_embeddings, pos_one_hots, sent_len, pose, m_length) + motion_annotation_list.append(em) + motion_pred_list.append(em_pred) + + if draw or savenpy: + pose = val_loader.dataset.inv_transform(pose.detach().cpu().numpy()) + pose_xyz = recover_from_ric(torch.from_numpy(pose).float().cuda(), num_joints) + + if savenpy: + for j in range(bs): + np.save(os.path.join(out_dir, name[j]+'_gt.npy'), pose_xyz[j][:m_length[j]].unsqueeze(0).cpu().numpy()) + + if draw: + for j in range(bs): + draw_org.append(pose_xyz[j][:m_length[j]].unsqueeze(0)) + draw_text.append(clip_text[j]) + + temp_R, temp_match = calculate_R_precision(et.cpu().numpy(), em.cpu().numpy(), top_k=3, sum_all=True) + R_precision_real += temp_R + matching_score_real += temp_match + temp_R, temp_match = calculate_R_precision(et_pred.cpu().numpy(), em_pred.cpu().numpy(), top_k=3, sum_all=True) + R_precision += temp_R + matching_score_pred += temp_match + + nb_sample += bs + + motion_multimodality.append(torch.cat(motion_multimodality_batch, dim=1)) + + motion_annotation_np = torch.cat(motion_annotation_list, dim=0).cpu().numpy() + motion_pred_np = torch.cat(motion_pred_list, dim=0).cpu().numpy() + gt_mu, gt_cov = calculate_activation_statistics(motion_annotation_np) + mu, cov= calculate_activation_statistics(motion_pred_np) + + diversity_real = calculate_diversity(motion_annotation_np, 300 if nb_sample > 300 else 100) + diversity = calculate_diversity(motion_pred_np, 300 if nb_sample > 300 else 100) + + R_precision_real = R_precision_real / nb_sample + R_precision = R_precision / nb_sample + + matching_score_real = matching_score_real / nb_sample + matching_score_pred = matching_score_pred / nb_sample + + multimodality = 0 + motion_multimodality = torch.cat(motion_multimodality, dim=0).cpu().numpy() + multimodality = calculate_multimodality(motion_multimodality, 10) + + fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov) + + msg = f"--> \t Eva. Iter {nb_iter} :, FID. {fid:.4f}, Diversity Real. {diversity_real:.4f}, Diversity. {diversity:.4f}, R_precision_real. {R_precision_real}, R_precision. {R_precision}, matching_score_real. {matching_score_real}, matching_score_pred. {matching_score_pred}, multimodality. {multimodality:.4f}" + logger.info(msg) + + + if draw: + for ii in range(len(draw_org)): + tensorborad_add_video_xyz(writer, draw_org[ii], nb_iter, tag='./Vis/'+draw_name[ii]+'_org', nb_vis=1, title_batch=[draw_text[ii]], outname=[os.path.join(out_dir, draw_name[ii]+'_skel_gt.gif')] if savegif else None) + + tensorborad_add_video_xyz(writer, draw_pred[ii], nb_iter, tag='./Vis/'+draw_name[ii]+'_pred', nb_vis=1, title_batch=[draw_text_pred[ii]], outname=[os.path.join(out_dir, draw_name[ii]+'_skel_pred.gif')] if savegif else None) + + trans.train() + return fid, best_iter, diversity, R_precision[0], R_precision[1], R_precision[2], matching_score_pred, multimodality, writer, logger + +# (X - X_train)*(X - X_train) = -2X*X_train + X*X + X_train*X_train +def euclidean_distance_matrix(matrix1, matrix2): + """ + Params: + -- matrix1: N1 x D + -- matrix2: N2 x D + Returns: + -- dist: N1 x N2 + dist[i, j] == distance(matrix1[i], matrix2[j]) + """ + assert matrix1.shape[1] == matrix2.shape[1] + d1 = -2 * np.dot(matrix1, matrix2.T) # shape (num_test, num_train) + d2 = np.sum(np.square(matrix1), axis=1, keepdims=True) # shape (num_test, 1) + d3 = np.sum(np.square(matrix2), axis=1) # shape (num_train, ) + dists = np.sqrt(d1 + d2 + d3) # broadcasting + return dists + + + +def calculate_top_k(mat, top_k): + size = mat.shape[0] + gt_mat = np.expand_dims(np.arange(size), 1).repeat(size, 1) + bool_mat = (mat == gt_mat) + correct_vec = False + top_k_list = [] + for i in range(top_k): +# print(correct_vec, bool_mat[:, i]) + correct_vec = (correct_vec | bool_mat[:, i]) + # print(correct_vec) + top_k_list.append(correct_vec[:, None]) + top_k_mat = np.concatenate(top_k_list, axis=1) + return top_k_mat + + +def calculate_R_precision(embedding1, embedding2, top_k, sum_all=False): + dist_mat = euclidean_distance_matrix(embedding1, embedding2) + matching_score = dist_mat.trace() + argmax = np.argsort(dist_mat, axis=1) + top_k_mat = calculate_top_k(argmax, top_k) + if sum_all: + return top_k_mat.sum(axis=0), matching_score + else: + return top_k_mat, matching_score + +def calculate_multimodality(activation, multimodality_times): + assert len(activation.shape) == 3 + assert activation.shape[1] > multimodality_times + num_per_sent = activation.shape[1] + + first_dices = np.random.choice(num_per_sent, multimodality_times, replace=False) + second_dices = np.random.choice(num_per_sent, multimodality_times, replace=False) + dist = linalg.norm(activation[:, first_dices] - activation[:, second_dices], axis=2) + return dist.mean() + + +def calculate_diversity(activation, diversity_times): + assert len(activation.shape) == 2 + assert activation.shape[0] > diversity_times + num_samples = activation.shape[0] + + first_indices = np.random.choice(num_samples, diversity_times, replace=False) + second_indices = np.random.choice(num_samples, diversity_times, replace=False) + dist = linalg.norm(activation[first_indices] - activation[second_indices], axis=1) + return dist.mean() + + + +def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): + + mu1 = np.atleast_1d(mu1) + mu2 = np.atleast_1d(mu2) + + sigma1 = np.atleast_2d(sigma1) + sigma2 = np.atleast_2d(sigma2) + + assert mu1.shape == mu2.shape, \ + 'Training and test mean vectors have different lengths' + assert sigma1.shape == sigma2.shape, \ + 'Training and test covariances have different dimensions' + + diff = mu1 - mu2 + + # Product might be almost singular + covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) + if not np.isfinite(covmean).all(): + msg = ('fid calculation produces singular product; ' + 'adding %s to diagonal of cov estimates') % eps + print(msg) + offset = np.eye(sigma1.shape[0]) * eps + covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) + + # Numerical error might give slight imaginary component + if np.iscomplexobj(covmean): + if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): + m = np.max(np.abs(covmean.imag)) + raise ValueError('Imaginary component {}'.format(m)) + covmean = covmean.real + + tr_covmean = np.trace(covmean) + + return (diff.dot(diff) + np.trace(sigma1) + + np.trace(sigma2) - 2 * tr_covmean) + + + +def calculate_activation_statistics(activations): + + mu = np.mean(activations, axis=0) + cov = np.cov(activations, rowvar=False) + return mu, cov + + +def calculate_frechet_feature_distance(feature_list1, feature_list2): + feature_list1 = np.stack(feature_list1) + feature_list2 = np.stack(feature_list2) + + # normalize the scale + mean = np.mean(feature_list1, axis=0) + std = np.std(feature_list1, axis=0) + 1e-10 + feature_list1 = (feature_list1 - mean) / std + feature_list2 = (feature_list2 - mean) / std + + dist = calculate_frechet_distance( + mu1=np.mean(feature_list1, axis=0), + sigma1=np.cov(feature_list1, rowvar=False), + mu2=np.mean(feature_list2, axis=0), + sigma2=np.cov(feature_list2, rowvar=False), + ) + return dist \ No newline at end of file diff --git a/VQ-Trans/utils/losses.py b/VQ-Trans/utils/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..1998161032731fc2c3edae701700679c00fd00d0 --- /dev/null +++ b/VQ-Trans/utils/losses.py @@ -0,0 +1,30 @@ +import torch +import torch.nn as nn + +class ReConsLoss(nn.Module): + def __init__(self, recons_loss, nb_joints): + super(ReConsLoss, self).__init__() + + if recons_loss == 'l1': + self.Loss = torch.nn.L1Loss() + elif recons_loss == 'l2' : + self.Loss = torch.nn.MSELoss() + elif recons_loss == 'l1_smooth' : + self.Loss = torch.nn.SmoothL1Loss() + + # 4 global motion associated to root + # 12 local motion (3 local xyz, 3 vel xyz, 6 rot6d) + # 3 global vel xyz + # 4 foot contact + self.nb_joints = nb_joints + self.motion_dim = (nb_joints - 1) * 12 + 4 + 3 + 4 + + def forward(self, motion_pred, motion_gt) : + loss = self.Loss(motion_pred[..., : self.motion_dim], motion_gt[..., :self.motion_dim]) + return loss + + def forward_vel(self, motion_pred, motion_gt) : + loss = self.Loss(motion_pred[..., 4 : (self.nb_joints - 1) * 3 + 4], motion_gt[..., 4 : (self.nb_joints - 1) * 3 + 4]) + return loss + + \ No newline at end of file diff --git a/VQ-Trans/utils/motion_process.py b/VQ-Trans/utils/motion_process.py new file mode 100644 index 0000000000000000000000000000000000000000..7819c8b3cc61b6e48c65d1a456342119060696ea --- /dev/null +++ b/VQ-Trans/utils/motion_process.py @@ -0,0 +1,59 @@ +import torch +from utils.quaternion import quaternion_to_cont6d, qrot, qinv + +def recover_root_rot_pos(data): + rot_vel = data[..., 0] + r_rot_ang = torch.zeros_like(rot_vel).to(data.device) + '''Get Y-axis rotation from rotation velocity''' + r_rot_ang[..., 1:] = rot_vel[..., :-1] + r_rot_ang = torch.cumsum(r_rot_ang, dim=-1) + + r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device) + r_rot_quat[..., 0] = torch.cos(r_rot_ang) + r_rot_quat[..., 2] = torch.sin(r_rot_ang) + + r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device) + r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3] + '''Add Y-axis rotation to root position''' + r_pos = qrot(qinv(r_rot_quat), r_pos) + + r_pos = torch.cumsum(r_pos, dim=-2) + + r_pos[..., 1] = data[..., 3] + return r_rot_quat, r_pos + + +def recover_from_rot(data, joints_num, skeleton): + r_rot_quat, r_pos = recover_root_rot_pos(data) + + r_rot_cont6d = quaternion_to_cont6d(r_rot_quat) + + start_indx = 1 + 2 + 1 + (joints_num - 1) * 3 + end_indx = start_indx + (joints_num - 1) * 6 + cont6d_params = data[..., start_indx:end_indx] + # print(r_rot_cont6d.shape, cont6d_params.shape, r_pos.shape) + cont6d_params = torch.cat([r_rot_cont6d, cont6d_params], dim=-1) + cont6d_params = cont6d_params.view(-1, joints_num, 6) + + positions = skeleton.forward_kinematics_cont6d(cont6d_params, r_pos) + + return positions + + +def recover_from_ric(data, joints_num): + r_rot_quat, r_pos = recover_root_rot_pos(data) + positions = data[..., 4:(joints_num - 1) * 3 + 4] + positions = positions.view(positions.shape[:-1] + (-1, 3)) + + '''Add Y-axis rotation to local joints''' + positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions) + + '''Add root XZ to joints''' + positions[..., 0] += r_pos[..., 0:1] + positions[..., 2] += r_pos[..., 2:3] + + '''Concate root and joints''' + positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2) + + return positions + \ No newline at end of file diff --git a/VQ-Trans/utils/paramUtil.py b/VQ-Trans/utils/paramUtil.py new file mode 100644 index 0000000000000000000000000000000000000000..a9f1708b85ca80a9051cb3675cec9b999a0d0e2b --- /dev/null +++ b/VQ-Trans/utils/paramUtil.py @@ -0,0 +1,63 @@ +import numpy as np + +# Define a kinematic tree for the skeletal struture +kit_kinematic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]] + +kit_raw_offsets = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 1, 0], + [0, 1, 0], + [0, 1, 0], + [1, 0, 0], + [0, -1, 0], + [0, -1, 0], + [-1, 0, 0], + [0, -1, 0], + [0, -1, 0], + [1, 0, 0], + [0, -1, 0], + [0, -1, 0], + [0, 0, 1], + [0, 0, 1], + [-1, 0, 0], + [0, -1, 0], + [0, -1, 0], + [0, 0, 1], + [0, 0, 1] + ] +) + +t2m_raw_offsets = np.array([[0,0,0], + [1,0,0], + [-1,0,0], + [0,1,0], + [0,-1,0], + [0,-1,0], + [0,1,0], + [0,-1,0], + [0,-1,0], + [0,1,0], + [0,0,1], + [0,0,1], + [0,1,0], + [1,0,0], + [-1,0,0], + [0,0,1], + [0,-1,0], + [0,-1,0], + [0,-1,0], + [0,-1,0], + [0,-1,0], + [0,-1,0]]) + +t2m_kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]] +t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]] +t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]] + + +kit_tgt_skel_id = '03950' + +t2m_tgt_skel_id = '000021' + diff --git a/VQ-Trans/utils/quaternion.py b/VQ-Trans/utils/quaternion.py new file mode 100644 index 0000000000000000000000000000000000000000..e2daa00aef1df60e43775864d1dd3d551f89ded8 --- /dev/null +++ b/VQ-Trans/utils/quaternion.py @@ -0,0 +1,423 @@ +# Copyright (c) 2018-present, Facebook, Inc. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# + +import torch +import numpy as np + +_EPS4 = np.finfo(float).eps * 4.0 + +_FLOAT_EPS = np.finfo(np.float).eps + +# PyTorch-backed implementations +def qinv(q): + assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' + mask = torch.ones_like(q) + mask[..., 1:] = -mask[..., 1:] + return q * mask + + +def qinv_np(q): + assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' + return qinv(torch.from_numpy(q).float()).numpy() + + +def qnormalize(q): + assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' + return q / torch.norm(q, dim=-1, keepdim=True) + + +def qmul(q, r): + """ + Multiply quaternion(s) q with quaternion(s) r. + Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions. + Returns q*r as a tensor of shape (*, 4). + """ + assert q.shape[-1] == 4 + assert r.shape[-1] == 4 + + original_shape = q.shape + + # Compute outer product + terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4)) + + w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3] + x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2] + y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1] + z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0] + return torch.stack((w, x, y, z), dim=1).view(original_shape) + + +def qrot(q, v): + """ + Rotate vector(s) v about the rotation described by quaternion(s) q. + Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v, + where * denotes any number of dimensions. + Returns a tensor of shape (*, 3). + """ + assert q.shape[-1] == 4 + assert v.shape[-1] == 3 + assert q.shape[:-1] == v.shape[:-1] + + original_shape = list(v.shape) + # print(q.shape) + q = q.contiguous().view(-1, 4) + v = v.contiguous().view(-1, 3) + + qvec = q[:, 1:] + uv = torch.cross(qvec, v, dim=1) + uuv = torch.cross(qvec, uv, dim=1) + return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape) + + +def qeuler(q, order, epsilon=0, deg=True): + """ + Convert quaternion(s) q to Euler angles. + Expects a tensor of shape (*, 4), where * denotes any number of dimensions. + Returns a tensor of shape (*, 3). + """ + assert q.shape[-1] == 4 + + original_shape = list(q.shape) + original_shape[-1] = 3 + q = q.view(-1, 4) + + q0 = q[:, 0] + q1 = q[:, 1] + q2 = q[:, 2] + q3 = q[:, 3] + + if order == 'xyz': + x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) + y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon)) + z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) + elif order == 'yzx': + x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) + y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3)) + z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon)) + elif order == 'zxy': + x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon)) + y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) + z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3)) + elif order == 'xzy': + x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) + y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3)) + z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon)) + elif order == 'yxz': + x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon)) + y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2)) + z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3)) + elif order == 'zyx': + x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) + y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon)) + z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) + else: + raise + + if deg: + return torch.stack((x, y, z), dim=1).view(original_shape) * 180 / np.pi + else: + return torch.stack((x, y, z), dim=1).view(original_shape) + + +# Numpy-backed implementations + +def qmul_np(q, r): + q = torch.from_numpy(q).contiguous().float() + r = torch.from_numpy(r).contiguous().float() + return qmul(q, r).numpy() + + +def qrot_np(q, v): + q = torch.from_numpy(q).contiguous().float() + v = torch.from_numpy(v).contiguous().float() + return qrot(q, v).numpy() + + +def qeuler_np(q, order, epsilon=0, use_gpu=False): + if use_gpu: + q = torch.from_numpy(q).cuda().float() + return qeuler(q, order, epsilon).cpu().numpy() + else: + q = torch.from_numpy(q).contiguous().float() + return qeuler(q, order, epsilon).numpy() + + +def qfix(q): + """ + Enforce quaternion continuity across the time dimension by selecting + the representation (q or -q) with minimal distance (or, equivalently, maximal dot product) + between two consecutive frames. + + Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints. + Returns a tensor of the same shape. + """ + assert len(q.shape) == 3 + assert q.shape[-1] == 4 + + result = q.copy() + dot_products = np.sum(q[1:] * q[:-1], axis=2) + mask = dot_products < 0 + mask = (np.cumsum(mask, axis=0) % 2).astype(bool) + result[1:][mask] *= -1 + return result + + +def euler2quat(e, order, deg=True): + """ + Convert Euler angles to quaternions. + """ + assert e.shape[-1] == 3 + + original_shape = list(e.shape) + original_shape[-1] = 4 + + e = e.view(-1, 3) + + ## if euler angles in degrees + if deg: + e = e * np.pi / 180. + + x = e[:, 0] + y = e[:, 1] + z = e[:, 2] + + rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1) + ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1) + rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1) + + result = None + for coord in order: + if coord == 'x': + r = rx + elif coord == 'y': + r = ry + elif coord == 'z': + r = rz + else: + raise + if result is None: + result = r + else: + result = qmul(result, r) + + # Reverse antipodal representation to have a non-negative "w" + if order in ['xyz', 'yzx', 'zxy']: + result *= -1 + + return result.view(original_shape) + + +def expmap_to_quaternion(e): + """ + Convert axis-angle rotations (aka exponential maps) to quaternions. + Stable formula from "Practical Parameterization of Rotations Using the Exponential Map". + Expects a tensor of shape (*, 3), where * denotes any number of dimensions. + Returns a tensor of shape (*, 4). + """ + assert e.shape[-1] == 3 + + original_shape = list(e.shape) + original_shape[-1] = 4 + e = e.reshape(-1, 3) + + theta = np.linalg.norm(e, axis=1).reshape(-1, 1) + w = np.cos(0.5 * theta).reshape(-1, 1) + xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e + return np.concatenate((w, xyz), axis=1).reshape(original_shape) + + +def euler_to_quaternion(e, order): + """ + Convert Euler angles to quaternions. + """ + assert e.shape[-1] == 3 + + original_shape = list(e.shape) + original_shape[-1] = 4 + + e = e.reshape(-1, 3) + + x = e[:, 0] + y = e[:, 1] + z = e[:, 2] + + rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1) + ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1) + rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1) + + result = None + for coord in order: + if coord == 'x': + r = rx + elif coord == 'y': + r = ry + elif coord == 'z': + r = rz + else: + raise + if result is None: + result = r + else: + result = qmul_np(result, r) + + # Reverse antipodal representation to have a non-negative "w" + if order in ['xyz', 'yzx', 'zxy']: + result *= -1 + + return result.reshape(original_shape) + + +def quaternion_to_matrix(quaternions): + """ + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def quaternion_to_matrix_np(quaternions): + q = torch.from_numpy(quaternions).contiguous().float() + return quaternion_to_matrix(q).numpy() + + +def quaternion_to_cont6d_np(quaternions): + rotation_mat = quaternion_to_matrix_np(quaternions) + cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1) + return cont_6d + + +def quaternion_to_cont6d(quaternions): + rotation_mat = quaternion_to_matrix(quaternions) + cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1) + return cont_6d + + +def cont6d_to_matrix(cont6d): + assert cont6d.shape[-1] == 6, "The last dimension must be 6" + x_raw = cont6d[..., 0:3] + y_raw = cont6d[..., 3:6] + + x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True) + z = torch.cross(x, y_raw, dim=-1) + z = z / torch.norm(z, dim=-1, keepdim=True) + + y = torch.cross(z, x, dim=-1) + + x = x[..., None] + y = y[..., None] + z = z[..., None] + + mat = torch.cat([x, y, z], dim=-1) + return mat + + +def cont6d_to_matrix_np(cont6d): + q = torch.from_numpy(cont6d).contiguous().float() + return cont6d_to_matrix(q).numpy() + + +def qpow(q0, t, dtype=torch.float): + ''' q0 : tensor of quaternions + t: tensor of powers + ''' + q0 = qnormalize(q0) + theta0 = torch.acos(q0[..., 0]) + + ## if theta0 is close to zero, add epsilon to avoid NaNs + mask = (theta0 <= 10e-10) * (theta0 >= -10e-10) + theta0 = (1 - mask) * theta0 + mask * 10e-10 + v0 = q0[..., 1:] / torch.sin(theta0).view(-1, 1) + + if isinstance(t, torch.Tensor): + q = torch.zeros(t.shape + q0.shape) + theta = t.view(-1, 1) * theta0.view(1, -1) + else: ## if t is a number + q = torch.zeros(q0.shape) + theta = t * theta0 + + q[..., 0] = torch.cos(theta) + q[..., 1:] = v0 * torch.sin(theta).unsqueeze(-1) + + return q.to(dtype) + + +def qslerp(q0, q1, t): + ''' + q0: starting quaternion + q1: ending quaternion + t: array of points along the way + + Returns: + Tensor of Slerps: t.shape + q0.shape + ''' + + q0 = qnormalize(q0) + q1 = qnormalize(q1) + q_ = qpow(qmul(q1, qinv(q0)), t) + + return qmul(q_, + q0.contiguous().view(torch.Size([1] * len(t.shape)) + q0.shape).expand(t.shape + q0.shape).contiguous()) + + +def qbetween(v0, v1): + ''' + find the quaternion used to rotate v0 to v1 + ''' + assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)' + assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)' + + v = torch.cross(v0, v1) + w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1, + keepdim=True) + return qnormalize(torch.cat([w, v], dim=-1)) + + +def qbetween_np(v0, v1): + ''' + find the quaternion used to rotate v0 to v1 + ''' + assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)' + assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)' + + v0 = torch.from_numpy(v0).float() + v1 = torch.from_numpy(v1).float() + return qbetween(v0, v1).numpy() + + +def lerp(p0, p1, t): + if not isinstance(t, torch.Tensor): + t = torch.Tensor([t]) + + new_shape = t.shape + p0.shape + new_view_t = t.shape + torch.Size([1] * len(p0.shape)) + new_view_p = torch.Size([1] * len(t.shape)) + p0.shape + p0 = p0.view(new_view_p).expand(new_shape) + p1 = p1.view(new_view_p).expand(new_shape) + t = t.view(new_view_t).expand(new_shape) + + return p0 + t * (p1 - p0) diff --git a/VQ-Trans/utils/rotation_conversions.py b/VQ-Trans/utils/rotation_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..1006e8a3117b231a7a456d5b826e76347fe0bfd4 --- /dev/null +++ b/VQ-Trans/utils/rotation_conversions.py @@ -0,0 +1,532 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# Check PYTORCH3D_LICENCE before use + +import functools +from typing import Optional + +import torch +import torch.nn.functional as F + + +""" +The transformation matrices returned from the functions in this file assume +the points on which the transformation will be applied are column vectors. +i.e. the R matrix is structured as + R = [ + [Rxx, Rxy, Rxz], + [Ryx, Ryy, Ryz], + [Rzx, Rzy, Rzz], + ] # (3, 3) +This matrix can be applied to column vectors by post multiplication +by the points e.g. + points = [[0], [1], [2]] # (3 x 1) xyz coordinates of a point + transformed_points = R * points +To apply the same matrix to points which are row vectors, the R matrix +can be transposed and pre multiplied by the points: +e.g. + points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point + transformed_points = points * R.transpose(1, 0) +""" + + +def quaternion_to_matrix(quaternions): + """ + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def _copysign(a, b): + """ + Return a tensor where each element has the absolute value taken from the, + corresponding element of a, with sign taken from the corresponding + element of b. This is like the standard copysign floating-point operation, + but is not careful about negative 0 and NaN. + Args: + a: source tensor. + b: tensor whose signs will be used, of the same shape as a. + Returns: + Tensor of the same shape as a with the signs of b. + """ + signs_differ = (a < 0) != (b < 0) + return torch.where(signs_differ, -a, a) + + +def _sqrt_positive_part(x): + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + ret[positive_mask] = torch.sqrt(x[positive_mask]) + return ret + + +def matrix_to_quaternion(matrix): + """ + Convert rotations given as rotation matrices to quaternions. + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") + m00 = matrix[..., 0, 0] + m11 = matrix[..., 1, 1] + m22 = matrix[..., 2, 2] + o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22) + x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22) + y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22) + z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22) + o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2]) + o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0]) + o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1]) + return torch.stack((o0, o1, o2, o3), -1) + + +def _axis_angle_rotation(axis: str, angle): + """ + Return the rotation matrices for one of the rotations about an axis + of which Euler angles describe, for each value of the angle given. + Args: + axis: Axis label "X" or "Y or "Z". + angle: any shape tensor of Euler angles in radians + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + + cos = torch.cos(angle) + sin = torch.sin(angle) + one = torch.ones_like(angle) + zero = torch.zeros_like(angle) + + if axis == "X": + R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) + if axis == "Y": + R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) + if axis == "Z": + R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) + + return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) + + +def euler_angles_to_matrix(euler_angles, convention: str): + """ + Convert rotations given as Euler angles in radians to rotation matrices. + Args: + euler_angles: Euler angles in radians as tensor of shape (..., 3). + convention: Convention string of three uppercase letters from + {"X", "Y", and "Z"}. + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: + raise ValueError("Invalid input euler angles.") + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1)) + return functools.reduce(torch.matmul, matrices) + + +def _angle_from_tan( + axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool +): + """ + Extract the first or third Euler angle from the two members of + the matrix which are positive constant times its sine and cosine. + Args: + axis: Axis label "X" or "Y or "Z" for the angle we are finding. + other_axis: Axis label "X" or "Y or "Z" for the middle axis in the + convention. + data: Rotation matrices as tensor of shape (..., 3, 3). + horizontal: Whether we are looking for the angle for the third axis, + which means the relevant entries are in the same row of the + rotation matrix. If not, they are in the same column. + tait_bryan: Whether the first and third axes in the convention differ. + Returns: + Euler Angles in radians for each matrix in data as a tensor + of shape (...). + """ + + i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] + if horizontal: + i2, i1 = i1, i2 + even = (axis + other_axis) in ["XY", "YZ", "ZX"] + if horizontal == even: + return torch.atan2(data[..., i1], data[..., i2]) + if tait_bryan: + return torch.atan2(-data[..., i2], data[..., i1]) + return torch.atan2(data[..., i2], -data[..., i1]) + + +def _index_from_letter(letter: str): + if letter == "X": + return 0 + if letter == "Y": + return 1 + if letter == "Z": + return 2 + + +def matrix_to_euler_angles(matrix, convention: str): + """ + Convert rotations given as rotation matrices to Euler angles in radians. + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + convention: Convention string of three uppercase letters. + Returns: + Euler angles in radians as tensor of shape (..., 3). + """ + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") + i0 = _index_from_letter(convention[0]) + i2 = _index_from_letter(convention[2]) + tait_bryan = i0 != i2 + if tait_bryan: + central_angle = torch.asin( + matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0) + ) + else: + central_angle = torch.acos(matrix[..., i0, i0]) + + o = ( + _angle_from_tan( + convention[0], convention[1], matrix[..., i2], False, tait_bryan + ), + central_angle, + _angle_from_tan( + convention[2], convention[1], matrix[..., i0, :], True, tait_bryan + ), + ) + return torch.stack(o, -1) + + +def random_quaternions( + n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False +): + """ + Generate random quaternions representing rotations, + i.e. versors with nonnegative real part. + Args: + n: Number of quaternions in a batch to return. + dtype: Type to return. + device: Desired device of returned tensor. Default: + uses the current device for the default tensor type. + requires_grad: Whether the resulting tensor should have the gradient + flag set. + Returns: + Quaternions as tensor of shape (N, 4). + """ + o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad) + s = (o * o).sum(1) + o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] + return o + + +def random_rotations( + n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False +): + """ + Generate random rotations as 3x3 rotation matrices. + Args: + n: Number of rotation matrices in a batch to return. + dtype: Type to return. + device: Device of returned tensor. Default: if None, + uses the current device for the default tensor type. + requires_grad: Whether the resulting tensor should have the gradient + flag set. + Returns: + Rotation matrices as tensor of shape (n, 3, 3). + """ + quaternions = random_quaternions( + n, dtype=dtype, device=device, requires_grad=requires_grad + ) + return quaternion_to_matrix(quaternions) + + +def random_rotation( + dtype: Optional[torch.dtype] = None, device=None, requires_grad=False +): + """ + Generate a single random 3x3 rotation matrix. + Args: + dtype: Type to return + device: Device of returned tensor. Default: if None, + uses the current device for the default tensor type + requires_grad: Whether the resulting tensor should have the gradient + flag set + Returns: + Rotation matrix as tensor of shape (3, 3). + """ + return random_rotations(1, dtype, device, requires_grad)[0] + + +def standardize_quaternion(quaternions): + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + Args: + quaternions: Quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions) + + +def quaternion_raw_multiply(a, b): + """ + Multiply two quaternions. + Usual torch rules for broadcasting apply. + Args: + a: Quaternions as tensor of shape (..., 4), real part first. + b: Quaternions as tensor of shape (..., 4), real part first. + Returns: + The product of a and b, a tensor of quaternions shape (..., 4). + """ + aw, ax, ay, az = torch.unbind(a, -1) + bw, bx, by, bz = torch.unbind(b, -1) + ow = aw * bw - ax * bx - ay * by - az * bz + ox = aw * bx + ax * bw + ay * bz - az * by + oy = aw * by - ax * bz + ay * bw + az * bx + oz = aw * bz + ax * by - ay * bx + az * bw + return torch.stack((ow, ox, oy, oz), -1) + + +def quaternion_multiply(a, b): + """ + Multiply two quaternions representing rotations, returning the quaternion + representing their composition, i.e. the versor with nonnegative real part. + Usual torch rules for broadcasting apply. + Args: + a: Quaternions as tensor of shape (..., 4), real part first. + b: Quaternions as tensor of shape (..., 4), real part first. + Returns: + The product of a and b, a tensor of quaternions of shape (..., 4). + """ + ab = quaternion_raw_multiply(a, b) + return standardize_quaternion(ab) + + +def quaternion_invert(quaternion): + """ + Given a quaternion representing rotation, get the quaternion representing + its inverse. + Args: + quaternion: Quaternions as tensor of shape (..., 4), with real part + first, which must be versors (unit quaternions). + Returns: + The inverse, a tensor of quaternions of shape (..., 4). + """ + + return quaternion * quaternion.new_tensor([1, -1, -1, -1]) + + +def quaternion_apply(quaternion, point): + """ + Apply the rotation given by a quaternion to a 3D point. + Usual torch rules for broadcasting apply. + Args: + quaternion: Tensor of quaternions, real part first, of shape (..., 4). + point: Tensor of 3D points of shape (..., 3). + Returns: + Tensor of rotated points of shape (..., 3). + """ + if point.size(-1) != 3: + raise ValueError(f"Points are not in 3D, f{point.shape}.") + real_parts = point.new_zeros(point.shape[:-1] + (1,)) + point_as_quaternion = torch.cat((real_parts, point), -1) + out = quaternion_raw_multiply( + quaternion_raw_multiply(quaternion, point_as_quaternion), + quaternion_invert(quaternion), + ) + return out[..., 1:] + + +def axis_angle_to_matrix(axis_angle): + """ + Convert rotations given as axis/angle to rotation matrices. + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) + + +def matrix_to_axis_angle(matrix): + """ + Convert rotations given as rotation matrices to axis/angle. + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + return quaternion_to_axis_angle(matrix_to_quaternion(matrix)) + + +def axis_angle_to_quaternion(axis_angle): + """ + Convert rotations given as axis/angle to quaternions. + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) + half_angles = 0.5 * angles + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + quaternions = torch.cat( + [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 + ) + return quaternions + + +def quaternion_to_axis_angle(quaternions): + """ + Convert rotations given as quaternions to axis/angle. + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) + half_angles = torch.atan2(norms, quaternions[..., :1]) + angles = 2 * half_angles + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + return quaternions[..., 1:] / sin_half_angles_over_angles + + +def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: + """ + Converts 6D rotation representation by Zhou et al. [1] to rotation matrix + using Gram--Schmidt orthogonalisation per Section B of [1]. + Args: + d6: 6D rotation representation, of size (*, 6) + Returns: + batch of rotation matrices of size (*, 3, 3) + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + + a1, a2 = d6[..., :3], d6[..., 3:] + b1 = F.normalize(a1, dim=-1) + b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 + b2 = F.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack((b1, b2, b3), dim=-2) + + +def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor: + """ + Converts rotation matrices to 6D rotation representation by Zhou et al. [1] + by dropping the last row. Note that 6D representation is not unique. + Args: + matrix: batch of rotation matrices of size (*, 3, 3) + Returns: + 6D rotation representation, of size (*, 6) + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6) + +def canonicalize_smplh(poses, trans = None): + bs, nframes, njoints = poses.shape[:3] + + global_orient = poses[:, :, 0] + + # first global rotations + rot2d = matrix_to_axis_angle(global_orient[:, 0]) + #rot2d[:, :2] = 0 # Remove the rotation along the vertical axis + rot2d = axis_angle_to_matrix(rot2d) + + # Rotate the global rotation to eliminate Z rotations + global_orient = torch.einsum("ikj,imkl->imjl", rot2d, global_orient) + + # Construct canonicalized version of x + xc = torch.cat((global_orient[:, :, None], poses[:, :, 1:]), dim=2) + + if trans is not None: + vel = trans[:, 1:] - trans[:, :-1] + # Turn the translation as well + vel = torch.einsum("ikj,ilk->ilj", rot2d, vel) + trans = torch.cat((torch.zeros(bs, 1, 3, device=vel.device), + torch.cumsum(vel, 1)), 1) + return xc, trans + else: + return xc + + \ No newline at end of file diff --git a/VQ-Trans/utils/skeleton.py b/VQ-Trans/utils/skeleton.py new file mode 100644 index 0000000000000000000000000000000000000000..6de56af0c29ae7cccbd7178f912459413f87c646 --- /dev/null +++ b/VQ-Trans/utils/skeleton.py @@ -0,0 +1,199 @@ +from utils.quaternion import * +import scipy.ndimage.filters as filters + +class Skeleton(object): + def __init__(self, offset, kinematic_tree, device): + self.device = device + self._raw_offset_np = offset.numpy() + self._raw_offset = offset.clone().detach().to(device).float() + self._kinematic_tree = kinematic_tree + self._offset = None + self._parents = [0] * len(self._raw_offset) + self._parents[0] = -1 + for chain in self._kinematic_tree: + for j in range(1, len(chain)): + self._parents[chain[j]] = chain[j-1] + + def njoints(self): + return len(self._raw_offset) + + def offset(self): + return self._offset + + def set_offset(self, offsets): + self._offset = offsets.clone().detach().to(self.device).float() + + def kinematic_tree(self): + return self._kinematic_tree + + def parents(self): + return self._parents + + # joints (batch_size, joints_num, 3) + def get_offsets_joints_batch(self, joints): + assert len(joints.shape) == 3 + _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone() + for i in range(1, self._raw_offset.shape[0]): + _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i] + + self._offset = _offsets.detach() + return _offsets + + # joints (joints_num, 3) + def get_offsets_joints(self, joints): + assert len(joints.shape) == 2 + _offsets = self._raw_offset.clone() + for i in range(1, self._raw_offset.shape[0]): + # print(joints.shape) + _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i] + + self._offset = _offsets.detach() + return _offsets + + # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder + # joints (batch_size, joints_num, 3) + def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False): + assert len(face_joint_idx) == 4 + '''Get Forward Direction''' + l_hip, r_hip, sdr_r, sdr_l = face_joint_idx + across1 = joints[:, r_hip] - joints[:, l_hip] + across2 = joints[:, sdr_r] - joints[:, sdr_l] + across = across1 + across2 + across = across / np.sqrt((across**2).sum(axis=-1))[:, np.newaxis] + # print(across1.shape, across2.shape) + + # forward (batch_size, 3) + forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1) + if smooth_forward: + forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest') + # forward (batch_size, 3) + forward = forward / np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis] + + '''Get Root Rotation''' + target = np.array([[0,0,1]]).repeat(len(forward), axis=0) + root_quat = qbetween_np(forward, target) + + '''Inverse Kinematics''' + # quat_params (batch_size, joints_num, 4) + # print(joints.shape[:-1]) + quat_params = np.zeros(joints.shape[:-1] + (4,)) + # print(quat_params.shape) + root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]]) + quat_params[:, 0] = root_quat + # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]]) + for chain in self._kinematic_tree: + R = root_quat + for j in range(len(chain) - 1): + # (batch, 3) + u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0) + # print(u.shape) + # (batch, 3) + v = joints[:, chain[j+1]] - joints[:, chain[j]] + v = v / np.sqrt((v**2).sum(axis=-1))[:, np.newaxis] + # print(u.shape, v.shape) + rot_u_v = qbetween_np(u, v) + + R_loc = qmul_np(qinv_np(R), rot_u_v) + + quat_params[:,chain[j + 1], :] = R_loc + R = qmul_np(R, R_loc) + + return quat_params + + # Be sure root joint is at the beginning of kinematic chains + def forward_kinematics(self, quat_params, root_pos, skel_joints=None, do_root_R=True): + # quat_params (batch_size, joints_num, 4) + # joints (batch_size, joints_num, 3) + # root_pos (batch_size, 3) + if skel_joints is not None: + offsets = self.get_offsets_joints_batch(skel_joints) + if len(self._offset.shape) == 2: + offsets = self._offset.expand(quat_params.shape[0], -1, -1) + joints = torch.zeros(quat_params.shape[:-1] + (3,)).to(self.device) + joints[:, 0] = root_pos + for chain in self._kinematic_tree: + if do_root_R: + R = quat_params[:, 0] + else: + R = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(len(quat_params), -1).detach().to(self.device) + for i in range(1, len(chain)): + R = qmul(R, quat_params[:, chain[i]]) + offset_vec = offsets[:, chain[i]] + joints[:, chain[i]] = qrot(R, offset_vec) + joints[:, chain[i-1]] + return joints + + # Be sure root joint is at the beginning of kinematic chains + def forward_kinematics_np(self, quat_params, root_pos, skel_joints=None, do_root_R=True): + # quat_params (batch_size, joints_num, 4) + # joints (batch_size, joints_num, 3) + # root_pos (batch_size, 3) + if skel_joints is not None: + skel_joints = torch.from_numpy(skel_joints) + offsets = self.get_offsets_joints_batch(skel_joints) + if len(self._offset.shape) == 2: + offsets = self._offset.expand(quat_params.shape[0], -1, -1) + offsets = offsets.numpy() + joints = np.zeros(quat_params.shape[:-1] + (3,)) + joints[:, 0] = root_pos + for chain in self._kinematic_tree: + if do_root_R: + R = quat_params[:, 0] + else: + R = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(len(quat_params), axis=0) + for i in range(1, len(chain)): + R = qmul_np(R, quat_params[:, chain[i]]) + offset_vec = offsets[:, chain[i]] + joints[:, chain[i]] = qrot_np(R, offset_vec) + joints[:, chain[i - 1]] + return joints + + def forward_kinematics_cont6d_np(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True): + # cont6d_params (batch_size, joints_num, 6) + # joints (batch_size, joints_num, 3) + # root_pos (batch_size, 3) + if skel_joints is not None: + skel_joints = torch.from_numpy(skel_joints) + offsets = self.get_offsets_joints_batch(skel_joints) + if len(self._offset.shape) == 2: + offsets = self._offset.expand(cont6d_params.shape[0], -1, -1) + offsets = offsets.numpy() + joints = np.zeros(cont6d_params.shape[:-1] + (3,)) + joints[:, 0] = root_pos + for chain in self._kinematic_tree: + if do_root_R: + matR = cont6d_to_matrix_np(cont6d_params[:, 0]) + else: + matR = np.eye(3)[np.newaxis, :].repeat(len(cont6d_params), axis=0) + for i in range(1, len(chain)): + matR = np.matmul(matR, cont6d_to_matrix_np(cont6d_params[:, chain[i]])) + offset_vec = offsets[:, chain[i]][..., np.newaxis] + # print(matR.shape, offset_vec.shape) + joints[:, chain[i]] = np.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]] + return joints + + def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True): + # cont6d_params (batch_size, joints_num, 6) + # joints (batch_size, joints_num, 3) + # root_pos (batch_size, 3) + if skel_joints is not None: + # skel_joints = torch.from_numpy(skel_joints) + offsets = self.get_offsets_joints_batch(skel_joints) + if len(self._offset.shape) == 2: + offsets = self._offset.expand(cont6d_params.shape[0], -1, -1) + joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device) + joints[..., 0, :] = root_pos + for chain in self._kinematic_tree: + if do_root_R: + matR = cont6d_to_matrix(cont6d_params[:, 0]) + else: + matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device) + for i in range(1, len(chain)): + matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]])) + offset_vec = offsets[:, chain[i]].unsqueeze(-1) + # print(matR.shape, offset_vec.shape) + joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]] + return joints + + + + + diff --git a/VQ-Trans/utils/utils_model.py b/VQ-Trans/utils/utils_model.py new file mode 100644 index 0000000000000000000000000000000000000000..b3653a47ddb96f2ba27aae73b4eef8be904e9bf0 --- /dev/null +++ b/VQ-Trans/utils/utils_model.py @@ -0,0 +1,66 @@ +import numpy as np +import torch +import torch.optim as optim +import logging +import os +import sys + +def getCi(accLog): + + mean = np.mean(accLog) + std = np.std(accLog) + ci95 = 1.96*std/np.sqrt(len(accLog)) + + return mean, ci95 + +def get_logger(out_dir): + logger = logging.getLogger('Exp') + logger.setLevel(logging.INFO) + formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s") + + file_path = os.path.join(out_dir, "run.log") + file_hdlr = logging.FileHandler(file_path) + file_hdlr.setFormatter(formatter) + + strm_hdlr = logging.StreamHandler(sys.stdout) + strm_hdlr.setFormatter(formatter) + + logger.addHandler(file_hdlr) + logger.addHandler(strm_hdlr) + return logger + +## Optimizer +def initial_optim(decay_option, lr, weight_decay, net, optimizer) : + + if optimizer == 'adamw' : + optimizer_adam_family = optim.AdamW + elif optimizer == 'adam' : + optimizer_adam_family = optim.Adam + if decay_option == 'all': + #optimizer = optimizer_adam_family(net.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=weight_decay) + optimizer = optimizer_adam_family(net.parameters(), lr=lr, betas=(0.5, 0.9), weight_decay=weight_decay) + + elif decay_option == 'noVQ': + all_params = set(net.parameters()) + no_decay = set([net.vq_layer]) + + decay = all_params - no_decay + optimizer = optimizer_adam_family([ + {'params': list(no_decay), 'weight_decay': 0}, + {'params': list(decay), 'weight_decay' : weight_decay}], lr=lr) + + return optimizer + + +def get_motion_with_trans(motion, velocity) : + ''' + motion : torch.tensor, shape (batch_size, T, 72), with the global translation = 0 + velocity : torch.tensor, shape (batch_size, T, 3), contain the information of velocity = 0 + + ''' + trans = torch.cumsum(velocity, dim=1) + trans = trans - trans[:, :1] ## the first root is initialized at 0 (just for visualization) + trans = trans.repeat((1, 1, 21)) + motion_with_trans = motion + trans + return motion_with_trans + \ No newline at end of file diff --git a/VQ-Trans/utils/word_vectorizer.py b/VQ-Trans/utils/word_vectorizer.py new file mode 100644 index 0000000000000000000000000000000000000000..557ff97a9539c084167f3eca51fb50f53f33c8ea --- /dev/null +++ b/VQ-Trans/utils/word_vectorizer.py @@ -0,0 +1,99 @@ +import numpy as np +import pickle +from os.path import join as pjoin + +POS_enumerator = { + 'VERB': 0, + 'NOUN': 1, + 'DET': 2, + 'ADP': 3, + 'NUM': 4, + 'AUX': 5, + 'PRON': 6, + 'ADJ': 7, + 'ADV': 8, + 'Loc_VIP': 9, + 'Body_VIP': 10, + 'Obj_VIP': 11, + 'Act_VIP': 12, + 'Desc_VIP': 13, + 'OTHER': 14, +} + +Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward', + 'up', 'down', 'straight', 'curve') + +Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh') + +Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball') + +Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn', + 'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll', + 'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb') + +Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily', + 'angrily', 'sadly') + +VIP_dict = { + 'Loc_VIP': Loc_list, + 'Body_VIP': Body_list, + 'Obj_VIP': Obj_List, + 'Act_VIP': Act_list, + 'Desc_VIP': Desc_list, +} + + +class WordVectorizer(object): + def __init__(self, meta_root, prefix): + vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix)) + words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb')) + self.word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb')) + self.word2vec = {w: vectors[self.word2idx[w]] for w in words} + + def _get_pos_ohot(self, pos): + pos_vec = np.zeros(len(POS_enumerator)) + if pos in POS_enumerator: + pos_vec[POS_enumerator[pos]] = 1 + else: + pos_vec[POS_enumerator['OTHER']] = 1 + return pos_vec + + def __len__(self): + return len(self.word2vec) + + def __getitem__(self, item): + word, pos = item.split('/') + if word in self.word2vec: + word_vec = self.word2vec[word] + vip_pos = None + for key, values in VIP_dict.items(): + if word in values: + vip_pos = key + break + if vip_pos is not None: + pos_vec = self._get_pos_ohot(vip_pos) + else: + pos_vec = self._get_pos_ohot(pos) + else: + word_vec = self.word2vec['unk'] + pos_vec = self._get_pos_ohot('OTHER') + return word_vec, pos_vec + + +class WordVectorizerV2(WordVectorizer): + def __init__(self, meta_root, prefix): + super(WordVectorizerV2, self).__init__(meta_root, prefix) + self.idx2word = {self.word2idx[w]: w for w in self.word2idx} + + def __getitem__(self, item): + word_vec, pose_vec = super(WordVectorizerV2, self).__getitem__(item) + word, pos = item.split('/') + if word in self.word2vec: + return word_vec, pose_vec, self.word2idx[word] + else: + return word_vec, pose_vec, self.word2idx['unk'] + + def itos(self, idx): + if idx == len(self.idx2word): + return "pad" + return self.idx2word[idx] \ No newline at end of file diff --git a/VQ-Trans/visualization/plot_3d_global.py b/VQ-Trans/visualization/plot_3d_global.py new file mode 100644 index 0000000000000000000000000000000000000000..42fea4efd366397e17bc74470d72d3313ae228d8 --- /dev/null +++ b/VQ-Trans/visualization/plot_3d_global.py @@ -0,0 +1,129 @@ +import torch +import matplotlib.pyplot as plt +import numpy as np +import io +import matplotlib +from mpl_toolkits.mplot3d.art3d import Poly3DCollection +import mpl_toolkits.mplot3d.axes3d as p3 +from textwrap import wrap +import imageio + +def plot_3d_motion(args, figsize=(10, 10), fps=120, radius=4): + matplotlib.use('Agg') + + + joints, out_name, title = args + + data = joints.copy().reshape(len(joints), -1, 3) + + nb_joints = joints.shape[1] + smpl_kinetic_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]] if nb_joints == 21 else [[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]] + limits = 1000 if nb_joints == 21 else 2 + MINS = data.min(axis=0).min(axis=0) + MAXS = data.max(axis=0).max(axis=0) + colors = ['red', 'blue', 'black', 'red', 'blue', + 'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue', + 'darkred', 'darkred', 'darkred', 'darkred', 'darkred'] + frame_number = data.shape[0] + # print(data.shape) + + height_offset = MINS[1] + data[:, :, 1] -= height_offset + trajec = data[:, 0, [0, 2]] + + data[..., 0] -= data[:, 0:1, 0] + data[..., 2] -= data[:, 0:1, 2] + + def update(index): + + def init(): + ax.set_xlim(-limits, limits) + ax.set_ylim(-limits, limits) + ax.set_zlim(0, limits) + ax.grid(b=False) + def plot_xzPlane(minx, maxx, miny, minz, maxz): + ## Plot a plane XZ + verts = [ + [minx, miny, minz], + [minx, miny, maxz], + [maxx, miny, maxz], + [maxx, miny, minz] + ] + xz_plane = Poly3DCollection([verts]) + xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5)) + ax.add_collection3d(xz_plane) + fig = plt.figure(figsize=(480/96., 320/96.), dpi=96) if nb_joints == 21 else plt.figure(figsize=(10, 10), dpi=96) + if title is not None : + wraped_title = '\n'.join(wrap(title, 40)) + fig.suptitle(wraped_title, fontsize=16) + ax = p3.Axes3D(fig) + + init() + + ax.lines = [] + ax.collections = [] + ax.view_init(elev=110, azim=-90) + ax.dist = 7.5 + # ax = + plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1], + MAXS[2] - trajec[index, 1]) + # ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3) + + if index > 1: + ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]), + trajec[:index, 1] - trajec[index, 1], linewidth=1.0, + color='blue') + # ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2]) + + for i, (chain, color) in enumerate(zip(smpl_kinetic_chain, colors)): + # print(color) + if i < 5: + linewidth = 4.0 + else: + linewidth = 2.0 + ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth, + color=color) + # print(trajec[:index, 0].shape) + + plt.axis('off') + ax.set_xticklabels([]) + ax.set_yticklabels([]) + ax.set_zticklabels([]) + + if out_name is not None : + plt.savefig(out_name, dpi=96) + plt.close() + + else : + io_buf = io.BytesIO() + fig.savefig(io_buf, format='raw', dpi=96) + io_buf.seek(0) + # print(fig.bbox.bounds) + arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8), + newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1)) + io_buf.close() + plt.close() + return arr + + out = [] + for i in range(frame_number) : + out.append(update(i)) + out = np.stack(out, axis=0) + return torch.from_numpy(out) + + +def draw_to_batch(smpl_joints_batch, title_batch=None, outname=None) : + + batch_size = len(smpl_joints_batch) + out = [] + for i in range(batch_size) : + out.append(plot_3d_motion([smpl_joints_batch[i], None, title_batch[i] if title_batch is not None else None])) + if outname is not None: + imageio.mimsave(outname[i], np.array(out[-1]), fps=20) + out = torch.stack(out, axis=0) + return out + + + + + diff --git a/VQ-Trans/visualize/joints2smpl/smpl_models/SMPL_downsample_index.pkl b/VQ-Trans/visualize/joints2smpl/smpl_models/SMPL_downsample_index.pkl new file mode 100644 index 0000000000000000000000000000000000000000..7bb54c4f1e03340ad58b60485abaed1641d68d47 --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/smpl_models/SMPL_downsample_index.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e5b783c1677079397ee4bc26df5c72d73b8bb393bea41fa295b951187443daec +size 3556 diff --git a/VQ-Trans/visualize/joints2smpl/smpl_models/gmm_08.pkl b/VQ-Trans/visualize/joints2smpl/smpl_models/gmm_08.pkl new file mode 100644 index 0000000000000000000000000000000000000000..c97a1d7ef396581e56ce74a12cc39175680ce028 --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/smpl_models/gmm_08.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e1374908aae055a2afa01a2cd9a169bc6cfec1ceb7aa590e201a47b383060491 +size 839127 diff --git a/VQ-Trans/visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5 b/VQ-Trans/visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5 new file mode 100644 index 0000000000000000000000000000000000000000..b6ecce2a748128cfde09b219ccc74307de50bbae --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ac9b474c74daec0253ed084720f662059336e976850f08a4a9a3f76d06613776 +size 4848 diff --git a/VQ-Trans/visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl b/VQ-Trans/visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl new file mode 100644 index 0000000000000000000000000000000000000000..77ce98631741ba3887d689077baf35422d39299d --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb69c10801205c9cfb5353fdeb1b9cc5ade53d14c265c3339421cdde8b9c91e7 +size 1323168 diff --git a/VQ-Trans/visualize/joints2smpl/src/config.py b/VQ-Trans/visualize/joints2smpl/src/config.py new file mode 100644 index 0000000000000000000000000000000000000000..1021115a53f19974fbea3d3768c25874a4ae5d38 --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/src/config.py @@ -0,0 +1,40 @@ +import numpy as np + +# Map joints Name to SMPL joints idx +JOINT_MAP = { +'MidHip': 0, +'LHip': 1, 'LKnee': 4, 'LAnkle': 7, 'LFoot': 10, +'RHip': 2, 'RKnee': 5, 'RAnkle': 8, 'RFoot': 11, +'LShoulder': 16, 'LElbow': 18, 'LWrist': 20, 'LHand': 22, +'RShoulder': 17, 'RElbow': 19, 'RWrist': 21, 'RHand': 23, +'spine1': 3, 'spine2': 6, 'spine3': 9, 'Neck': 12, 'Head': 15, +'LCollar':13, 'Rcollar' :14, +'Nose':24, 'REye':26, 'LEye':26, 'REar':27, 'LEar':28, +'LHeel': 31, 'RHeel': 34, +'OP RShoulder': 17, 'OP LShoulder': 16, +'OP RHip': 2, 'OP LHip': 1, +'OP Neck': 12, +} + +full_smpl_idx = range(24) +key_smpl_idx = [0, 1, 4, 7, 2, 5, 8, 17, 19, 21, 16, 18, 20] + + +AMASS_JOINT_MAP = { +'MidHip': 0, +'LHip': 1, 'LKnee': 4, 'LAnkle': 7, 'LFoot': 10, +'RHip': 2, 'RKnee': 5, 'RAnkle': 8, 'RFoot': 11, +'LShoulder': 16, 'LElbow': 18, 'LWrist': 20, +'RShoulder': 17, 'RElbow': 19, 'RWrist': 21, +'spine1': 3, 'spine2': 6, 'spine3': 9, 'Neck': 12, 'Head': 15, +'LCollar':13, 'Rcollar' :14, +} +amass_idx = range(22) +amass_smpl_idx = range(22) + + +SMPL_MODEL_DIR = "./body_models/" +GMM_MODEL_DIR = "./visualize/joints2smpl/smpl_models/" +SMPL_MEAN_FILE = "./visualize/joints2smpl/smpl_models/neutral_smpl_mean_params.h5" +# for collsion +Part_Seg_DIR = "./visualize/joints2smpl/smpl_models/smplx_parts_segm.pkl" \ No newline at end of file diff --git a/VQ-Trans/visualize/joints2smpl/src/customloss.py b/VQ-Trans/visualize/joints2smpl/src/customloss.py new file mode 100644 index 0000000000000000000000000000000000000000..880ab4861c58cec9faeb086e430fde7387c5cc9e --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/src/customloss.py @@ -0,0 +1,222 @@ +import torch +import torch.nn.functional as F +from visualize.joints2smpl.src import config + +# Guassian +def gmof(x, sigma): + """ + Geman-McClure error function + """ + x_squared = x ** 2 + sigma_squared = sigma ** 2 + return (sigma_squared * x_squared) / (sigma_squared + x_squared) + +# angle prior +def angle_prior(pose): + """ + Angle prior that penalizes unnatural bending of the knees and elbows + """ + # We subtract 3 because pose does not include the global rotation of the model + return torch.exp( + pose[:, [55 - 3, 58 - 3, 12 - 3, 15 - 3]] * torch.tensor([1., -1., -1, -1.], device=pose.device)) ** 2 + + +def perspective_projection(points, rotation, translation, + focal_length, camera_center): + """ + This function computes the perspective projection of a set of points. + Input: + points (bs, N, 3): 3D points + rotation (bs, 3, 3): Camera rotation + translation (bs, 3): Camera translation + focal_length (bs,) or scalar: Focal length + camera_center (bs, 2): Camera center + """ + batch_size = points.shape[0] + K = torch.zeros([batch_size, 3, 3], device=points.device) + K[:, 0, 0] = focal_length + K[:, 1, 1] = focal_length + K[:, 2, 2] = 1. + K[:, :-1, -1] = camera_center + + # Transform points + points = torch.einsum('bij,bkj->bki', rotation, points) + points = points + translation.unsqueeze(1) + + # Apply perspective distortion + projected_points = points / points[:, :, -1].unsqueeze(-1) + + # Apply camera intrinsics + projected_points = torch.einsum('bij,bkj->bki', K, projected_points) + + return projected_points[:, :, :-1] + + +def body_fitting_loss(body_pose, betas, model_joints, camera_t, camera_center, + joints_2d, joints_conf, pose_prior, + focal_length=5000, sigma=100, pose_prior_weight=4.78, + shape_prior_weight=5, angle_prior_weight=15.2, + output='sum'): + """ + Loss function for body fitting + """ + batch_size = body_pose.shape[0] + rotation = torch.eye(3, device=body_pose.device).unsqueeze(0).expand(batch_size, -1, -1) + + projected_joints = perspective_projection(model_joints, rotation, camera_t, + focal_length, camera_center) + + # Weighted robust reprojection error + reprojection_error = gmof(projected_joints - joints_2d, sigma) + reprojection_loss = (joints_conf ** 2) * reprojection_error.sum(dim=-1) + + # Pose prior loss + pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas) + + # Angle prior for knees and elbows + angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1) + + # Regularizer to prevent betas from taking large values + shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1) + + total_loss = reprojection_loss.sum(dim=-1) + pose_prior_loss + angle_prior_loss + shape_prior_loss + + if output == 'sum': + return total_loss.sum() + elif output == 'reprojection': + return reprojection_loss + + +# --- get camera fitting loss ----- +def camera_fitting_loss(model_joints, camera_t, camera_t_est, camera_center, + joints_2d, joints_conf, + focal_length=5000, depth_loss_weight=100): + """ + Loss function for camera optimization. + """ + # Project model joints + batch_size = model_joints.shape[0] + rotation = torch.eye(3, device=model_joints.device).unsqueeze(0).expand(batch_size, -1, -1) + projected_joints = perspective_projection(model_joints, rotation, camera_t, + focal_length, camera_center) + + # get the indexed four + op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder'] + op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints] + gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder'] + gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints] + + reprojection_error_op = (joints_2d[:, op_joints_ind] - + projected_joints[:, op_joints_ind]) ** 2 + reprojection_error_gt = (joints_2d[:, gt_joints_ind] - + projected_joints[:, gt_joints_ind]) ** 2 + + # Check if for each example in the batch all 4 OpenPose detections are valid, otherwise use the GT detections + # OpenPose joints are more reliable for this task, so we prefer to use them if possible + is_valid = (joints_conf[:, op_joints_ind].min(dim=-1)[0][:, None, None] > 0).float() + reprojection_loss = (is_valid * reprojection_error_op + (1 - is_valid) * reprojection_error_gt).sum(dim=(1, 2)) + + # Loss that penalizes deviation from depth estimate + depth_loss = (depth_loss_weight ** 2) * (camera_t[:, 2] - camera_t_est[:, 2]) ** 2 + + total_loss = reprojection_loss + depth_loss + return total_loss.sum() + + + + # #####--- body fitiing loss ----- +def body_fitting_loss_3d(body_pose, preserve_pose, + betas, model_joints, camera_translation, + j3d, pose_prior, + joints3d_conf, + sigma=100, pose_prior_weight=4.78*1.5, + shape_prior_weight=5.0, angle_prior_weight=15.2, + joint_loss_weight=500.0, + pose_preserve_weight=0.0, + use_collision=False, + model_vertices=None, model_faces=None, + search_tree=None, pen_distance=None, filter_faces=None, + collision_loss_weight=1000 + ): + """ + Loss function for body fitting + """ + batch_size = body_pose.shape[0] + + #joint3d_loss = (joint_loss_weight ** 2) * gmof((model_joints + camera_translation) - j3d, sigma).sum(dim=-1) + + joint3d_error = gmof((model_joints + camera_translation) - j3d, sigma) + + joint3d_loss_part = (joints3d_conf ** 2) * joint3d_error.sum(dim=-1) + joint3d_loss = ((joint_loss_weight ** 2) * joint3d_loss_part).sum(dim=-1) + + # Pose prior loss + pose_prior_loss = (pose_prior_weight ** 2) * pose_prior(body_pose, betas) + # Angle prior for knees and elbows + angle_prior_loss = (angle_prior_weight ** 2) * angle_prior(body_pose).sum(dim=-1) + # Regularizer to prevent betas from taking large values + shape_prior_loss = (shape_prior_weight ** 2) * (betas ** 2).sum(dim=-1) + + collision_loss = 0.0 + # Calculate the loss due to interpenetration + if use_collision: + triangles = torch.index_select( + model_vertices, 1, + model_faces).view(batch_size, -1, 3, 3) + + with torch.no_grad(): + collision_idxs = search_tree(triangles) + + # Remove unwanted collisions + if filter_faces is not None: + collision_idxs = filter_faces(collision_idxs) + + if collision_idxs.ge(0).sum().item() > 0: + collision_loss = torch.sum(collision_loss_weight * pen_distance(triangles, collision_idxs)) + + pose_preserve_loss = (pose_preserve_weight ** 2) * ((body_pose - preserve_pose) ** 2).sum(dim=-1) + + # print('joint3d_loss', joint3d_loss.shape) + # print('pose_prior_loss', pose_prior_loss.shape) + # print('angle_prior_loss', angle_prior_loss.shape) + # print('shape_prior_loss', shape_prior_loss.shape) + # print('collision_loss', collision_loss) + # print('pose_preserve_loss', pose_preserve_loss.shape) + + total_loss = joint3d_loss + pose_prior_loss + angle_prior_loss + shape_prior_loss + collision_loss + pose_preserve_loss + + return total_loss.sum() + + +# #####--- get camera fitting loss ----- +def camera_fitting_loss_3d(model_joints, camera_t, camera_t_est, + j3d, joints_category="orig", depth_loss_weight=100.0): + """ + Loss function for camera optimization. + """ + model_joints = model_joints + camera_t + # # get the indexed four + # op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder'] + # op_joints_ind = [config.JOINT_MAP[joint] for joint in op_joints] + # + # j3d_error_loss = (j3d[:, op_joints_ind] - + # model_joints[:, op_joints_ind]) ** 2 + + gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder'] + gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints] + + if joints_category=="orig": + select_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints] + elif joints_category=="AMASS": + select_joints_ind = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints] + else: + print("NO SUCH JOINTS CATEGORY!") + + j3d_error_loss = (j3d[:, select_joints_ind] - + model_joints[:, gt_joints_ind]) ** 2 + + # Loss that penalizes deviation from depth estimate + depth_loss = (depth_loss_weight**2) * (camera_t - camera_t_est)**2 + + total_loss = j3d_error_loss + depth_loss + return total_loss.sum() diff --git a/VQ-Trans/visualize/joints2smpl/src/prior.py b/VQ-Trans/visualize/joints2smpl/src/prior.py new file mode 100644 index 0000000000000000000000000000000000000000..7f13806dd1f6607507b0c7e5ad463b3fb0026be8 --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/src/prior.py @@ -0,0 +1,230 @@ +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from __future__ import absolute_import +from __future__ import print_function +from __future__ import division + +import sys +import os + +import time +import pickle + +import numpy as np + +import torch +import torch.nn as nn + +DEFAULT_DTYPE = torch.float32 + + +def create_prior(prior_type, **kwargs): + if prior_type == 'gmm': + prior = MaxMixturePrior(**kwargs) + elif prior_type == 'l2': + return L2Prior(**kwargs) + elif prior_type == 'angle': + return SMPLifyAnglePrior(**kwargs) + elif prior_type == 'none' or prior_type is None: + # Don't use any pose prior + def no_prior(*args, **kwargs): + return 0.0 + prior = no_prior + else: + raise ValueError('Prior {}'.format(prior_type) + ' is not implemented') + return prior + + +class SMPLifyAnglePrior(nn.Module): + def __init__(self, dtype=torch.float32, **kwargs): + super(SMPLifyAnglePrior, self).__init__() + + # Indices for the roration angle of + # 55: left elbow, 90deg bend at -np.pi/2 + # 58: right elbow, 90deg bend at np.pi/2 + # 12: left knee, 90deg bend at np.pi/2 + # 15: right knee, 90deg bend at np.pi/2 + angle_prior_idxs = np.array([55, 58, 12, 15], dtype=np.int64) + angle_prior_idxs = torch.tensor(angle_prior_idxs, dtype=torch.long) + self.register_buffer('angle_prior_idxs', angle_prior_idxs) + + angle_prior_signs = np.array([1, -1, -1, -1], + dtype=np.float32 if dtype == torch.float32 + else np.float64) + angle_prior_signs = torch.tensor(angle_prior_signs, + dtype=dtype) + self.register_buffer('angle_prior_signs', angle_prior_signs) + + def forward(self, pose, with_global_pose=False): + ''' Returns the angle prior loss for the given pose + + Args: + pose: (Bx[23 + 1] * 3) torch tensor with the axis-angle + representation of the rotations of the joints of the SMPL model. + Kwargs: + with_global_pose: Whether the pose vector also contains the global + orientation of the SMPL model. If not then the indices must be + corrected. + Returns: + A sze (B) tensor containing the angle prior loss for each element + in the batch. + ''' + angle_prior_idxs = self.angle_prior_idxs - (not with_global_pose) * 3 + return torch.exp(pose[:, angle_prior_idxs] * + self.angle_prior_signs).pow(2) + + +class L2Prior(nn.Module): + def __init__(self, dtype=DEFAULT_DTYPE, reduction='sum', **kwargs): + super(L2Prior, self).__init__() + + def forward(self, module_input, *args): + return torch.sum(module_input.pow(2)) + + +class MaxMixturePrior(nn.Module): + + def __init__(self, prior_folder='prior', + num_gaussians=6, dtype=DEFAULT_DTYPE, epsilon=1e-16, + use_merged=True, + **kwargs): + super(MaxMixturePrior, self).__init__() + + if dtype == DEFAULT_DTYPE: + np_dtype = np.float32 + elif dtype == torch.float64: + np_dtype = np.float64 + else: + print('Unknown float type {}, exiting!'.format(dtype)) + sys.exit(-1) + + self.num_gaussians = num_gaussians + self.epsilon = epsilon + self.use_merged = use_merged + gmm_fn = 'gmm_{:02d}.pkl'.format(num_gaussians) + + full_gmm_fn = os.path.join(prior_folder, gmm_fn) + if not os.path.exists(full_gmm_fn): + print('The path to the mixture prior "{}"'.format(full_gmm_fn) + + ' does not exist, exiting!') + sys.exit(-1) + + with open(full_gmm_fn, 'rb') as f: + gmm = pickle.load(f, encoding='latin1') + + if type(gmm) == dict: + means = gmm['means'].astype(np_dtype) + covs = gmm['covars'].astype(np_dtype) + weights = gmm['weights'].astype(np_dtype) + elif 'sklearn.mixture.gmm.GMM' in str(type(gmm)): + means = gmm.means_.astype(np_dtype) + covs = gmm.covars_.astype(np_dtype) + weights = gmm.weights_.astype(np_dtype) + else: + print('Unknown type for the prior: {}, exiting!'.format(type(gmm))) + sys.exit(-1) + + self.register_buffer('means', torch.tensor(means, dtype=dtype)) + + self.register_buffer('covs', torch.tensor(covs, dtype=dtype)) + + precisions = [np.linalg.inv(cov) for cov in covs] + precisions = np.stack(precisions).astype(np_dtype) + + self.register_buffer('precisions', + torch.tensor(precisions, dtype=dtype)) + + # The constant term: + sqrdets = np.array([(np.sqrt(np.linalg.det(c))) + for c in gmm['covars']]) + const = (2 * np.pi)**(69 / 2.) + + nll_weights = np.asarray(gmm['weights'] / (const * + (sqrdets / sqrdets.min()))) + nll_weights = torch.tensor(nll_weights, dtype=dtype).unsqueeze(dim=0) + self.register_buffer('nll_weights', nll_weights) + + weights = torch.tensor(gmm['weights'], dtype=dtype).unsqueeze(dim=0) + self.register_buffer('weights', weights) + + self.register_buffer('pi_term', + torch.log(torch.tensor(2 * np.pi, dtype=dtype))) + + cov_dets = [np.log(np.linalg.det(cov.astype(np_dtype)) + epsilon) + for cov in covs] + self.register_buffer('cov_dets', + torch.tensor(cov_dets, dtype=dtype)) + + # The dimensionality of the random variable + self.random_var_dim = self.means.shape[1] + + def get_mean(self): + ''' Returns the mean of the mixture ''' + mean_pose = torch.matmul(self.weights, self.means) + return mean_pose + + def merged_log_likelihood(self, pose, betas): + diff_from_mean = pose.unsqueeze(dim=1) - self.means + + prec_diff_prod = torch.einsum('mij,bmj->bmi', + [self.precisions, diff_from_mean]) + diff_prec_quadratic = (prec_diff_prod * diff_from_mean).sum(dim=-1) + + curr_loglikelihood = 0.5 * diff_prec_quadratic - \ + torch.log(self.nll_weights) + # curr_loglikelihood = 0.5 * (self.cov_dets.unsqueeze(dim=0) + + # self.random_var_dim * self.pi_term + + # diff_prec_quadratic + # ) - torch.log(self.weights) + + min_likelihood, _ = torch.min(curr_loglikelihood, dim=1) + return min_likelihood + + def log_likelihood(self, pose, betas, *args, **kwargs): + ''' Create graph operation for negative log-likelihood calculation + ''' + likelihoods = [] + + for idx in range(self.num_gaussians): + mean = self.means[idx] + prec = self.precisions[idx] + cov = self.covs[idx] + diff_from_mean = pose - mean + + curr_loglikelihood = torch.einsum('bj,ji->bi', + [diff_from_mean, prec]) + curr_loglikelihood = torch.einsum('bi,bi->b', + [curr_loglikelihood, + diff_from_mean]) + cov_term = torch.log(torch.det(cov) + self.epsilon) + curr_loglikelihood += 0.5 * (cov_term + + self.random_var_dim * + self.pi_term) + likelihoods.append(curr_loglikelihood) + + log_likelihoods = torch.stack(likelihoods, dim=1) + min_idx = torch.argmin(log_likelihoods, dim=1) + weight_component = self.nll_weights[:, min_idx] + weight_component = -torch.log(weight_component) + + return weight_component + log_likelihoods[:, min_idx] + + def forward(self, pose, betas): + if self.use_merged: + return self.merged_log_likelihood(pose, betas) + else: + return self.log_likelihood(pose, betas) \ No newline at end of file diff --git a/VQ-Trans/visualize/joints2smpl/src/smplify.py b/VQ-Trans/visualize/joints2smpl/src/smplify.py new file mode 100644 index 0000000000000000000000000000000000000000..580efef98dfdcf6e7486b7f5c5436820edfb6c4b --- /dev/null +++ b/VQ-Trans/visualize/joints2smpl/src/smplify.py @@ -0,0 +1,279 @@ +import torch +import os, sys +import pickle +import smplx +import numpy as np + +sys.path.append(os.path.dirname(__file__)) +from customloss import (camera_fitting_loss, + body_fitting_loss, + camera_fitting_loss_3d, + body_fitting_loss_3d, + ) +from prior import MaxMixturePrior +from visualize.joints2smpl.src import config + + + +@torch.no_grad() +def guess_init_3d(model_joints, + j3d, + joints_category="orig"): + """Initialize the camera translation via triangle similarity, by using the torso joints . + :param model_joints: SMPL model with pre joints + :param j3d: 25x3 array of Kinect Joints + :returns: 3D vector corresponding to the estimated camera translation + """ + # get the indexed four + gt_joints = ['RHip', 'LHip', 'RShoulder', 'LShoulder'] + gt_joints_ind = [config.JOINT_MAP[joint] for joint in gt_joints] + + if joints_category=="orig": + joints_ind_category = [config.JOINT_MAP[joint] for joint in gt_joints] + elif joints_category=="AMASS": + joints_ind_category = [config.AMASS_JOINT_MAP[joint] for joint in gt_joints] + else: + print("NO SUCH JOINTS CATEGORY!") + + sum_init_t = (j3d[:, joints_ind_category] - model_joints[:, gt_joints_ind]).sum(dim=1) + init_t = sum_init_t / 4.0 + return init_t + + +# SMPLIfy 3D +class SMPLify3D(): + """Implementation of SMPLify, use 3D joints.""" + + def __init__(self, + smplxmodel, + step_size=1e-2, + batch_size=1, + num_iters=100, + use_collision=False, + use_lbfgs=True, + joints_category="orig", + device=torch.device('cuda:0'), + ): + + # Store options + self.batch_size = batch_size + self.device = device + self.step_size = step_size + + self.num_iters = num_iters + # --- choose optimizer + self.use_lbfgs = use_lbfgs + # GMM pose prior + self.pose_prior = MaxMixturePrior(prior_folder=config.GMM_MODEL_DIR, + num_gaussians=8, + dtype=torch.float32).to(device) + # collision part + self.use_collision = use_collision + if self.use_collision: + self.part_segm_fn = config.Part_Seg_DIR + + # reLoad SMPL-X model + self.smpl = smplxmodel + + self.model_faces = smplxmodel.faces_tensor.view(-1) + + # select joint joint_category + self.joints_category = joints_category + + if joints_category=="orig": + self.smpl_index = config.full_smpl_idx + self.corr_index = config.full_smpl_idx + elif joints_category=="AMASS": + self.smpl_index = config.amass_smpl_idx + self.corr_index = config.amass_idx + else: + self.smpl_index = None + self.corr_index = None + print("NO SUCH JOINTS CATEGORY!") + + # ---- get the man function here ------ + def __call__(self, init_pose, init_betas, init_cam_t, j3d, conf_3d=1.0, seq_ind=0): + """Perform body fitting. + Input: + init_pose: SMPL pose estimate + init_betas: SMPL betas estimate + init_cam_t: Camera translation estimate + j3d: joints 3d aka keypoints + conf_3d: confidence for 3d joints + seq_ind: index of the sequence + Returns: + vertices: Vertices of optimized shape + joints: 3D joints of optimized shape + pose: SMPL pose parameters of optimized shape + betas: SMPL beta parameters of optimized shape + camera_translation: Camera translation + """ + + # # # add the mesh inter-section to avoid + search_tree = None + pen_distance = None + filter_faces = None + + if self.use_collision: + from mesh_intersection.bvh_search_tree import BVH + import mesh_intersection.loss as collisions_loss + from mesh_intersection.filter_faces import FilterFaces + + search_tree = BVH(max_collisions=8) + + pen_distance = collisions_loss.DistanceFieldPenetrationLoss( + sigma=0.5, point2plane=False, vectorized=True, penalize_outside=True) + + if self.part_segm_fn: + # Read the part segmentation + part_segm_fn = os.path.expandvars(self.part_segm_fn) + with open(part_segm_fn, 'rb') as faces_parents_file: + face_segm_data = pickle.load(faces_parents_file, encoding='latin1') + faces_segm = face_segm_data['segm'] + faces_parents = face_segm_data['parents'] + # Create the module used to filter invalid collision pairs + filter_faces = FilterFaces( + faces_segm=faces_segm, faces_parents=faces_parents, + ign_part_pairs=None).to(device=self.device) + + + # Split SMPL pose to body pose and global orientation + body_pose = init_pose[:, 3:].detach().clone() + global_orient = init_pose[:, :3].detach().clone() + betas = init_betas.detach().clone() + + # use guess 3d to get the initial + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas) + model_joints = smpl_output.joints + + init_cam_t = guess_init_3d(model_joints, j3d, self.joints_category).unsqueeze(1).detach() + camera_translation = init_cam_t.clone() + + preserve_pose = init_pose[:, 3:].detach().clone() + # -------------Step 1: Optimize camera translation and body orientation-------- + # Optimize only camera translation and body orientation + body_pose.requires_grad = False + betas.requires_grad = False + global_orient.requires_grad = True + camera_translation.requires_grad = True + + camera_opt_params = [global_orient, camera_translation] + + if self.use_lbfgs: + camera_optimizer = torch.optim.LBFGS(camera_opt_params, max_iter=self.num_iters, + lr=self.step_size, line_search_fn='strong_wolfe') + for i in range(10): + def closure(): + camera_optimizer.zero_grad() + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas) + model_joints = smpl_output.joints + # print('model_joints', model_joints.shape) + # print('camera_translation', camera_translation.shape) + # print('init_cam_t', init_cam_t.shape) + # print('j3d', j3d.shape) + loss = camera_fitting_loss_3d(model_joints, camera_translation, + init_cam_t, j3d, self.joints_category) + loss.backward() + return loss + + camera_optimizer.step(closure) + else: + camera_optimizer = torch.optim.Adam(camera_opt_params, lr=self.step_size, betas=(0.9, 0.999)) + + for i in range(20): + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas) + model_joints = smpl_output.joints + + loss = camera_fitting_loss_3d(model_joints[:, self.smpl_index], camera_translation, + init_cam_t, j3d[:, self.corr_index], self.joints_category) + camera_optimizer.zero_grad() + loss.backward() + camera_optimizer.step() + + # Fix camera translation after optimizing camera + # --------Step 2: Optimize body joints -------------------------- + # Optimize only the body pose and global orientation of the body + body_pose.requires_grad = True + global_orient.requires_grad = True + camera_translation.requires_grad = True + + # --- if we use the sequence, fix the shape + if seq_ind == 0: + betas.requires_grad = True + body_opt_params = [body_pose, betas, global_orient, camera_translation] + else: + betas.requires_grad = False + body_opt_params = [body_pose, global_orient, camera_translation] + + if self.use_lbfgs: + body_optimizer = torch.optim.LBFGS(body_opt_params, max_iter=self.num_iters, + lr=self.step_size, line_search_fn='strong_wolfe') + for i in range(self.num_iters): + def closure(): + body_optimizer.zero_grad() + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas) + model_joints = smpl_output.joints + model_vertices = smpl_output.vertices + + loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation, + j3d[:, self.corr_index], self.pose_prior, + joints3d_conf=conf_3d, + joint_loss_weight=600.0, + pose_preserve_weight=5.0, + use_collision=self.use_collision, + model_vertices=model_vertices, model_faces=self.model_faces, + search_tree=search_tree, pen_distance=pen_distance, filter_faces=filter_faces) + loss.backward() + return loss + + body_optimizer.step(closure) + else: + body_optimizer = torch.optim.Adam(body_opt_params, lr=self.step_size, betas=(0.9, 0.999)) + + for i in range(self.num_iters): + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas) + model_joints = smpl_output.joints + model_vertices = smpl_output.vertices + + loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation, + j3d[:, self.corr_index], self.pose_prior, + joints3d_conf=conf_3d, + joint_loss_weight=600.0, + use_collision=self.use_collision, + model_vertices=model_vertices, model_faces=self.model_faces, + search_tree=search_tree, pen_distance=pen_distance, filter_faces=filter_faces) + body_optimizer.zero_grad() + loss.backward() + body_optimizer.step() + + # Get final loss value + with torch.no_grad(): + smpl_output = self.smpl(global_orient=global_orient, + body_pose=body_pose, + betas=betas, return_full_pose=True) + model_joints = smpl_output.joints + model_vertices = smpl_output.vertices + + final_loss = body_fitting_loss_3d(body_pose, preserve_pose, betas, model_joints[:, self.smpl_index], camera_translation, + j3d[:, self.corr_index], self.pose_prior, + joints3d_conf=conf_3d, + joint_loss_weight=600.0, + use_collision=self.use_collision, model_vertices=model_vertices, model_faces=self.model_faces, + search_tree=search_tree, pen_distance=pen_distance, filter_faces=filter_faces) + + vertices = smpl_output.vertices.detach() + joints = smpl_output.joints.detach() + pose = torch.cat([global_orient, body_pose], dim=-1).detach() + betas = betas.detach() + + return vertices, joints, pose, betas, camera_translation, final_loss diff --git a/VQ-Trans/visualize/render_mesh.py b/VQ-Trans/visualize/render_mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..d44d04f551ccb4f1ffc9efb4cb1a44c407ede836 --- /dev/null +++ b/VQ-Trans/visualize/render_mesh.py @@ -0,0 +1,33 @@ +import argparse +import os +from visualize import vis_utils +import shutil +from tqdm import tqdm + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--input_path", type=str, required=True, help='stick figure mp4 file to be rendered.') + parser.add_argument("--cuda", type=bool, default=True, help='') + parser.add_argument("--device", type=int, default=0, help='') + params = parser.parse_args() + + assert params.input_path.endswith('.mp4') + parsed_name = os.path.basename(params.input_path).replace('.mp4', '').replace('sample', '').replace('rep', '') + sample_i, rep_i = [int(e) for e in parsed_name.split('_')] + npy_path = os.path.join(os.path.dirname(params.input_path), 'results.npy') + out_npy_path = params.input_path.replace('.mp4', '_smpl_params.npy') + assert os.path.exists(npy_path) + results_dir = params.input_path.replace('.mp4', '_obj') + if os.path.exists(results_dir): + shutil.rmtree(results_dir) + os.makedirs(results_dir) + + npy2obj = vis_utils.npy2obj(npy_path, sample_i, rep_i, + device=params.device, cuda=params.cuda) + + print('Saving obj files to [{}]'.format(os.path.abspath(results_dir))) + for frame_i in tqdm(range(npy2obj.real_num_frames)): + npy2obj.save_obj(os.path.join(results_dir, 'frame{:03d}.obj'.format(frame_i)), frame_i) + + print('Saving SMPL params to [{}]'.format(os.path.abspath(out_npy_path))) + npy2obj.save_npy(out_npy_path) diff --git a/VQ-Trans/visualize/simplify_loc2rot.py b/VQ-Trans/visualize/simplify_loc2rot.py new file mode 100644 index 0000000000000000000000000000000000000000..5d3d4411310876033cb50d998ad64557a9c4b0c1 --- /dev/null +++ b/VQ-Trans/visualize/simplify_loc2rot.py @@ -0,0 +1,131 @@ +import numpy as np +import os +import torch +from visualize.joints2smpl.src import config +import smplx +import h5py +from visualize.joints2smpl.src.smplify import SMPLify3D +from tqdm import tqdm +import utils.rotation_conversions as geometry +import argparse + + +class joints2smpl: + + def __init__(self, num_frames, device_id, cuda=True): + self.device = torch.device("cuda:" + str(device_id) if cuda else "cpu") + # self.device = torch.device("cpu") + self.batch_size = num_frames + self.num_joints = 22 # for HumanML3D + self.joint_category = "AMASS" + self.num_smplify_iters = 150 + self.fix_foot = False + print(config.SMPL_MODEL_DIR) + smplmodel = smplx.create(config.SMPL_MODEL_DIR, + model_type="smpl", gender="neutral", ext="pkl", + batch_size=self.batch_size).to(self.device) + + # ## --- load the mean pose as original ---- + smpl_mean_file = config.SMPL_MEAN_FILE + + file = h5py.File(smpl_mean_file, 'r') + self.init_mean_pose = torch.from_numpy(file['pose'][:]).unsqueeze(0).repeat(self.batch_size, 1).float().to(self.device) + self.init_mean_shape = torch.from_numpy(file['shape'][:]).unsqueeze(0).repeat(self.batch_size, 1).float().to(self.device) + self.cam_trans_zero = torch.Tensor([0.0, 0.0, 0.0]).unsqueeze(0).to(self.device) + # + + # # #-------------initialize SMPLify + self.smplify = SMPLify3D(smplxmodel=smplmodel, + batch_size=self.batch_size, + joints_category=self.joint_category, + num_iters=self.num_smplify_iters, + device=self.device) + + + def npy2smpl(self, npy_path): + out_path = npy_path.replace('.npy', '_rot.npy') + motions = np.load(npy_path, allow_pickle=True)[None][0] + # print_batch('', motions) + n_samples = motions['motion'].shape[0] + all_thetas = [] + for sample_i in tqdm(range(n_samples)): + thetas, _ = self.joint2smpl(motions['motion'][sample_i].transpose(2, 0, 1)) # [nframes, njoints, 3] + all_thetas.append(thetas.cpu().numpy()) + motions['motion'] = np.concatenate(all_thetas, axis=0) + print('motions', motions['motion'].shape) + + print(f'Saving [{out_path}]') + np.save(out_path, motions) + exit() + + + + def joint2smpl(self, input_joints, init_params=None): + _smplify = self.smplify # if init_params is None else self.smplify_fast + pred_pose = torch.zeros(self.batch_size, 72).to(self.device) + pred_betas = torch.zeros(self.batch_size, 10).to(self.device) + pred_cam_t = torch.zeros(self.batch_size, 3).to(self.device) + keypoints_3d = torch.zeros(self.batch_size, self.num_joints, 3).to(self.device) + + # run the whole seqs + num_seqs = input_joints.shape[0] + + + # joints3d = input_joints[idx] # *1.2 #scale problem [check first] + keypoints_3d = torch.Tensor(input_joints).to(self.device).float() + + # if idx == 0: + if init_params is None: + pred_betas = self.init_mean_shape + pred_pose = self.init_mean_pose + pred_cam_t = self.cam_trans_zero + else: + pred_betas = init_params['betas'] + pred_pose = init_params['pose'] + pred_cam_t = init_params['cam'] + + if self.joint_category == "AMASS": + confidence_input = torch.ones(self.num_joints) + # make sure the foot and ankle + if self.fix_foot == True: + confidence_input[7] = 1.5 + confidence_input[8] = 1.5 + confidence_input[10] = 1.5 + confidence_input[11] = 1.5 + else: + print("Such category not settle down!") + + new_opt_vertices, new_opt_joints, new_opt_pose, new_opt_betas, \ + new_opt_cam_t, new_opt_joint_loss = _smplify( + pred_pose.detach(), + pred_betas.detach(), + pred_cam_t.detach(), + keypoints_3d, + conf_3d=confidence_input.to(self.device), + # seq_ind=idx + ) + + thetas = new_opt_pose.reshape(self.batch_size, 24, 3) + thetas = geometry.matrix_to_rotation_6d(geometry.axis_angle_to_matrix(thetas)) # [bs, 24, 6] + root_loc = torch.tensor(keypoints_3d[:, 0]) # [bs, 3] + root_loc = torch.cat([root_loc, torch.zeros_like(root_loc)], dim=-1).unsqueeze(1) # [bs, 1, 6] + thetas = torch.cat([thetas, root_loc], dim=1).unsqueeze(0).permute(0, 2, 3, 1) # [1, 25, 6, 196] + + return thetas.clone().detach(), {'pose': new_opt_joints[0, :24].flatten().clone().detach(), 'betas': new_opt_betas.clone().detach(), 'cam': new_opt_cam_t.clone().detach()} + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--input_path", type=str, required=True, help='Blender file or dir with blender files') + parser.add_argument("--cuda", type=bool, default=True, help='') + parser.add_argument("--device", type=int, default=0, help='') + params = parser.parse_args() + + simplify = joints2smpl(device_id=params.device, cuda=params.cuda) + + if os.path.isfile(params.input_path) and params.input_path.endswith('.npy'): + simplify.npy2smpl(params.input_path) + elif os.path.isdir(params.input_path): + files = [os.path.join(params.input_path, f) for f in os.listdir(params.input_path) if f.endswith('.npy')] + for f in files: + simplify.npy2smpl(f) \ No newline at end of file diff --git a/VQ-Trans/visualize/vis_utils.py b/VQ-Trans/visualize/vis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..05728b38e3d6be4bfd83324907e3fa7a3f358071 --- /dev/null +++ b/VQ-Trans/visualize/vis_utils.py @@ -0,0 +1,66 @@ +from model.rotation2xyz import Rotation2xyz +import numpy as np +from trimesh import Trimesh +import os +import torch +from visualize.simplify_loc2rot import joints2smpl + +class npy2obj: + def __init__(self, npy_path, sample_idx, rep_idx, device=0, cuda=True): + self.npy_path = npy_path + self.motions = np.load(self.npy_path, allow_pickle=True) + if self.npy_path.endswith('.npz'): + self.motions = self.motions['arr_0'] + self.motions = self.motions[None][0] + self.rot2xyz = Rotation2xyz(device='cpu') + self.faces = self.rot2xyz.smpl_model.faces + self.bs, self.njoints, self.nfeats, self.nframes = self.motions['motion'].shape + self.opt_cache = {} + self.sample_idx = sample_idx + self.total_num_samples = self.motions['num_samples'] + self.rep_idx = rep_idx + self.absl_idx = self.rep_idx*self.total_num_samples + self.sample_idx + self.num_frames = self.motions['motion'][self.absl_idx].shape[-1] + self.j2s = joints2smpl(num_frames=self.num_frames, device_id=device, cuda=cuda) + + if self.nfeats == 3: + print(f'Running SMPLify For sample [{sample_idx}], repetition [{rep_idx}], it may take a few minutes.') + motion_tensor, opt_dict = self.j2s.joint2smpl(self.motions['motion'][self.absl_idx].transpose(2, 0, 1)) # [nframes, njoints, 3] + self.motions['motion'] = motion_tensor.cpu().numpy() + elif self.nfeats == 6: + self.motions['motion'] = self.motions['motion'][[self.absl_idx]] + self.bs, self.njoints, self.nfeats, self.nframes = self.motions['motion'].shape + self.real_num_frames = self.motions['lengths'][self.absl_idx] + + self.vertices = self.rot2xyz(torch.tensor(self.motions['motion']), mask=None, + pose_rep='rot6d', translation=True, glob=True, + jointstype='vertices', + # jointstype='smpl', # for joint locations + vertstrans=True) + self.root_loc = self.motions['motion'][:, -1, :3, :].reshape(1, 1, 3, -1) + self.vertices += self.root_loc + + def get_vertices(self, sample_i, frame_i): + return self.vertices[sample_i, :, :, frame_i].squeeze().tolist() + + def get_trimesh(self, sample_i, frame_i): + return Trimesh(vertices=self.get_vertices(sample_i, frame_i), + faces=self.faces) + + def save_obj(self, save_path, frame_i): + mesh = self.get_trimesh(0, frame_i) + with open(save_path, 'w') as fw: + mesh.export(fw, 'obj') + return save_path + + def save_npy(self, save_path): + data_dict = { + 'motion': self.motions['motion'][0, :, :, :self.real_num_frames], + 'thetas': self.motions['motion'][0, :-1, :, :self.real_num_frames], + 'root_translation': self.motions['motion'][0, -1, :3, :self.real_num_frames], + 'faces': self.faces, + 'vertices': self.vertices[0, :, :, :self.real_num_frames], + 'text': self.motions['text'][0], + 'length': self.real_num_frames, + } + np.save(save_path, data_dict)