├── visualizer ├── tools │ ├── __init__.py │ ├── cfg_parser.py │ ├── objectmodel.py │ ├── utils.py │ └── meshviewer.py └── vis │ ├── visualize_finger_results.py │ ├── visualize_navigation_results.py │ ├── visualize_interaction_results.py │ └── visualize_long_sequence_results.py ├── physics_tracking ├── .gitignore ├── assets │ ├── largebox_cube.xml │ ├── largebox_1_cube.xml │ ├── largebox_2_cube.xml │ ├── largebox_3_cube.xml │ ├── largebox_cube.urdf │ ├── largebox_1_cube.urdf │ ├── largebox_2_cube.urdf │ └── largebox_3_cube.urdf ├── README.md ├── config_tracking_seq1_cube.py ├── config_tracking.py ├── utils.py └── models.py ├── manip ├── utils │ ├── visualize │ │ └── tools │ │ │ ├── __init__.py │ │ │ ├── cfg_parser.py │ │ │ ├── objectmodel.py │ │ │ ├── utils.py │ │ │ └── meshviewer.py │ ├── smplx_utils.py │ └── model_utils.py ├── inertialize │ ├── spring.py │ └── inert.py └── vis │ └── blender_vis_mesh_motion.py ├── bps.pt ├── grasp_generation ├── .gitignore ├── mano │ ├── MANO_RIGHT.pkl │ ├── pose_distrib.pt │ └── contact_indices.json └── utils │ ├── logger.py │ ├── energy.py │ ├── save.py │ ├── optimizer.py │ ├── object_model.py │ └── initializations.py ├── teaser.png ├── rest_hand_pose.pkl ├── local_joint_rot_mat.pkl ├── all_object_data_dict_for_eval.pkl ├── .gitignore ├── data └── smpl_all_models │ ├── smpl_parents_22.npy │ ├── smpl_parents_24.npy │ ├── smplx_parents_52.npy │ ├── MANO_SMPLX_face_ids.pkl │ ├── palm_sample_indices.pkl │ ├── MANO_SMPLX_vertex_ids.pkl │ ├── palm_sample_vertices.pkl │ └── smplx_parents_onlyhand_32.npy ├── scripts ├── sample.sh ├── train │ ├── train_coarsenet.sh │ └── train_refinenet.sh ├── download_data.sh └── download_training_data.sh ├── README.md ├── environment.yml └── argument_parser.py /visualizer/tools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /physics_tracking/.gitignore: -------------------------------------------------------------------------------- 1 | ckpt* -------------------------------------------------------------------------------- /manip/utils/visualize/tools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /bps.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/bps.pt -------------------------------------------------------------------------------- /grasp_generation/.gitignore: -------------------------------------------------------------------------------- 1 | experiments 2 | 3 | objects/cache/ 4 | thirdparty/ 5 | -------------------------------------------------------------------------------- /teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/teaser.png -------------------------------------------------------------------------------- /rest_hand_pose.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/rest_hand_pose.pkl -------------------------------------------------------------------------------- /local_joint_rot_mat.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/local_joint_rot_mat.pkl -------------------------------------------------------------------------------- /all_object_data_dict_for_eval.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/all_object_data_dict_for_eval.pkl -------------------------------------------------------------------------------- /grasp_generation/mano/MANO_RIGHT.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/grasp_generation/mano/MANO_RIGHT.pkl -------------------------------------------------------------------------------- /grasp_generation/mano/pose_distrib.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/grasp_generation/mano/pose_distrib.pt -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .vscode 2 | __pycache__ 3 | 4 | data/* 5 | results/ 6 | experiments/ 7 | 8 | visualizer_results/ 9 | 10 | removed/ 11 | 12 | -------------------------------------------------------------------------------- /data/smpl_all_models/smpl_parents_22.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/smpl_parents_22.npy -------------------------------------------------------------------------------- /data/smpl_all_models/smpl_parents_24.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/smpl_parents_24.npy -------------------------------------------------------------------------------- /data/smpl_all_models/smplx_parents_52.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/smplx_parents_52.npy -------------------------------------------------------------------------------- /data/smpl_all_models/MANO_SMPLX_face_ids.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/MANO_SMPLX_face_ids.pkl -------------------------------------------------------------------------------- /data/smpl_all_models/palm_sample_indices.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/palm_sample_indices.pkl -------------------------------------------------------------------------------- /data/smpl_all_models/MANO_SMPLX_vertex_ids.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/MANO_SMPLX_vertex_ids.pkl -------------------------------------------------------------------------------- /data/smpl_all_models/palm_sample_vertices.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/palm_sample_vertices.pkl -------------------------------------------------------------------------------- /data/smpl_all_models/smplx_parents_onlyhand_32.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhenkirito123/hoifhli_release/HEAD/data/smpl_all_models/smplx_parents_onlyhand_32.npy -------------------------------------------------------------------------------- /scripts/sample.sh: -------------------------------------------------------------------------------- 1 | python sample.py \ 2 | --window=120 \ 3 | --data_root_folder="./data/processed_data" \ 4 | --project="./experiments" \ 5 | --test_sample_res \ 6 | --use_long_planned_path \ 7 | --add_interaction_root_xy_ori --add_interaction_feet_contact \ 8 | --add_finger_motion \ 9 | --use_guidance_in_denoising \ 10 | --vis_wdir="moving_box" -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_1_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_2_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_3_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /scripts/train/train_coarsenet.sh: -------------------------------------------------------------------------------- 1 | python trainer_interaction_motion_diffusion.py \ 2 | --window=120 \ 3 | --batch_size=32 \ 4 | --data_root_folder="./data/processed_data" \ 5 | --project="./experiments" \ 6 | --exp_name="train_cnet_release_test" \ 7 | --wandb_pj_name="final_interaction_diffusion" \ 8 | --add_object_in_wrist_loss \ 9 | --loss_w_feet=1 \ 10 | --loss_w_fk=0.5 \ 11 | --loss_w_obj_pts=1 -------------------------------------------------------------------------------- /grasp_generation/mano/contact_indices.json: -------------------------------------------------------------------------------- 1 | [ 2 | 699, 700, 753, 754, 714, 741, 755, 757, 739, 756, 760, 740, 762, 763, 3 | 194, 195, 165, 48, 49, 166, 46, 47, 280, 237, 238, 340, 341, 330, 342, 328, 343, 4 | 375, 386, 387, 358, 359, 376, 356, 357, 402, 396, 397, 452, 453, 440, 454, 438, 455, 5 | 485, 496, 497, 470, 471, 486, 468, 469, 513, 506, 507, 563, 564, 551, 565, 549, 566, 6 | 614, 615, 582, 583, 580, 581, 681, 681, 625, 666, 683, 7 | 73, 96, 98, 99, 772, 774, 775, 777 8 | ] -------------------------------------------------------------------------------- /scripts/train/train_refinenet.sh: -------------------------------------------------------------------------------- 1 | python trainer_interaction_motion_diffusion.py \ 2 | --window=120 \ 3 | --batch_size=32 \ 4 | --data_root_folder="./data/processed_data" \ 5 | --project="./experiments" \ 6 | --exp_name="train_rnet_release_test" \ 7 | --wandb_pj_name="final_interaction_diffusion" \ 8 | --add_object_in_wrist_loss \ 9 | --loss_w_feet=2 \ 10 | --loss_w_fk=0.5 \ 11 | --loss_w_obj_pts=1 \ 12 | --add_wrist_relative --add_object_static \ 13 | --add_interaction_root_xy_ori --add_interaction_feet_contact 14 | -------------------------------------------------------------------------------- /scripts/download_data.sh: -------------------------------------------------------------------------------- 1 | cd ./data 2 | gdown --folder https://drive.google.com/drive/folders/1LDo6RMjJ5Wr2bmL7hdTavMAU5z1Ex0GA?usp=drive_link 3 | gdown --folder https://drive.google.com/drive/folders/1-pexA27aGR2U9ucEufPfzPV_BRLKvDPG?usp=drive_link 4 | gdown --folder https://drive.google.com/drive/folders/12bRrFYV24P6_f6Xbklr7h7f0CD7TRmXC?usp=drive_link 5 | gdown --folder https://drive.google.com/drive/folders/1tXrrP-ph8NBckYfjAriwKe_9Lybl718L?usp=drive_link 6 | 7 | cd .. 8 | gdown --folder https://drive.google.com/drive/folders/1pJkRpKfLoolybgVE5PplNH1rhQ3XOgTE?usp=drive_link 9 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_cube.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_1_cube.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_2_cube.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /physics_tracking/assets/largebox_3_cube.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /physics_tracking/README.md: -------------------------------------------------------------------------------- 1 | ### Environment Setup 2 | Download and setup [Isaac Gym](https://developer.nvidia.com/isaac-gym). 3 | 4 | ### Example Usage 5 | #### Training 6 | ``` 7 | python main.py config_tracking_seq1_cube.py --ckpt ckpt_seq1_cube/ 8 | ``` 9 | 10 | You can monitor the training process using TensorBoard: 11 | ``` 12 | tensorboard --logdir ./ 13 | ``` 14 | 15 | #### Testing 16 | ``` 17 | python main.py config_tracking_seq1_cube.py --ckpt ckpt_seq1_cube/ --test 18 | ``` 19 | 20 | ### Custom Data 21 | You can use the motion generated by the previous step for training and testing. 22 | They can be found in `results/motion_params/seq_*.json`. 23 | -------------------------------------------------------------------------------- /physics_tracking/config_tracking_seq1_cube.py: -------------------------------------------------------------------------------- 1 | from config_tracking import * 2 | # env_cls = "ICCGANHumanoidDemo" 3 | env_params["character_model"] += ["assets/largebox_cube.urdf", "assets/largebox_1_cube.urdf", "assets/largebox_2_cube.urdf", "assets/largebox_3_cube.urdf"] 4 | env_params["motion_file"] = "seq_1.json" 5 | env_params["contactable_links"]["largebox"] = -1000 6 | env_params["contactable_links"]["largebox_1"] = -1000 7 | env_params["contactable_links"]["largebox_2"] = -1000 8 | env_params["contactable_links"]["largebox_3"] = -1000 9 | for s in ["L_", "R_"]: 10 | env_params["contactable_links"][s+"Wrist"] = 0.1 11 | for i in range(1, 5): 12 | for f in ["Index", "Middle", "Ring", "Pinky", "Thumb"]: 13 | env_params["contactable_links"][s+f+str(i)] = 0.1 14 | -------------------------------------------------------------------------------- /scripts/download_training_data.sh: -------------------------------------------------------------------------------- 1 | cd ./data 2 | rm -rf ./processed_data 3 | 4 | gdown --fuzzy https://drive.google.com/file/d/1vtx8rjTGwy6hAxnB1QsMPzZGBJNr1Ffa/view?usp=drive_link 5 | gdown --fuzzy https://drive.google.com/file/d/18JGdLLZwAIYNud6wBx8JZhHnAnSPXJmM/view?usp=drive_link 6 | gdown --fuzzy https://drive.google.com/file/d/1npkrQd2SM14wsFxKDrIE5fNMxzE0GWo-/view?usp=sharing 7 | gdown --fuzzy https://drive.google.com/file/d/1BYl1VUeeC-Vf_V9Sq1gw6Ft2Bcj_LyB_/view?usp=sharing 8 | gdown --fuzzy https://drive.google.com/file/d/1dxhsEQkJzio52jUSsUkCc5VCV-FrgSRx/view?usp=sharing 9 | gdown --fuzzy https://drive.google.com/file/d/17HjKySkDaZ-eJQ0jb3tuGtNC2ZuHnRRN/view?usp=sharing 10 | gdown --fuzzy https://drive.google.com/file/d/1CUOEbO8dPwIY8BrovszmzU5NqVfvcggo/view?usp=sharing 11 | 12 | cat processed_data.tar.gz.part* | tar -xzvf - 13 | 14 | rm -f processed_data.tar.gz.partaa 15 | rm -f processed_data.tar.gz.partab 16 | rm -f processed_data.tar.gz.partac 17 | rm -f processed_data.tar.gz.partad 18 | rm -f processed_data.tar.gz.partae 19 | rm -f processed_data.tar.gz.partaf 20 | rm -f processed_data.tar.gz.partag 21 | 22 | -------------------------------------------------------------------------------- /manip/inertialize/spring.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file is inspired by https://theorangeduck.com/page/spring-roll-call 3 | """ 4 | 5 | import math 6 | 7 | import numpy as np 8 | import torch 9 | 10 | # Constants 11 | LN2f = 0.69314718 12 | PIf = 3.14159265 13 | 14 | 15 | # Helper functions 16 | def square(x): 17 | return x * x 18 | 19 | 20 | def neg_exp(x): 21 | return math.exp(-x) 22 | 23 | 24 | def lerp(a, b, t): 25 | return a + t * (b - a) 26 | 27 | 28 | def clamp(value, min, max): 29 | return torch.clamp(value, min, max) 30 | 31 | 32 | def halflife_to_damping(halflife, eps=1e-5): 33 | return (4.0 * LN2f) / (halflife + eps) 34 | 35 | 36 | def damping_to_halflife(damping, eps=1e-5): 37 | return (4.0 * LN2f) / (damping + eps) 38 | 39 | 40 | def frequency_to_stiffness(frequency): 41 | return square(2.0 * PIf * frequency) 42 | 43 | 44 | def stiffness_to_frequency(stiffness): 45 | return torch.sqrt(stiffness) / (2.0 * PIf) 46 | 47 | 48 | def decay_spring_damper_exact(x, v, halflife, dt): 49 | y = halflife_to_damping(halflife) / 2.0 50 | j1 = v + x * y 51 | eydt = neg_exp(y * dt) 52 | 53 | x_new = eydt * (x + j1 * dt) 54 | return x_new 55 | 56 | 57 | def decay_spring_damper_exact_cubic(x, v, blendtime, dt, eps=1e-8): 58 | t = np.clip(dt / (blendtime + eps), 0.0, 1.0) 59 | 60 | d = x 61 | c = v * blendtime 62 | b = -3 * d - 2 * c 63 | a = 2 * d + c 64 | 65 | return a * t * t * t + b * t * t + c * t + d 66 | 67 | 68 | def inertialize_transition(off_x, off_v, src_x, src_v, dst_x, dst_v): 69 | off_x_new = (src_x + off_x) - dst_x 70 | off_v_new = (src_v + off_v) - dst_v 71 | return off_x_new, off_v_new 72 | 73 | 74 | def inertialize_update(out_x, out_v, off_x, off_v, in_x, in_v, halflife, dt): 75 | off_x_new, off_v_new = decay_spring_damper_exact(off_x, off_v, halflife, dt) 76 | out_x_new = in_x + off_x_new 77 | out_v_new = in_v + off_v_new 78 | return out_x_new, out_v_new 79 | -------------------------------------------------------------------------------- /manip/utils/smplx_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | import torch.nn.functional as F 4 | 5 | def batch_rodrigues( 6 | rot_vecs: Tensor, 7 | epsilon: float = 1e-8, 8 | ) -> Tensor: 9 | ''' Calculates the rotation matrices for a batch of rotation vectors 10 | Parameters 11 | ---------- 12 | rot_vecs: torch.tensor Nx3 13 | array of N axis-angle vectors 14 | Returns 15 | ------- 16 | R: torch.tensor Nx3x3 17 | The rotation matrices for the given axis-angle parameters 18 | ''' 19 | 20 | batch_size = rot_vecs.shape[0] 21 | device, dtype = rot_vecs.device, rot_vecs.dtype 22 | 23 | angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True) 24 | rot_dir = rot_vecs / angle 25 | 26 | cos = torch.unsqueeze(torch.cos(angle), dim=1) 27 | sin = torch.unsqueeze(torch.sin(angle), dim=1) 28 | 29 | # Bx1 arrays 30 | rx, ry, rz = torch.split(rot_dir, 1, dim=1) 31 | K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device) 32 | 33 | zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device) 34 | K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \ 35 | .view((batch_size, 3, 3)) 36 | 37 | ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0) 38 | rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K) 39 | return rot_mat 40 | 41 | class Struct(object): 42 | def __init__(self, **kwargs): 43 | for key, val in kwargs.items(): 44 | setattr(self, key, val) 45 | 46 | def transform_mat(R: Tensor, t: Tensor) -> Tensor: 47 | ''' Creates a batch of transformation matrices 48 | Args: 49 | - R: Bx3x3 array of a batch of rotation matrices 50 | - t: Bx3x1 array of a batch of translation vectors 51 | Returns: 52 | - T: Bx4x4 Transformation matrix 53 | ''' 54 | # No padding left or right, only add an extra row 55 | return torch.cat([F.pad(R, [0, 0, 0, 1]), 56 | F.pad(t, [0, 0, 0, 1], value=1)], dim=2) -------------------------------------------------------------------------------- /visualizer/tools/cfg_parser.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | import yaml 16 | import os 17 | 18 | class Config(dict): 19 | 20 | 21 | def __init__(self,default_cfg_path=None,**kwargs): 22 | 23 | default_cfg = {} 24 | if default_cfg_path is not None and os.path.exists(default_cfg_path): 25 | default_cfg = self.load_cfg(default_cfg_path) 26 | 27 | super(Config,self).__init__(**kwargs) 28 | 29 | default_cfg.update(self) 30 | self.update(default_cfg) 31 | self.default_cfg = default_cfg 32 | 33 | def load_cfg(self,load_path): 34 | with open(load_path, 'r') as infile: 35 | cfg = yaml.safe_load(infile) 36 | return cfg if cfg is not None else {} 37 | 38 | def write_cfg(self,write_path=None): 39 | 40 | if write_path is None: 41 | write_path = 'yaml_config.yaml' 42 | 43 | dump_dict = {k:v for k,v in self.items() if k!='default_cfg'} 44 | makepath(write_path, isfile=True) 45 | with open(write_path, 'w') as outfile: 46 | yaml.safe_dump(dump_dict, outfile, default_flow_style=False) 47 | 48 | def __getattr__(self, key): 49 | try: 50 | return self[key] 51 | except KeyError: 52 | raise AttributeError(key) 53 | 54 | __setattr__ = dict.__setitem__ 55 | __delattr__ = dict.__delitem__ 56 | 57 | def makepath(desired_path, isfile = False): 58 | ''' 59 | if the path does not exist make it 60 | :param desired_path: can be path to a file or a folder name 61 | :return: 62 | ''' 63 | import os 64 | if isfile: 65 | if not os.path.exists(os.path.dirname(desired_path)):os.makedirs(os.path.dirname(desired_path)) 66 | else: 67 | if not os.path.exists(desired_path): os.makedirs(desired_path) 68 | return desired_path 69 | 70 | if __name__ == '__main__': 71 | 72 | cfg = { 73 | 'intent': 'all', 74 | 'only_contact': True, 75 | 'save_body_verts': False, 76 | 'save_object_verts': False, 77 | 'save_contact': False, 78 | } 79 | 80 | cfg = Config(**cfg) 81 | cfg.write_cfg() -------------------------------------------------------------------------------- /manip/utils/visualize/tools/cfg_parser.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | import yaml 16 | import os 17 | 18 | class Config(dict): 19 | 20 | 21 | def __init__(self,default_cfg_path=None,**kwargs): 22 | 23 | default_cfg = {} 24 | if default_cfg_path is not None and os.path.exists(default_cfg_path): 25 | default_cfg = self.load_cfg(default_cfg_path) 26 | 27 | super(Config,self).__init__(**kwargs) 28 | 29 | default_cfg.update(self) 30 | self.update(default_cfg) 31 | self.default_cfg = default_cfg 32 | 33 | def load_cfg(self,load_path): 34 | with open(load_path, 'r') as infile: 35 | cfg = yaml.safe_load(infile) 36 | return cfg if cfg is not None else {} 37 | 38 | def write_cfg(self,write_path=None): 39 | 40 | if write_path is None: 41 | write_path = 'yaml_config.yaml' 42 | 43 | dump_dict = {k:v for k,v in self.items() if k!='default_cfg'} 44 | makepath(write_path, isfile=True) 45 | with open(write_path, 'w') as outfile: 46 | yaml.safe_dump(dump_dict, outfile, default_flow_style=False) 47 | 48 | def __getattr__(self, key): 49 | try: 50 | return self[key] 51 | except KeyError: 52 | raise AttributeError(key) 53 | 54 | __setattr__ = dict.__setitem__ 55 | __delattr__ = dict.__delitem__ 56 | 57 | def makepath(desired_path, isfile = False): 58 | ''' 59 | if the path does not exist make it 60 | :param desired_path: can be path to a file or a folder name 61 | :return: 62 | ''' 63 | import os 64 | if isfile: 65 | if not os.path.exists(os.path.dirname(desired_path)):os.makedirs(os.path.dirname(desired_path)) 66 | else: 67 | if not os.path.exists(desired_path): os.makedirs(desired_path) 68 | return desired_path 69 | 70 | if __name__ == '__main__': 71 | 72 | cfg = { 73 | 'intent': 'all', 74 | 'only_contact': True, 75 | 'save_body_verts': False, 76 | 'save_object_verts': False, 77 | 'save_contact': False, 78 | } 79 | 80 | cfg = Config(**cfg) 81 | cfg.write_cfg() -------------------------------------------------------------------------------- /grasp_generation/utils/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/PKU-EPIC/DexGraspNet 3 | """ 4 | 5 | from torch.utils.tensorboard.writer import SummaryWriter 6 | 7 | 8 | class Logger: 9 | def __init__(self, log_dir, thres_fc=0.3, thres_dis=0.005, thres_pen=0.02): 10 | """ 11 | Create a Logger to log tensorboard scalars 12 | 13 | Parameters 14 | ---------- 15 | log_dir: str 16 | directory for logs 17 | thres_fc: float 18 | E_fc threshold for success estimation 19 | thres_dis: float 20 | E_dis threshold for success estimation 21 | thres_pen: float 22 | E_pen threshold for data filtering 23 | """ 24 | self.writer = SummaryWriter(log_dir=log_dir) 25 | self.thres_fc = thres_fc 26 | self.thres_dis = thres_dis 27 | self.thres_pen = thres_pen 28 | 29 | def log(self, energy, E_fc, E_dis, E_pen, E_prior, E_spen, step, show=False): 30 | """ 31 | Log energy terms and estimate success rate using energy thresholds 32 | 33 | Parameters 34 | ---------- 35 | energy: torch.Tensor 36 | weighted sum of all terms 37 | E_fc: torch.Tensor 38 | E_dis: torch.Tensor 39 | E_pen: torch.Tensor 40 | E_prior: torch.Tensor 41 | E_spen: torch.Tensor 42 | step: int 43 | current iteration of optimization 44 | show: bool 45 | whether to print current energy terms to console 46 | """ 47 | success_fc = E_fc < self.thres_fc 48 | success_dis = E_dis < self.thres_dis 49 | success_pen = E_pen < self.thres_pen 50 | success = success_fc * success_dis * success_pen 51 | self.writer.add_scalar("Energy/energy", energy.mean(), step) 52 | self.writer.add_scalar("Energy/fc", E_fc.mean(), step) 53 | self.writer.add_scalar("Energy/dis", E_dis.mean(), step) 54 | self.writer.add_scalar("Energy/pen", E_pen.mean(), step) 55 | self.writer.add_scalar("Energy/prior", E_prior.mean(), step) 56 | self.writer.add_scalar("Energy/spen", E_spen.mean(), step) 57 | 58 | self.writer.add_scalar("Success/success", success.float().mean(), step) 59 | self.writer.add_scalar("Success/fc", success_fc.float().mean(), step) 60 | self.writer.add_scalar("Success/dis", success_dis.float().mean(), step) 61 | self.writer.add_scalar("Success/pen", success_pen.float().mean(), step) 62 | 63 | if show: 64 | print( 65 | f"Step %d energy: %f fc: %f dis: %f pen: %f" 66 | % (step, energy.mean(), E_fc.mean(), E_dis.mean(), E_pen.mean()) 67 | ) 68 | print( 69 | f"success: %f fc: %f dis: %f pen: %f" 70 | % ( 71 | success.float().mean(), 72 | success_fc.float().mean(), 73 | success_dis.float().mean(), 74 | success_pen.float().mean(), 75 | ) 76 | ) 77 | -------------------------------------------------------------------------------- /grasp_generation/utils/energy.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/PKU-EPIC/DexGraspNet 3 | """ 4 | 5 | import torch 6 | 7 | 8 | def cal_energy( 9 | hand_model, 10 | object_model, 11 | w_dis=100.0, 12 | w_pen=100.0, 13 | w_prior=0.5, 14 | w_spen=10.0, 15 | verbose=False, 16 | no_fc=False, 17 | correct_initial_pose=True, 18 | ): 19 | w_fc = 0.0 if no_fc else 1.0 20 | # E_dis 21 | batch_size, n_contact, _ = hand_model.contact_points.shape 22 | device = object_model.device 23 | distance, contact_normal = object_model.cal_distance(hand_model.contact_points) 24 | E_dis = torch.sum(distance.abs(), dim=-1, dtype=torch.float).to(device) 25 | 26 | # E_fc 27 | contact_normal = contact_normal.reshape(batch_size, 1, 3 * n_contact) 28 | transformation_matrix = torch.tensor( 29 | [ 30 | [0, 0, 0, 0, 0, -1, 0, 1, 0], 31 | [0, 0, 1, 0, 0, 0, -1, 0, 0], 32 | [0, -1, 0, 1, 0, 0, 0, 0, 0], 33 | ], 34 | dtype=torch.float, 35 | device=device, 36 | ) 37 | g = ( 38 | torch.cat( 39 | [ 40 | torch.eye(3, dtype=torch.float, device=device) 41 | .expand(batch_size, n_contact, 3, 3) 42 | .reshape(batch_size, 3 * n_contact, 3), 43 | (hand_model.contact_points @ transformation_matrix).view( 44 | batch_size, 3 * n_contact, 3 45 | ), 46 | ], 47 | dim=2, 48 | ) 49 | .float() 50 | .to(device) 51 | ) 52 | norm = torch.norm(contact_normal @ g, dim=[1, 2]) 53 | E_fc = norm * norm 54 | # E_pen 55 | object_scale = object_model.object_scale_tensor.flatten().unsqueeze(1).unsqueeze(2) 56 | object_surface_points = ( 57 | object_model.surface_points_tensor * object_scale 58 | ) # (n_objects * batch_size_each, num_samples, 3) 59 | distances = hand_model.cal_distance(object_surface_points) 60 | distances[distances <= 0] = 0 61 | E_pen = distances.sum(-1) 62 | 63 | # E_prior 64 | E_prior = torch.norm( 65 | (hand_model.hand_pose[:, 6:] - hand_model.pose_distrib[0]) 66 | / hand_model.pose_distrib[1], 67 | dim=-1, 68 | ) 69 | if correct_initial_pose: 70 | E_prior += ( 71 | torch.norm( 72 | (hand_model.hand_pose[:, :3] - hand_model.initial_translation), dim=-1 73 | ) 74 | * 20 75 | ) 76 | 77 | # E_spen 78 | E_spen = hand_model.self_penetration() 79 | 80 | if verbose: 81 | return ( 82 | w_fc * E_fc 83 | + w_dis * E_dis 84 | + w_pen * E_pen 85 | + w_prior * E_prior 86 | + w_spen * E_spen, 87 | E_fc, 88 | E_dis, 89 | E_pen, 90 | E_prior, 91 | E_spen, 92 | ) 93 | else: 94 | return ( 95 | w_fc * E_fc 96 | + w_dis * E_dis 97 | + w_pen * E_pen 98 | + w_prior * E_prior 99 | + w_spen * E_spen 100 | ) 101 | -------------------------------------------------------------------------------- /grasp_generation/utils/save.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | from typing import List, Tuple 4 | 5 | import numpy as np 6 | import trimesh 7 | 8 | 9 | def save_hand_meshes( 10 | hand_model, object_model, result_path, args, successes, data_lists, lefthand 11 | ) -> Tuple[List[List[str]], List[str], List[np.ndarray]]: 12 | all_hand_mesh_paths = [] 13 | all_object_mesh_paths = [] 14 | all_hand_poses = [] 15 | for i in range(len(args.object_code_list)): 16 | hand_poses = [] 17 | hand_mesh_paths = [] 18 | min_fail_hand_mesh_path = None 19 | min_fail_hand_pose = None 20 | min_fail_energy = float("inf") 21 | 22 | mesh_path = os.path.join(result_path, args.object_code_list[i] + f"_{i}") 23 | hand = "left_hand" if lefthand else "right_hand" 24 | os.makedirs(mesh_path, exist_ok=True) 25 | if os.path.exists(os.path.join(mesh_path, hand)): 26 | shutil.rmtree(os.path.join(mesh_path, hand)) 27 | os.makedirs(os.path.join(mesh_path, hand), exist_ok=True) 28 | os.makedirs(os.path.join(mesh_path, hand, "succ"), exist_ok=True) 29 | os.makedirs(os.path.join(mesh_path, hand, "fail"), exist_ok=True) 30 | 31 | obj_mesh = trimesh.Trimesh( 32 | vertices=object_model.object_mesh_list[i].vertices, 33 | faces=object_model.object_mesh_list[i].faces, 34 | ) 35 | obj_mesh.export(os.path.join(mesh_path, "object.obj")) 36 | for j in range(args.batch_size): 37 | idx = i * args.batch_size + j 38 | hand_mesh = trimesh.Trimesh( 39 | hand_model.vertices[idx].detach().cpu().numpy(), 40 | hand_model.hand_faces.detach().cpu().numpy(), 41 | ) 42 | if successes[idx]: 43 | path = os.path.join( 44 | mesh_path, hand, "succ", "{}_{}.obj".format(hand, j) 45 | ) 46 | hand_mesh.export(path) 47 | hand_mesh_paths.append(path) 48 | hand_poses.append(hand_model.hand_pose[idx].detach().cpu().numpy()) 49 | else: 50 | path = os.path.join( 51 | mesh_path, hand, "fail", "{}_{}.obj".format(hand, j) 52 | ) 53 | hand_mesh.export(path) 54 | if data_lists[i][j]["energy"] < min_fail_energy: 55 | min_fail_energy = data_lists[i][j]["energy"] 56 | min_fail_hand_mesh_path = path 57 | min_fail_hand_pose = ( 58 | hand_model.hand_pose[idx].detach().cpu().numpy() 59 | ) 60 | 61 | data_list = data_lists[i] 62 | np.save( 63 | os.path.join(mesh_path, hand, args.object_code_list[i] + ".npy"), 64 | data_list, 65 | allow_pickle=True, 66 | ) 67 | print("Save hand meshes to", os.path.join(mesh_path, hand)) 68 | 69 | if min_fail_hand_mesh_path is not None: 70 | hand_mesh_paths.append(min_fail_hand_mesh_path) 71 | hand_poses.append(min_fail_hand_pose) 72 | 73 | all_object_mesh_paths.append(os.path.join(mesh_path, "object.obj")) 74 | all_hand_mesh_paths.append(hand_mesh_paths) 75 | all_hand_poses.append(hand_poses) 76 | 77 | return all_hand_mesh_paths, all_object_mesh_paths, all_hand_poses 78 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # Human-Object Interaction from Human-Level Instructions (ICCV 2025) 3 | 4 | This is the official implementation for the ICCV 2025 [paper](https://arxiv.org/abs/2406.17840). For more information, please check the [project webpage](https://hoifhli.github.io/). 5 | 6 | ![CHOIS Teaser](teaser.png) 7 | 8 | ### Environment Setup 9 | > Note: This code was developed on Ubuntu 20.04 with Python 3.8, CUDA 12.4, and PyTorch 1.11.0. 10 | 11 | Clone the repo. 12 | ``` 13 | git clone https://github.com/zhenkirito123/hoifhli_release.git 14 | cd hoifhli_release/ 15 | ``` 16 | Create a virtual environment using Conda and activate the environment. 17 | ``` 18 | conda env create -f environment.yml 19 | conda activate hoifhli_env 20 | ``` 21 | Install TorchSDF based on the instruction from [here](https://github.com/PKU-EPIC/DexGraspNet/blob/bd1b13d7248729af117e1d46aaa6266b147a3c7b/grasp_generation/README.md?plain=1#L30). 22 | 23 | Install human_body_prior. 24 | ``` 25 | git clone https://github.com/nghorbani/human_body_prior.git 26 | pip install tqdm dotmap PyYAML omegaconf loguru 27 | cd human_body_prior/ 28 | python setup.py develop 29 | ``` 30 | Install BPS. 31 | ``` 32 | pip install git+https://github.com/otaheri/chamfer_distance 33 | pip install git+https://github.com/otaheri/bps_torch 34 | ``` 35 | Install smplx. 36 | ``` 37 | pip install "smplx[all]" 38 | ``` 39 | 40 | 41 | 42 | 43 | ### Prerequisites 44 | Please download [SMPL-X](https://smpl-x.is.tue.mpg.de/index.html) and put the model to ```data/smpl_all_models/```. The file structure should look like this: 45 | 46 | ``` 47 | data/ 48 | ├── smpl_all_models/ 49 | │ ├── smplx/ 50 | │ │ ├── SMPLX_FEMALE.npz 51 | │ │ ├── SMPLX_MALE.npz 52 | │ │ ├── SMPLX_NEUTRAL.npz 53 | │ │ ├── SMPLX_FEMALE.pkl 54 | │ │ ├── SMPLX_MALE.pkl 55 | │ │ ├── SMPLX_NEUTRAL.pkl 56 | ``` 57 | 58 | Then run the following script to download the data. 59 | ``` 60 | bash scripts/download_data.sh 61 | ``` 62 | 63 | 64 | ### Sampling Long Sequences 65 | Run the following script to sample long sequences. 66 | ``` 67 | bash scripts/sample.sh 68 | ``` 69 | By default, it will generate interaction motion with a box. You can change the object by modifying the ```--test_object_name``` argument. 70 | 71 | ### Training 72 | You need to download the training data first. Run the following script to download the data. (**Note**: This script will delete `/data/processed_data` folder before re-downloading the training data. If you have any custom data stored there, make sure to back it up beforehand.) 73 | ``` 74 | bash scripts/download_training_data.sh 75 | ``` 76 | 77 | #### Training CoarseNet 78 | ``` 79 | bash scripts/train/train_coarsenet.sh 80 | ``` 81 | 82 | #### Training RefineNet 83 | ``` 84 | bash scripts/train/train_refinenet.sh 85 | ``` 86 | 87 | ### Physics Tracking 88 | Please refer to [here](./physics_tracking/README.md). 89 | 90 | 91 | ### Citation 92 | ``` 93 | @inproceedings{hoifhli, 94 | title={Human-Object Interaction from Human-Level Instructions}, 95 | author={Wu, Zhen and Li, Jiaman and Xu, Pei and Liu, C. Karen}, 96 | booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, 97 | month={October}, 98 | year={2025} 99 | } 100 | ``` 101 | 102 | ### Related Repos 103 | We adapted some code from other repos in data processing, learning, evaluation, etc. We would like to thank the authors and contributors of these repositories for their valuable work and resources. 104 | ``` 105 | https://github.com/lijiaman/chois_release 106 | https://github.com/otaheri/GRAB 107 | https://github.com/nghorbani/human_body_prior 108 | https://github.com/PKU-EPIC/DexGraspNet 109 | ``` -------------------------------------------------------------------------------- /manip/utils/model_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import pytorch3d.transforms as transforms 3 | 4 | 5 | def wxyz_to_xyzw(input_quat): 6 | # input_quat: 1 X w X 22 X 4 7 | w = input_quat[..., 0:1] 8 | x = input_quat[..., 1:2] 9 | y = input_quat[..., 2:3] 10 | z = input_quat[..., 3:4] 11 | 12 | out_quat = torch.cat((x, y, z, w), dim=-1) # 1 X w X 22 X 4 13 | 14 | return out_quat 15 | 16 | 17 | def xyzw_to_wxyz(input_quat): 18 | # input_quat: 1 X w X 22 X 4 19 | x = input_quat[..., 0:1] 20 | y = input_quat[..., 1:2] 21 | z = input_quat[..., 2:3] 22 | w = input_quat[..., 3:4] 23 | 24 | out_quat = torch.cat((w, x, y, z), dim=-1) # 1 X w X 22 X 4 25 | 26 | return out_quat 27 | 28 | 29 | def apply_rotation_to_data(ds, trans2joint, cano_rot_mat, new_obj_rot_mat, curr_x): 30 | # cano_rot_mat:BS X 3 X 3, convert from the coodinate frame which canonicalize the first frame of a sequence to 31 | # the frame that canonicalize the first frame of a window. 32 | # trans2joint: BS X 3 33 | # new_obj_rot_mat: BS X 10(overlapped -length) X 3 X 3 34 | # curr_x: BS X window_size X D 35 | # This function is to convert window data to sequence data. 36 | bs, timesteps, _ = curr_x.shape 37 | 38 | pred_obj_normalized_com_pos = curr_x[:, :, :3] # BS X window_size X 3 39 | pred_obj_com_pos = ds.de_normalize_obj_pos_min_max(pred_obj_normalized_com_pos) 40 | pred_obj_rel_rot_mat = curr_x[:, :, 3:12].reshape(bs, timesteps, 3, 3) # BS X window_size X 3 X 3, relative rotation wrt current window's first frames's cano. 41 | pred_obj_rot_mat = ds.rel_rot_to_seq(pred_obj_rel_rot_mat, new_obj_rot_mat) # BS X window_size X 3 X 3 42 | pred_human_normalized_jpos = curr_x[:, :, 12:12+24*3] # BS X window_size X (24*3) 43 | pred_human_jpos = ds.de_normalize_jpos_min_max(pred_human_normalized_jpos.reshape(bs, timesteps, 24, 3)) # BS X window_size X 24 X 3 44 | pred_human_rot_6d = curr_x[:, :, 12+24*3:] # BS X window_size X (22*6) 45 | 46 | pred_human_rot_mat = transforms.rotation_6d_to_matrix(pred_human_rot_6d.reshape(bs, timesteps, 22, 6)) # BS X T X 22 X 3 X 3 47 | 48 | converted_obj_com_pos = torch.matmul( 49 | cano_rot_mat[:, None, :, :].repeat(1, timesteps, 1, 1).transpose(2, 3), 50 | pred_obj_com_pos[:, :, :, None] 51 | ).squeeze(-1) # BS X window_size X 3 52 | 53 | converted_obj_rot_mat = torch.matmul( 54 | cano_rot_mat[:, None, :, :].repeat(1, timesteps,1, 1).transpose(2, 3), 55 | pred_obj_rot_mat 56 | ) # BS X window_size X 3 X 3 57 | 58 | converted_human_jpos = torch.matmul( 59 | cano_rot_mat[:, None, None, :, :].repeat(1, timesteps, 24, 1, 1).transpose(3, 4), 60 | pred_human_jpos[:, :, :, :, None] 61 | ).squeeze(-1) # BS X T X 24 X 3 62 | converted_rot_mat = torch.matmul( 63 | cano_rot_mat[:, None, None, :, :].repeat(1, timesteps, 22, 1, 1).transpose(3, 4), 64 | pred_human_rot_mat 65 | ) # BS X T X 22 X 3 X 3 66 | 67 | converted_rot_6d = transforms.matrix_to_rotation_6d(converted_rot_mat) 68 | 69 | return converted_obj_com_pos, converted_obj_rot_mat, converted_human_jpos, converted_rot_6d 70 | 71 | 72 | def calculate_obj_kpts_in_wrist( 73 | global_wrist_jpos, 74 | global_joint_rot_mat, 75 | seq_obj_kpts 76 | ): 77 | """ 78 | global_wrist_jpos: BS X T X 2 X 3 79 | global_joint_rot_mat: BS X T X 2 X 3 X 3 80 | seq_obj_kpts: BS X T X K X 3 81 | 82 | pred_seq_obj_kpts_in_hand: BS X T X 2 X K X 3 83 | """ 84 | K = seq_obj_kpts.shape[2] 85 | 86 | global_wrist_jpos_expand = global_wrist_jpos.unsqueeze(3).repeat(1, 1, 1, K, 1) # BS X T X 2 X K X 3 87 | seq_obj_kpts_expand = seq_obj_kpts.unsqueeze(2).repeat(1, 1, 2, 1, 1) # BS X T X 2 X K X 3 88 | 89 | # (inv(R) * (p - t)^T)^T = (R^T * (p - t)^T)^T = (p - t) * R 90 | seq_obj_kpts_in_hand = torch.matmul( 91 | seq_obj_kpts_expand - global_wrist_jpos_expand, global_joint_rot_mat 92 | ) 93 | 94 | return seq_obj_kpts_in_hand 95 | -------------------------------------------------------------------------------- /physics_tracking/config_tracking.py: -------------------------------------------------------------------------------- 1 | env_cls = "TrackingHumanoid" 2 | # env_cls = "ICCGANHumanoidDemo" 3 | 4 | env_params = dict( 5 | fps = 30, 6 | random_init = True, 7 | episode_length = 500, 8 | character_model = ["assets/humanoid.xml"], # first character must be the humanoid character 9 | 10 | contactable_links = dict( # link_name: contact_height_threshold (default: 0.2) 11 | L_Ankle = -1000, R_Ankle = -1000, # allow to contact always 12 | L_Toe_tip = -1000, R_Toe_tip = -1000, 13 | ), 14 | key_link_weights_orient = dict( 15 | Pelvis = 1.0, 16 | L_Hip = 0.5, L_Knee = 0.3, L_Ankle = 0.2, #L_Toe = 0.1, 17 | R_Hip = 0.5, R_Knee = 0.3, R_Ankle = 0.2, #R_Toe = 0.1, 18 | Torso = 0.2, Spine = 0.2, Chest = 0.2, Neck = 0.2, Head = 0.2, 19 | L_Thorax = 0.1, L_Shoulder = 0.2, L_Elbow = 0.2, L_Wrist = 0.3, 20 | L_Index1 = 0, L_Index2 = 0, L_Index3 = 0, L_Index4 = 0, 21 | L_Middle1 = 0, L_Middle2 = 0, L_Middle3 = 0, L_Middle4 = 0, 22 | L_Pinky1 = 0, L_Pinky2 = 0, L_Pinky3 = 0, L_Pinky4 = 0, 23 | L_Ring1 = 0, L_Ring2 = 0, L_Ring3 = 0, L_Ring4 = 0, 24 | L_Thumb1 = 0, L_Thumb2 = 0, L_Thumb3 = 0, L_Thumb4 = 0, 25 | R_Thorax = 0.1, R_Shoulder = 0.2, R_Elbow = 0.2, R_Wrist = 0.3, 26 | R_Index1 = 0, R_Index2 = 0, R_Index3 = 0, R_Index4 = 0, 27 | R_Middle1 = 0, R_Middle2 = 0, R_Middle3 = 0, R_Middle4 = 0, 28 | R_Pinky1 = 0, R_Pinky2 = 0, R_Pinky3 = 0, R_Pinky4 = 0, 29 | R_Ring1 = 0, R_Ring2 = 0, R_Ring3 = 0, R_Ring4 = 0, 30 | R_Thumb1 = 0, R_Thumb2 = 0, R_Thumb3 = 0, R_Thumb4 = 0, 31 | object = 1 32 | ), 33 | key_link_weights_pos = dict( 34 | Pelvis = 1, 35 | L_Hip = 0, L_Knee = 0, L_Ankle = 0.1, #L_Toe = 0, 36 | R_Hip = 0, R_Knee = 0, R_Ankle = 0.1, #R_Toe = 0, 37 | Torso = 0, Spine = 0, Chest = 0, Neck = 0, Head = 0, 38 | L_Thorax = 0, L_Shoulder = 0, L_Elbow = 0, L_Wrist = 0.3, 39 | L_Index1 = 0, L_Index2 = 0, L_Index3 = 0, L_Index4 = 0, 40 | L_Middle1 = 0, L_Middle2 = 0, L_Middle3 = 0, L_Middle4 = 0, 41 | L_Pinky1 = 0, L_Pinky2 = 0, L_Pinky3 = 0, L_Pinky4 = 0, 42 | L_Ring1 = 0, L_Ring2 = 0, L_Ring3 = 0, L_Ring4 = 0, 43 | L_Thumb1 = 0, L_Thumb2 = 0, L_Thumb3 = 0, L_Thumb4 = 0, 44 | R_Thorax = 0, R_Shoulder = 0, R_Elbow = 0, R_Wrist = 0.3, 45 | R_Index1 = 0, R_Index2 = 0, R_Index3 = 0, R_Index4 = 0, 46 | R_Middle1 = 0, R_Middle2 = 0, R_Middle3 = 0, R_Middle4 = 0, 47 | R_Pinky1 = 0, R_Pinky2 = 0, R_Pinky3 = 0, R_Pinky4 = 0, 48 | R_Ring1 = 0, R_Ring2 = 0, R_Ring3 = 0, R_Ring4 = 0, 49 | R_Thumb1 = 0, R_Thumb2 = 0, R_Thumb3 = 0, R_Thumb4 = 0, 50 | object =1 51 | ), 52 | key_link_weights_pos_related = dict( 53 | Pelvis = 0, 54 | L_Hip = 0, L_Knee = 0, L_Ankle = 0, #L_Toe = 0, 55 | R_Hip = 0, R_Knee = 0, R_Ankle = 0, # R_Toe = 0, 56 | Torso = 0, Spine = 0, Chest = 0, Neck = 0, Head = 0, 57 | L_Thorax = 0, L_Shoulder = 0, L_Elbow = 0, L_Wrist = 1, 58 | L_Index1= .3, L_Index2= .3, L_Index3= .3, L_Index4= .3, 59 | L_Middle1= .3, L_Middle2= .3, L_Middle3= .3, L_Middle4= .3, 60 | L_Pinky1= .3, L_Pinky2= .3, L_Pinky3= .3, L_Pinky4= .3, 61 | L_Ring1= .3, L_Ring2= .3, L_Ring3= .3, L_Ring4= .3, 62 | L_Thumb1= .3, L_Thumb2= .3, L_Thumb3= .3, L_Thumb4= .3, 63 | R_Thorax = 0, R_Shoulder = 0, R_Elbow = 0, R_Wrist = 1, 64 | R_Index1= .3, R_Index2= .3, R_Index3= .3, R_Index4= .3, 65 | R_Middle1= .3, R_Middle2= .3, R_Middle3= .3, R_Middle4= .3, 66 | R_Pinky1= .3, R_Pinky2= .3, R_Pinky3= .3, R_Pinky4= .3, 67 | R_Ring1= .3, R_Ring2= .3, R_Ring3= .3, R_Ring4= .3, 68 | R_Thumb1= .3, R_Thumb2= .3, R_Thumb3= .3, R_Thumb4= .3 69 | ), 70 | key_link_weights_acc_penalty = dict( 71 | L_Ankle=1, L_Toe_tip=1, 72 | R_Ankle=1, R_Toe_tip=1, 73 | L_Wrist=1, R_Wrist=1 74 | ) 75 | ) 76 | 77 | training_params = dict( 78 | max_epochs = 1000000, 79 | save_interval = 50000 80 | ) 81 | 82 | discriminators = {} 83 | -------------------------------------------------------------------------------- /visualizer/tools/objectmodel.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | 16 | import numpy as np 17 | 18 | import torch 19 | import torch.nn as nn 20 | from smplx.lbs import batch_rodrigues 21 | from collections import namedtuple 22 | 23 | model_output = namedtuple('output', ['vertices', 'global_orient', 'transl']) 24 | 25 | class ObjectModel(nn.Module): 26 | 27 | def __init__(self, 28 | v_template, 29 | batch_size=1, 30 | dtype=torch.float32): 31 | ''' 3D rigid object model 32 | 33 | Parameters 34 | ---------- 35 | v_template: np.array Vx3, dtype = np.float32 36 | The vertices of the object 37 | batch_size: int, N, optional 38 | The batch size used for creating the model variables 39 | 40 | dtype: torch.dtype 41 | The data type for the created variables 42 | ''' 43 | 44 | super(ObjectModel, self).__init__() 45 | 46 | 47 | self.dtype = dtype 48 | 49 | # Mean template vertices 50 | v_template = np.repeat(v_template[np.newaxis], batch_size, axis=0) 51 | self.register_buffer('v_template', torch.tensor(v_template, dtype=dtype)) 52 | 53 | transl = torch.tensor(np.zeros((batch_size, 3)), dtype=dtype, requires_grad=True) 54 | self.register_parameter('transl', nn.Parameter(transl, requires_grad=True)) 55 | 56 | global_orient = torch.tensor(np.zeros((batch_size, 3)), dtype=dtype, requires_grad=True) 57 | self.register_parameter('global_orient', nn.Parameter(global_orient, requires_grad=True)) 58 | 59 | self.batch_size = batch_size 60 | 61 | 62 | def forward(self, global_orient=None, transl=None, v_template=None, **kwargs): 63 | 64 | ''' Forward pass for the object model 65 | 66 | Parameters 67 | ---------- 68 | global_orient: torch.tensor, optional, shape Bx3 69 | If given, ignore the member variable and use it as the global 70 | rotation of the body. Useful if someone wishes to predicts this 71 | with an external model. (default=None) 72 | 73 | transl: torch.tensor, optional, shape Bx3 74 | If given, ignore the member variable `transl` and use it 75 | instead. For example, it can used if the translation 76 | `transl` is predicted from some external model. 77 | (default=None) 78 | v_template: torch.tensor, optional, shape BxVx3 79 | The new object vertices to overwrite the default vertices 80 | 81 | Returns 82 | ------- 83 | output: ModelOutput 84 | A named tuple of type `ModelOutput` 85 | ''' 86 | if global_orient is None: 87 | global_orient = self.global_orient 88 | if transl is None: 89 | transl = self.transl 90 | if v_template is None: 91 | v_template = self.v_template 92 | 93 | rot_mats = batch_rodrigues(global_orient.view(-1, 3)).view([self.batch_size, 3, 3]) 94 | 95 | vertices = torch.matmul(v_template, rot_mats) + transl.unsqueeze(dim=1) 96 | 97 | output = model_output(vertices=vertices, 98 | global_orient=global_orient, 99 | transl=transl) 100 | 101 | return output 102 | 103 | -------------------------------------------------------------------------------- /manip/utils/visualize/tools/objectmodel.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | 16 | import numpy as np 17 | 18 | import torch 19 | import torch.nn as nn 20 | from smplx.lbs import batch_rodrigues 21 | from collections import namedtuple 22 | 23 | model_output = namedtuple('output', ['vertices', 'global_orient', 'transl']) 24 | 25 | class ObjectModel(nn.Module): 26 | 27 | def __init__(self, 28 | v_template, 29 | batch_size=1, 30 | dtype=torch.float32): 31 | ''' 3D rigid object model 32 | 33 | Parameters 34 | ---------- 35 | v_template: np.array Vx3, dtype = np.float32 36 | The vertices of the object 37 | batch_size: int, N, optional 38 | The batch size used for creating the model variables 39 | 40 | dtype: torch.dtype 41 | The data type for the created variables 42 | ''' 43 | 44 | super(ObjectModel, self).__init__() 45 | 46 | 47 | self.dtype = dtype 48 | 49 | # Mean template vertices 50 | v_template = np.repeat(v_template[np.newaxis], batch_size, axis=0) 51 | self.register_buffer('v_template', torch.tensor(v_template, dtype=dtype)) 52 | 53 | transl = torch.tensor(np.zeros((batch_size, 3)), dtype=dtype, requires_grad=True) 54 | self.register_parameter('transl', nn.Parameter(transl, requires_grad=True)) 55 | 56 | global_orient = torch.tensor(np.zeros((batch_size, 3)), dtype=dtype, requires_grad=True) 57 | self.register_parameter('global_orient', nn.Parameter(global_orient, requires_grad=True)) 58 | 59 | self.batch_size = batch_size 60 | 61 | 62 | def forward(self, global_orient=None, transl=None, v_template=None, **kwargs): 63 | 64 | ''' Forward pass for the object model 65 | 66 | Parameters 67 | ---------- 68 | global_orient: torch.tensor, optional, shape Bx3 69 | If given, ignore the member variable and use it as the global 70 | rotation of the body. Useful if someone wishes to predicts this 71 | with an external model. (default=None) 72 | 73 | transl: torch.tensor, optional, shape Bx3 74 | If given, ignore the member variable `transl` and use it 75 | instead. For example, it can used if the translation 76 | `transl` is predicted from some external model. 77 | (default=None) 78 | v_template: torch.tensor, optional, shape BxVx3 79 | The new object vertices to overwrite the default vertices 80 | 81 | Returns 82 | ------- 83 | output: ModelOutput 84 | A named tuple of type `ModelOutput` 85 | ''' 86 | 87 | 88 | if global_orient is None: 89 | global_orient = self.global_orient 90 | if transl is None: 91 | transl = self.transl 92 | if v_template is None: 93 | v_template = self.v_template 94 | 95 | rot_mats = batch_rodrigues(global_orient.view(-1, 3)).view([self.batch_size, 3, 3]) 96 | 97 | vertices = torch.matmul(v_template, rot_mats) + transl.unsqueeze(dim=1) 98 | 99 | output = model_output(vertices=vertices, 100 | global_orient=global_orient, 101 | transl=transl) 102 | 103 | return output 104 | 105 | -------------------------------------------------------------------------------- /manip/vis/blender_vis_mesh_motion.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | 4 | import imageio 5 | import numpy as np 6 | import trimesh 7 | 8 | 9 | def concat_multiple_videos(input_files, output_file): 10 | # List of input files 11 | # input_files = ['video1.mp4', 'video2.mp4'] 12 | 13 | # Output file 14 | # output_file = 'output.mp4' 15 | 16 | # Step 1: Convert each video to a consistent FPS (e.g., 30 fps) and save to a temp file. 17 | temp_files = [] 18 | target_fps = 30 19 | 20 | temp_folder = output_file.replace(".mp4", "_tmp") 21 | if not os.path.exists(temp_folder): 22 | os.makedirs(temp_folder) 23 | for i, file in enumerate(input_files): 24 | temp_filename = os.path.join(temp_folder, str(i) + ".mp4") 25 | 26 | temp_files.append(temp_filename) 27 | reader = imageio.get_reader(file) 28 | writer = imageio.get_writer(temp_filename, fps=target_fps) 29 | 30 | for frame in reader: 31 | writer.append_data(frame) 32 | writer.close() 33 | 34 | # Step 2: Concatenate the temporary files. 35 | with imageio.get_writer(output_file, fps=target_fps) as final_writer: 36 | for temp_file in temp_files: 37 | reader = imageio.get_reader(temp_file) 38 | for frame in reader: 39 | final_writer.append_data(frame) 40 | 41 | # Step 3: Cleanup temp files. 42 | # for temp_file in temp_files: 43 | # os.remove(temp_file) 44 | # # shutil.rmtree(temp_folder) 45 | # shutil.rmtree(temp_folder) 46 | 47 | 48 | def images_to_video(img_folder, output_vid_file): 49 | os.makedirs(img_folder, exist_ok=True) 50 | 51 | command = [ 52 | "ffmpeg", 53 | "-r", 54 | "30", 55 | "-y", 56 | "-threads", 57 | "16", 58 | "-i", 59 | f"{img_folder}/%05d.png", 60 | "-profile:v", 61 | "baseline", 62 | "-level", 63 | "3.0", 64 | "-c:v", 65 | "libx264", 66 | "-pix_fmt", 67 | "yuv420p", 68 | "-an", 69 | "-v", 70 | "error", 71 | output_vid_file, 72 | ] 73 | 74 | # command = [ 75 | # 'ffmpeg', '-r', '30', '-y', '-threads', '16', '-i', f'{img_folder}/%05d.png', output_vid_file, 76 | # ] 77 | 78 | print(f'Running "{" ".join(command)}"') 79 | subprocess.call(command) 80 | 81 | 82 | def images_to_video_w_imageio(img_folder, output_vid_file, fps=30): 83 | img_files = os.listdir(img_folder) 84 | img_files.sort() 85 | im_arr = [] 86 | for img_name in img_files: 87 | img_path = os.path.join(img_folder, img_name) 88 | im = imageio.imread(img_path) 89 | im_arr.append(im) 90 | 91 | im_arr = np.asarray(im_arr) 92 | imageio.mimwrite(output_vid_file, im_arr, fps=fps, quality=8) 93 | 94 | 95 | def save_verts_faces_to_mesh_file( 96 | mesh_verts, mesh_faces, save_mesh_folder, save_gt=False 97 | ): 98 | # mesh_verts: T X Nv X 3 99 | # mesh_faces: Nf X 3 100 | if not os.path.exists(save_mesh_folder): 101 | os.makedirs(save_mesh_folder) 102 | 103 | num_meshes = mesh_verts.shape[0] 104 | for idx in range(num_meshes): 105 | mesh = trimesh.Trimesh(vertices=mesh_verts[idx], faces=mesh_faces) 106 | if save_gt: 107 | curr_mesh_path = os.path.join(save_mesh_folder, "%05d" % (idx) + "_gt.ply") 108 | else: 109 | curr_mesh_path = os.path.join(save_mesh_folder, "%05d" % (idx) + ".ply") 110 | mesh.export(curr_mesh_path) 111 | 112 | 113 | def save_verts_faces_to_mesh_file_w_object( 114 | mesh_verts, mesh_faces, obj_verts, obj_faces, save_mesh_folder 115 | ): 116 | # mesh_verts: T X Nv X 3 117 | # mesh_faces: Nf X 3 118 | if not os.path.exists(save_mesh_folder): 119 | os.makedirs(save_mesh_folder) 120 | 121 | num_meshes = mesh_verts.shape[0] 122 | for idx in range(num_meshes): 123 | mesh = trimesh.Trimesh(vertices=mesh_verts[idx], faces=mesh_faces) 124 | curr_mesh_path = os.path.join(save_mesh_folder, "%05d" % (idx) + ".ply") 125 | mesh.export(curr_mesh_path) 126 | 127 | obj_mesh = trimesh.Trimesh(vertices=obj_verts[idx], faces=obj_faces) 128 | curr_obj_mesh_path = os.path.join( 129 | save_mesh_folder, "%05d" % (idx) + "_object.ply" 130 | ) 131 | obj_mesh.export(curr_obj_mesh_path) 132 | -------------------------------------------------------------------------------- /physics_tracking/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | import torch 3 | 4 | 5 | @torch.jit.script 6 | def rotatepoint(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor: 7 | # q_v = [v[0], v[1], v[2], 0] 8 | # return quatmultiply(quatmultiply(q, q_v), quatconj(q))[:-1] 9 | # 10 | # https://fgiesen.wordpress.com/2019/02/09/rotating-a-single-vector-using-a-quaternion/ 11 | q_r = q[...,3:4] 12 | q_xyz = q[...,:3] 13 | t = 2*torch.linalg.cross(q_xyz, v) 14 | return v + q_r * t + torch.linalg.cross(q_xyz, t) 15 | 16 | @torch.jit.script 17 | def heading_zup(q: torch.Tensor) -> torch.Tensor: 18 | ref_dir = torch.zeros_like(q[...,:3]) 19 | ref_dir[..., 0] = 1 20 | ref_dir = rotatepoint(q, ref_dir) 21 | return torch.atan2(ref_dir[...,1], ref_dir[...,0]) 22 | 23 | @torch.jit.script 24 | def heading_yup(q: torch.Tensor) -> torch.Tensor: 25 | ref_dir = torch.zeros_like(q[...,:3]) 26 | ref_dir[..., 0] = 1 27 | ref_dir = rotatepoint(q, ref_dir) 28 | return torch.atan2(-ref_dir[...,2], ref_dir[...,0]) 29 | 30 | @torch.jit.script 31 | def quatnormalize(q: torch.Tensor) -> torch.Tensor: 32 | q = (1-2*(q[...,3:4]<0).to(q.dtype))*q 33 | return q / q.norm(p=2, dim=-1, keepdim=True) 34 | 35 | @torch.jit.script 36 | def quatmultiply(q0: torch.Tensor, q1: torch.Tensor): 37 | x0, y0, z0, w0 = torch.unbind(q0, -1) 38 | x1, y1, z1, w1 = torch.unbind(q1, -1) 39 | w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1 40 | x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1 41 | y = w0 * y1 - x0 * z1 + y0 * w1 + z0 * x1 42 | z = w0 * z1 + x0 * y1 - y0 * x1 + z0 * w1 43 | return quatnormalize(torch.stack((x, y, z, w), -1)) 44 | 45 | @torch.jit.script 46 | def quatconj(q: torch.Tensor): 47 | return torch.cat((-q[...,:3], q[...,-1:]), dim=-1) 48 | 49 | @torch.jit.script 50 | def axang2quat(axis: torch.Tensor, angle: torch.Tensor) -> torch.Tensor: 51 | # axis: n x 3 52 | # angle: n 53 | theta = (angle / 2).unsqueeze(-1) 54 | axis = axis / (axis.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-9)) 55 | xyz = axis * torch.sin(theta) 56 | w = torch.cos(theta) 57 | return quatnormalize(torch.cat((xyz, w), -1)) 58 | 59 | @torch.jit.script 60 | def quatdiff_normalized(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: 61 | # quaternion representation of the rotation from unit vector a to b 62 | # need to check if a == -b 63 | # if a == -b: q = *a, 0 # 180 degree around any axis 64 | w = (a*b).sum(-1).add_(1) 65 | xyz = torch.linalg.cross(a, b) 66 | q = torch.cat((xyz, w.unsqueeze_(-1)), -1) 67 | return quatnormalize(q) 68 | 69 | @torch.jit.script 70 | def wrap2pi(x: torch.Tensor) -> torch.Tensor: 71 | return torch.atan2(torch.sin(x), torch.cos(x)) 72 | 73 | @torch.jit.script 74 | def quat2axang(q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 75 | w = q[..., 3] 76 | 77 | sin = torch.sqrt(1 - w * w) 78 | mask = sin > 1e-5 79 | 80 | angle = 2 * torch.acos(w) 81 | angle = wrap2pi(angle) 82 | axis = q[..., 0:3] / sin.unsqueeze_(-1) 83 | 84 | z_axis = torch.zeros_like(axis) 85 | z_axis[..., -1] = 1 86 | 87 | angle = torch.where(mask, angle, z_axis[...,0]) 88 | axis = torch.where(mask.unsqueeze_(-1), axis, z_axis) 89 | return axis, angle 90 | 91 | @torch.jit.script 92 | def quat2expmap(q: torch.Tensor) -> torch.Tensor: 93 | ax, ang = quat2axang(q) 94 | return ang.unsqueeze(-1)*ax 95 | 96 | @torch.jit.script 97 | def slerp(q0, q1, frac): 98 | c = q0[..., 3]*q1[..., 3] + q0[..., 0]*q1[..., 0] + \ 99 | q0[..., 1]*q1[..., 1] + q0[..., 2]*q1[..., 2] 100 | q1 = torch.where(c.unsqueeze_(-1) < 0, -q1, q1) 101 | 102 | c = c.abs_() 103 | s = torch.sqrt(1.0 - c*c) 104 | t = torch.acos(c); 105 | 106 | c1 = torch.sin((1-frac)*t) / s 107 | c2 = torch.sin(frac*t) / s 108 | 109 | x = c1*q0[..., 0:1] + c2*q1[..., 0:1] 110 | y = c1*q0[..., 1:2] + c2*q1[..., 1:2] 111 | z = c1*q0[..., 2:3] + c2*q1[..., 2:3] 112 | w = c1*q0[..., 3:4] + c2*q1[..., 3:4] 113 | 114 | q = torch.cat((x, y, z, w), dim=-1) 115 | q = torch.where(s < 0.001, 0.5*q0+0.5*q1, q) 116 | q = torch.where(c >= 1, q0, q) 117 | return q 118 | 119 | 120 | @torch.jit.script 121 | def expmap2quat(rotvec: torch.Tensor) -> torch.Tensor: 122 | angle = rotvec.norm(p=2, dim=-1, keepdim=True) 123 | axis = rotvec / angle 124 | ref_angle = torch.zeros_like(angle) 125 | ref_axis = torch.zeros_like(rotvec) 126 | ref_axis[..., 0] = 1 127 | axis = torch.where(angle < 1e-5, ref_axis, axis) 128 | angle = torch.where(angle < 1e-5, ref_angle, angle) 129 | 130 | theta = angle / 2 131 | xyz = axis * torch.sin(theta) 132 | w = torch.cos(theta) 133 | return quatnormalize(torch.cat((xyz, w), -1)) -------------------------------------------------------------------------------- /grasp_generation/utils/optimizer.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/PKU-EPIC/DexGraspNet 3 | """ 4 | 5 | import random 6 | 7 | import torch 8 | 9 | 10 | class Annealing: 11 | def __init__( 12 | self, 13 | hand_model, 14 | switch_possibility=0.5, 15 | starting_temperature=18, 16 | temperature_decay=0.95, 17 | annealing_period=30, 18 | step_size=0.005, 19 | stepsize_period=50, 20 | mu=0.98, 21 | device="cpu", 22 | ): 23 | """ 24 | Create a optimizer 25 | 26 | Use random resampling to update contact point indices 27 | 28 | Use RMSProp to update translation, rotation, and joint angles, use step size decay 29 | 30 | Use Annealing to accept / reject parameter updates 31 | 32 | Parameters 33 | ---------- 34 | hand_model: hand_model.HandModel 35 | switch_possibility: float 36 | possibility to resample each contact point index each step 37 | starting_temperature: float 38 | temperature_decay: float 39 | temperature decay rate and step size decay rate 40 | annealing_period: int 41 | step_size: float 42 | stepsize_period: int 43 | mu: float 44 | `1 - decay_rate` of RMSProp 45 | """ 46 | 47 | self.hand_model = hand_model 48 | self.device = device 49 | self.switch_possibility = switch_possibility 50 | self.starting_temperature = torch.tensor( 51 | starting_temperature, dtype=torch.float, device=device 52 | ) 53 | self.temperature_decay = torch.tensor( 54 | temperature_decay, dtype=torch.float, device=device 55 | ) 56 | self.annealing_period = torch.tensor( 57 | annealing_period, dtype=torch.long, device=device 58 | ) 59 | self.step_size = torch.tensor(step_size, dtype=torch.float, device=device) 60 | self.step_size_period = torch.tensor( 61 | stepsize_period, dtype=torch.long, device=device 62 | ) 63 | self.mu = torch.tensor(mu, dtype=torch.float, device=device) 64 | self.step = 0 65 | 66 | self.old_hand_pose = None 67 | self.old_contact_point_indices = None 68 | self.old_global_transformation = None 69 | self.old_global_rotation = None 70 | self.old_current_status = None 71 | self.old_contact_points = None 72 | self.old_grad_hand_pose = None 73 | self.ema_grad_hand_pose = torch.zeros(51, dtype=torch.float, device=device) 74 | 75 | def try_step(self): 76 | """ 77 | Try to update translation, rotation, joint angles, and contact point indices 78 | 79 | Returns 80 | ------- 81 | s: torch.Tensor 82 | current step size 83 | """ 84 | 85 | s = self.step_size * self.temperature_decay ** torch.div( 86 | self.step, self.step_size_period, rounding_mode="floor" 87 | ) 88 | step_size = ( 89 | torch.zeros( 90 | *self.hand_model.hand_pose.shape, dtype=torch.float, device=self.device 91 | ) 92 | + s 93 | ) 94 | 95 | self.ema_grad_hand_pose = ( 96 | self.mu * (self.hand_model.hand_pose.grad**2).mean(0) 97 | + (1 - self.mu) * self.ema_grad_hand_pose 98 | ) 99 | 100 | hand_pose = ( 101 | self.hand_model.hand_pose 102 | - step_size 103 | * self.hand_model.hand_pose.grad 104 | / (torch.sqrt(self.ema_grad_hand_pose) + 1e-6) 105 | ) 106 | batch_size, n_contact = self.hand_model.contact_point_indices.shape 107 | switch_mask = ( 108 | torch.rand(batch_size, n_contact, dtype=torch.float, device=self.device) 109 | < self.switch_possibility 110 | ) 111 | contact_point_indices = self.hand_model.contact_point_indices.clone() 112 | contact_point_indices[switch_mask] = torch.randint( 113 | self.hand_model.n_contact_candidates, 114 | size=[switch_mask.sum()], 115 | device=self.device, 116 | ) 117 | 118 | # switch_mask = torch.rand(1).item() < self.switch_possibility 119 | 120 | # if switch_mask: 121 | # self.hand_model.other_idx = torch.randint(low=0, high=5, size=(1,)).item() 122 | 123 | # thumb_point_indices = torch.randint(len(self.hand_model.thumb_contact_indices), size=[batch_size, 2], device=self.device) 124 | # other_point_indices = torch.randint(len(self.hand_model.other_contact_indices[self.hand_model.other_idx]), size=[batch_size, n_contact - 2], device=self.device) 125 | # contact_point_indices = torch.cat([thumb_point_indices, other_point_indices], dim=1) 126 | 127 | self.old_hand_pose = self.hand_model.hand_pose 128 | self.old_contact_point_indices = self.hand_model.contact_point_indices 129 | self.old_vertices = self.hand_model.vertices 130 | self.old_contact_points = self.hand_model.contact_points 131 | self.old_grad_hand_pose = self.hand_model.hand_pose.grad 132 | self.hand_model.set_parameters(hand_pose, contact_point_indices) 133 | 134 | self.step += 1 135 | 136 | return s 137 | 138 | def accept_step(self, energy, new_energy): 139 | """ 140 | Accept / reject updates using annealing 141 | 142 | Returns 143 | ------- 144 | accept: (N,) torch.BoolTensor 145 | temperature: torch.Tensor 146 | current temperature 147 | """ 148 | 149 | batch_size = energy.shape[0] 150 | temperature = self.starting_temperature * self.temperature_decay ** torch.div( 151 | self.step, self.annealing_period, rounding_mode="floor" 152 | ) 153 | 154 | alpha = torch.rand(batch_size, dtype=torch.float, device=self.device) 155 | accept = alpha < torch.exp((energy - new_energy) / temperature) 156 | 157 | with torch.no_grad(): 158 | reject = ~accept 159 | self.hand_model.hand_pose[reject] = self.old_hand_pose[reject] 160 | self.hand_model.contact_point_indices[reject] = ( 161 | self.old_contact_point_indices[reject] 162 | ) 163 | self.hand_model.vertices[reject] = self.old_vertices[reject] 164 | self.hand_model.contact_points[reject] = self.old_contact_points[reject] 165 | self.hand_model.hand_pose.grad[reject] = self.old_grad_hand_pose[reject] 166 | 167 | return accept, temperature 168 | 169 | def zero_grad(self): 170 | """ 171 | Sets the gradients of translation, rotation, and joint angles to zero 172 | """ 173 | if self.hand_model.hand_pose.grad is not None: 174 | self.hand_model.hand_pose.grad.data.zero_() 175 | -------------------------------------------------------------------------------- /visualizer/tools/utils.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | 16 | import numpy as np 17 | import torch 18 | import logging 19 | from copy import copy 20 | 21 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 22 | to_cpu = lambda tensor: tensor.detach().cpu().numpy() 23 | 24 | def parse_npz(npz, allow_pickle=True): 25 | npz = np.load(npz, allow_pickle=allow_pickle) 26 | npz = {k: npz[k].item() for k in npz.files} 27 | return DotDict(npz) 28 | 29 | def params2torch(params, dtype = torch.float32): 30 | return {k: torch.from_numpy(v).type(dtype) for k, v in params.items()} 31 | 32 | def prepare_params(params, frame_mask, dtype = np.float32): 33 | return {k: v[frame_mask].astype(dtype) for k, v in params.items()} 34 | 35 | def DotDict(in_dict): 36 | 37 | out_dict = copy(in_dict) 38 | for k,v in out_dict.items(): 39 | if isinstance(v,dict): 40 | out_dict[k] = DotDict(v) 41 | return dotdict(out_dict) 42 | 43 | class dotdict(dict): 44 | """dot.notation access to dictionary attributes""" 45 | __getattr__ = dict.get 46 | __setattr__ = dict.__setitem__ 47 | __delattr__ = dict.__delitem__ 48 | 49 | 50 | def append2dict(source, data): 51 | for k in data.keys(): 52 | if isinstance(data[k], list): 53 | source[k] += data[k].astype(np.float32) 54 | else: 55 | source[k].append(data[k].astype(np.float32)) 56 | 57 | def np2torch(item): 58 | out = {} 59 | for k, v in item.items(): 60 | if v ==[]: 61 | continue 62 | if isinstance(v, list): 63 | try: 64 | out[k] = torch.from_numpy(np.concatenate(v)) 65 | except: 66 | out[k] = torch.from_numpy(np.array(v)) 67 | elif isinstance(v, dict): 68 | out[k] = np2torch(v) 69 | else: 70 | out[k] = torch.from_numpy(v) 71 | return out 72 | 73 | def to_tensor(array, dtype=torch.float32): 74 | if not torch.is_tensor(array): 75 | array = torch.tensor(array) 76 | return array.to(dtype) 77 | 78 | 79 | def to_np(array, dtype=np.float32): 80 | if 'scipy.sparse' in str(type(array)): 81 | array = np.array(array.todencse(), dtype=dtype) 82 | elif torch.is_tensor(array): 83 | array = array.detach().cpu().numpy() 84 | return array 85 | 86 | def makepath(desired_path, isfile = False): 87 | ''' 88 | if the path does not exist make it 89 | :param desired_path: can be path to a file or a folder name 90 | :return: 91 | ''' 92 | import os 93 | if isfile: 94 | if not os.path.exists(os.path.dirname(desired_path)):os.makedirs(os.path.dirname(desired_path)) 95 | else: 96 | if not os.path.exists(desired_path): os.makedirs(desired_path) 97 | return desired_path 98 | 99 | def makelogger(log_dir,mode='w'): 100 | 101 | makepath(log_dir, isfile=True) 102 | logger = logging.getLogger() 103 | logger.setLevel(logging.INFO) 104 | 105 | ch = logging.StreamHandler() 106 | ch.setLevel(logging.INFO) 107 | 108 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 109 | 110 | ch.setFormatter(formatter) 111 | 112 | logger.addHandler(ch) 113 | 114 | fh = logging.FileHandler('%s'%log_dir, mode=mode) 115 | fh.setFormatter(formatter) 116 | logger.addHandler(fh) 117 | 118 | return logger 119 | 120 | 121 | def euler(rots, order='xyz', units='deg'): 122 | 123 | rots = np.asarray(rots) 124 | single_val = False if len(rots.shape)>1 else True 125 | rots = rots.reshape(-1,3) 126 | rotmats = [] 127 | 128 | for xyz in rots: 129 | if units == 'deg': 130 | xyz = np.radians(xyz) 131 | r = np.eye(3) 132 | for theta, axis in zip(xyz,order): 133 | c = np.cos(theta) 134 | s = np.sin(theta) 135 | if axis=='x': 136 | r = np.dot(np.array([[1,0,0],[0,c,-s],[0,s,c]]), r) 137 | if axis=='y': 138 | r = np.dot(np.array([[c,0,s],[0,1,0],[-s,0,c]]), r) 139 | if axis=='z': 140 | r = np.dot(np.array([[c,-s,0],[s,c,0],[0,0,1]]), r) 141 | rotmats.append(r) 142 | rotmats = np.stack(rotmats).astype(np.float32) 143 | if single_val: 144 | return rotmats[0] 145 | else: 146 | return rotmats 147 | 148 | def create_video(path, fps=30,name='movie'): 149 | import os 150 | import subprocess 151 | 152 | src = os.path.join(path,'%*.png') 153 | movie_path = os.path.join(path,'%s.mp4'%name) 154 | i = 0 155 | while os.path.isfile(movie_path): 156 | movie_path = os.path.join(path,'%s_%02d.mp4'%(name,i)) 157 | i +=1 158 | 159 | cmd = 'ffmpeg -f image2 -r %d -i %s -b:v 6400k -pix_fmt yuv420p %s' % (fps, src, movie_path) 160 | 161 | subprocess.call(cmd.split(' ')) 162 | while not os.path.exists(movie_path): 163 | continue 164 | 165 | # mapping the contact ids to each body part in smplx 166 | contact_ids={'Body': 1, 167 | 'L_Thigh': 2, 168 | 'R_Thigh': 3, 169 | 'Spine': 4, 170 | 'L_Calf': 5, 171 | 'R_Calf': 6, 172 | 'Spine1': 7, 173 | 'L_Foot': 8, 174 | 'R_Foot': 9, 175 | 'Spine2': 10, 176 | 'L_Toes': 11, 177 | 'R_Toes': 12, 178 | 'Neck': 13, 179 | 'L_Shoulder': 14, 180 | 'R_Shoulder': 15, 181 | 'Head': 16, 182 | 'L_UpperArm': 17, 183 | 'R_UpperArm': 18, 184 | 'L_ForeArm': 19, 185 | 'R_ForeArm': 20, 186 | 'L_Hand': 21, 187 | 'R_Hand': 22, 188 | 'Jaw': 23, 189 | 'L_Eye': 24, 190 | 'R_Eye': 25, 191 | 'L_Index1': 26, 192 | 'L_Index2': 27, 193 | 'L_Index3': 28, 194 | 'L_Middle1': 29, 195 | 'L_Middle2': 30, 196 | 'L_Middle3': 31, 197 | 'L_Pinky1': 32, 198 | 'L_Pinky2': 33, 199 | 'L_Pinky3': 34, 200 | 'L_Ring1': 35, 201 | 'L_Ring2': 36, 202 | 'L_Ring3': 37, 203 | 'L_Thumb1': 38, 204 | 'L_Thumb2': 39, 205 | 'L_Thumb3': 40, 206 | 'R_Index1': 41, 207 | 'R_Index2': 42, 208 | 'R_Index3': 43, 209 | 'R_Middle1': 44, 210 | 'R_Middle2': 45, 211 | 'R_Middle3': 46, 212 | 'R_Pinky1': 47, 213 | 'R_Pinky2': 48, 214 | 'R_Pinky3': 49, 215 | 'R_Ring1': 50, 216 | 'R_Ring2': 51, 217 | 'R_Ring3': 52, 218 | 'R_Thumb1': 53, 219 | 'R_Thumb2': 54, 220 | 'R_Thumb3': 55} 221 | -------------------------------------------------------------------------------- /manip/utils/visualize/tools/utils.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | 15 | 16 | import numpy as np 17 | import torch 18 | import logging 19 | from copy import copy 20 | 21 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 22 | to_cpu = lambda tensor: tensor.detach().cpu().numpy() 23 | 24 | def parse_npz(npz, allow_pickle=True): 25 | npz = np.load(npz, allow_pickle=allow_pickle) 26 | npz = {k: npz[k].item() for k in npz.files} 27 | return DotDict(npz) 28 | 29 | def params2torch(params, dtype = torch.float32): 30 | return {k: torch.from_numpy(v).type(dtype) for k, v in params.items()} 31 | 32 | def prepare_params(params, frame_mask, dtype = np.float32): 33 | return {k: v[frame_mask].astype(dtype) for k, v in params.items()} 34 | 35 | def DotDict(in_dict): 36 | 37 | out_dict = copy(in_dict) 38 | for k,v in out_dict.items(): 39 | if isinstance(v,dict): 40 | out_dict[k] = DotDict(v) 41 | return dotdict(out_dict) 42 | 43 | class dotdict(dict): 44 | """dot.notation access to dictionary attributes""" 45 | __getattr__ = dict.get 46 | __setattr__ = dict.__setitem__ 47 | __delattr__ = dict.__delitem__ 48 | 49 | 50 | def append2dict(source, data): 51 | for k in data.keys(): 52 | if isinstance(data[k], list): 53 | source[k] += data[k].astype(np.float32) 54 | else: 55 | source[k].append(data[k].astype(np.float32)) 56 | 57 | def np2torch(item): 58 | out = {} 59 | for k, v in item.items(): 60 | if v ==[]: 61 | continue 62 | if isinstance(v, list): 63 | try: 64 | out[k] = torch.from_numpy(np.concatenate(v)) 65 | except: 66 | out[k] = torch.from_numpy(np.array(v)) 67 | elif isinstance(v, dict): 68 | out[k] = np2torch(v) 69 | else: 70 | out[k] = torch.from_numpy(v) 71 | return out 72 | 73 | def to_tensor(array, dtype=torch.float32): 74 | if not torch.is_tensor(array): 75 | array = torch.tensor(array) 76 | return array.to(dtype) 77 | 78 | 79 | def to_np(array, dtype=np.float32): 80 | if 'scipy.sparse' in str(type(array)): 81 | array = np.array(array.todencse(), dtype=dtype) 82 | elif torch.is_tensor(array): 83 | array = array.detach().cpu().numpy() 84 | return array 85 | 86 | def makepath(desired_path, isfile = False): 87 | ''' 88 | if the path does not exist make it 89 | :param desired_path: can be path to a file or a folder name 90 | :return: 91 | ''' 92 | import os 93 | if isfile: 94 | if not os.path.exists(os.path.dirname(desired_path)):os.makedirs(os.path.dirname(desired_path)) 95 | else: 96 | if not os.path.exists(desired_path): os.makedirs(desired_path) 97 | return desired_path 98 | 99 | def makelogger(log_dir,mode='w'): 100 | 101 | makepath(log_dir, isfile=True) 102 | logger = logging.getLogger() 103 | logger.setLevel(logging.INFO) 104 | 105 | ch = logging.StreamHandler() 106 | ch.setLevel(logging.INFO) 107 | 108 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 109 | 110 | ch.setFormatter(formatter) 111 | 112 | logger.addHandler(ch) 113 | 114 | fh = logging.FileHandler('%s'%log_dir, mode=mode) 115 | fh.setFormatter(formatter) 116 | logger.addHandler(fh) 117 | 118 | return logger 119 | 120 | 121 | def euler(rots, order='xyz', units='deg'): 122 | 123 | rots = np.asarray(rots) 124 | single_val = False if len(rots.shape)>1 else True 125 | rots = rots.reshape(-1,3) 126 | rotmats = [] 127 | 128 | for xyz in rots: 129 | if units == 'deg': 130 | xyz = np.radians(xyz) 131 | r = np.eye(3) 132 | for theta, axis in zip(xyz,order): 133 | c = np.cos(theta) 134 | s = np.sin(theta) 135 | if axis=='x': 136 | r = np.dot(np.array([[1,0,0],[0,c,-s],[0,s,c]]), r) 137 | if axis=='y': 138 | r = np.dot(np.array([[c,0,s],[0,1,0],[-s,0,c]]), r) 139 | if axis=='z': 140 | r = np.dot(np.array([[c,-s,0],[s,c,0],[0,0,1]]), r) 141 | rotmats.append(r) 142 | rotmats = np.stack(rotmats).astype(np.float32) 143 | if single_val: 144 | return rotmats[0] 145 | else: 146 | return rotmats 147 | 148 | def create_video(path, fps=30,name='movie'): 149 | import os 150 | import subprocess 151 | 152 | src = os.path.join(path,'%*.png') 153 | movie_path = os.path.join(path,'%s.mp4'%name) 154 | i = 0 155 | while os.path.isfile(movie_path): 156 | movie_path = os.path.join(path,'%s_%02d.mp4'%(name,i)) 157 | i +=1 158 | 159 | cmd = 'ffmpeg -f image2 -r %d -i %s -b:v 6400k -pix_fmt yuv420p %s' % (fps, src, movie_path) 160 | 161 | subprocess.call(cmd.split(' ')) 162 | while not os.path.exists(movie_path): 163 | continue 164 | 165 | # mapping the contact ids to each body part in smplx 166 | contact_ids={'Body': 1, 167 | 'L_Thigh': 2, 168 | 'R_Thigh': 3, 169 | 'Spine': 4, 170 | 'L_Calf': 5, 171 | 'R_Calf': 6, 172 | 'Spine1': 7, 173 | 'L_Foot': 8, 174 | 'R_Foot': 9, 175 | 'Spine2': 10, 176 | 'L_Toes': 11, 177 | 'R_Toes': 12, 178 | 'Neck': 13, 179 | 'L_Shoulder': 14, 180 | 'R_Shoulder': 15, 181 | 'Head': 16, 182 | 'L_UpperArm': 17, 183 | 'R_UpperArm': 18, 184 | 'L_ForeArm': 19, 185 | 'R_ForeArm': 20, 186 | 'L_Hand': 21, 187 | 'R_Hand': 22, 188 | 'Jaw': 23, 189 | 'L_Eye': 24, 190 | 'R_Eye': 25, 191 | 'L_Index1': 26, 192 | 'L_Index2': 27, 193 | 'L_Index3': 28, 194 | 'L_Middle1': 29, 195 | 'L_Middle2': 30, 196 | 'L_Middle3': 31, 197 | 'L_Pinky1': 32, 198 | 'L_Pinky2': 33, 199 | 'L_Pinky3': 34, 200 | 'L_Ring1': 35, 201 | 'L_Ring2': 36, 202 | 'L_Ring3': 37, 203 | 'L_Thumb1': 38, 204 | 'L_Thumb2': 39, 205 | 'L_Thumb3': 40, 206 | 'R_Index1': 41, 207 | 'R_Index2': 42, 208 | 'R_Index3': 43, 209 | 'R_Middle1': 44, 210 | 'R_Middle2': 45, 211 | 'R_Middle3': 46, 212 | 'R_Pinky1': 47, 213 | 'R_Pinky2': 48, 214 | 'R_Pinky3': 49, 215 | 'R_Ring1': 50, 216 | 'R_Ring2': 51, 217 | 'R_Ring3': 52, 218 | 'R_Thumb1': 53, 219 | 'R_Thumb2': 54, 220 | 'R_Thumb3': 55} 221 | -------------------------------------------------------------------------------- /manip/inertialize/inert.py: -------------------------------------------------------------------------------- 1 | import pytorch3d.transforms as transforms 2 | import torch 3 | 4 | from manip.inertialize.spring import ( 5 | decay_spring_damper_exact_cubic, 6 | ) 7 | from manip.lafan1.utils import quat_slerp 8 | from manip.utils.model_utils import wxyz_to_xyzw, xyzw_to_wxyz 9 | 10 | 11 | def quat_abs(q): 12 | q[q[..., 0] < 0] *= -1 13 | return q 14 | 15 | 16 | def apply_inertialize( 17 | prev_jpos=None, 18 | prev_rot_6d=None, 19 | window_jpos=None, 20 | window_rot_6d=None, 21 | ratio=0.5, 22 | prev_blend_time=0.2, 23 | window_blend_time=0.2, 24 | zero_velocity=False, 25 | ): 26 | """ 27 | prev_jpos: BS X T X 24 X 3 / BS X T X 3 28 | prev_rot_6d: BS X T X 22 X 6 / BS X T X 6 29 | window_jpos: BS X T X 24 X 3 / BS X T X 3 30 | window_rot_6d: BS X T X 22 X 6 / BS X T X 6 31 | ratio: 0 - all prev, 1 - all window 32 | """ 33 | 34 | dt = 1 / 30.0 35 | ##### pos ##### 36 | if prev_jpos is not None and window_jpos is not None: 37 | dst_pos = window_jpos[:, 0] 38 | dst_pos_next = window_jpos[:, 1] 39 | dst_vel = (dst_pos_next - dst_pos) / dt # BS X 24 X 3 40 | 41 | src_pos = prev_jpos[:, -1] 42 | src_pos_prev = prev_jpos[:, -2] 43 | src_vel = (src_pos - src_pos_prev) / dt # BS X 24 X 3 44 | 45 | diff_pos = src_pos - dst_pos 46 | diff_vel = src_vel - dst_vel 47 | 48 | if zero_velocity: 49 | diff_vel[:] = 0.0 50 | 51 | ##### interpolate ##### 52 | new_jpos = window_jpos.clone() 53 | new_prev_jpos = prev_jpos.clone() 54 | for i in range(new_jpos.shape[1]): 55 | offset = decay_spring_damper_exact_cubic( 56 | ratio * diff_pos, ratio * diff_vel, window_blend_time, i * dt 57 | ) 58 | new_jpos[:, i] += offset 59 | for i in range(prev_jpos.shape[1]): 60 | offset = decay_spring_damper_exact_cubic( 61 | (1 - ratio) * -diff_pos, 62 | (1 - ratio) * diff_vel, 63 | prev_blend_time, 64 | (prev_jpos.shape[1] - 1 - i) * dt, 65 | ) 66 | new_prev_jpos[:, i] += offset 67 | else: 68 | new_jpos = None 69 | new_prev_jpos = None 70 | 71 | ##### rot ##### 72 | if prev_rot_6d is not None and window_rot_6d is not None: 73 | dst_ori = transforms.matrix_to_quaternion( 74 | transforms.rotation_6d_to_matrix(window_rot_6d[:, 0]) 75 | ) 76 | dst_ori_next = transforms.matrix_to_quaternion( 77 | transforms.rotation_6d_to_matrix(window_rot_6d[:, 1]) 78 | ) 79 | dst_ang_vel = ( 80 | transforms.quaternion_to_axis_angle( 81 | quat_abs( 82 | transforms.quaternion_multiply( 83 | dst_ori_next, transforms.quaternion_invert(dst_ori) 84 | ) 85 | ) 86 | ) 87 | / dt 88 | ) 89 | 90 | src_ori = transforms.matrix_to_quaternion( 91 | transforms.rotation_6d_to_matrix(prev_rot_6d[:, -1]) 92 | ) 93 | src_ori_prev = transforms.matrix_to_quaternion( 94 | transforms.rotation_6d_to_matrix(prev_rot_6d[:, -2]) 95 | ) 96 | src_ang_vel = ( 97 | transforms.quaternion_to_axis_angle( 98 | quat_abs( 99 | transforms.quaternion_multiply( 100 | src_ori, transforms.quaternion_invert(src_ori_prev) 101 | ) 102 | ) 103 | ) 104 | / dt 105 | ) 106 | 107 | diff_ori = transforms.quaternion_to_axis_angle( 108 | quat_abs( 109 | transforms.quaternion_multiply( 110 | src_ori, transforms.quaternion_invert(dst_ori) 111 | ) 112 | ) 113 | ) 114 | diff_ang_vel = src_ang_vel - dst_ang_vel 115 | 116 | if zero_velocity: 117 | diff_ang_vel[:] = 0.0 118 | 119 | new_rot_6d = window_rot_6d.clone() 120 | new_prev_rot_6d = prev_rot_6d.clone() 121 | 122 | ##### interpolate ##### 123 | for i in range(new_rot_6d.shape[1]): 124 | offset = decay_spring_damper_exact_cubic( 125 | ratio * diff_ori, ratio * diff_ang_vel, window_blend_time, i * dt 126 | ) 127 | new_rot_6d[:, i] = transforms.matrix_to_rotation_6d( 128 | transforms.quaternion_to_matrix( 129 | transforms.quaternion_multiply( 130 | transforms.axis_angle_to_quaternion(offset), 131 | transforms.matrix_to_quaternion( 132 | transforms.rotation_6d_to_matrix(new_rot_6d[:, i]) 133 | ), 134 | ) 135 | ) 136 | ) 137 | for i in range(new_prev_rot_6d.shape[1]): 138 | offset = decay_spring_damper_exact_cubic( 139 | (1 - ratio) * -diff_ori, 140 | (1 - ratio) * diff_ang_vel, 141 | prev_blend_time, 142 | (new_prev_rot_6d.shape[1] - 1 - i) * dt, 143 | ) 144 | new_prev_rot_6d[:, i] = transforms.matrix_to_rotation_6d( 145 | transforms.quaternion_to_matrix( 146 | transforms.quaternion_multiply( 147 | transforms.axis_angle_to_quaternion(offset), 148 | transforms.matrix_to_quaternion( 149 | transforms.rotation_6d_to_matrix(new_prev_rot_6d[:, i]) 150 | ), 151 | ) 152 | ) 153 | ) 154 | else: 155 | new_rot_6d = None 156 | new_prev_rot_6d = None 157 | 158 | return new_jpos, new_rot_6d, new_prev_jpos, new_prev_rot_6d 159 | 160 | 161 | def apply_linear_offset( 162 | original_jpos=None, 163 | original_rot_6d=None, 164 | new_target_jpos=None, 165 | new_target_rot_6d=None, 166 | reversed=False, 167 | ): 168 | """ 169 | original_jpos: BS X T X 3 / BS X T X 24 X 3 170 | original_rot_6d: BS X T X 6 / BS X T X 22 X 6 171 | new_target_jpos: BS X 3 / BS X 24 X 3 172 | new_target_rot_6d: BS X 6 / BS X 22 X 6 173 | """ 174 | new_jpos, new_rot_6d = None, None 175 | if reversed: 176 | if original_jpos is not None: 177 | original_jpos = torch.flip(original_jpos, dims=[1]) 178 | if original_rot_6d is not None: 179 | original_rot_6d = torch.flip(original_rot_6d, dims=[1]) 180 | 181 | ##### pos ##### 182 | if original_jpos is not None: 183 | T = original_jpos.shape[1] 184 | new_jpos = original_jpos.clone() 185 | pos_offset = new_target_jpos - original_jpos[:, -1] 186 | for i in range(T): 187 | new_jpos[:, i] += (i + 1) / T * pos_offset 188 | 189 | ##### rot ##### 190 | if original_rot_6d is not None: 191 | T = original_rot_6d.shape[1] 192 | new_target_quat = transforms.matrix_to_quaternion( 193 | transforms.rotation_6d_to_matrix(new_target_rot_6d) 194 | ) 195 | original_quat = transforms.matrix_to_quaternion( 196 | transforms.rotation_6d_to_matrix(original_rot_6d) 197 | ) 198 | new_quat = original_quat.clone() 199 | quat_offset = transforms.quaternion_multiply( 200 | new_target_quat, transforms.quaternion_invert(original_quat[:, -1]) 201 | ) 202 | 203 | # quat_slerp needs xyzw 204 | quat_offset = wxyz_to_xyzw(quat_offset) 205 | zero_quat = torch.zeros_like(quat_offset) 206 | zero_quat[..., -1] = 1 207 | for i in range(T): 208 | cur_offset = quat_slerp(zero_quat, quat_offset, (i + 1) / T) 209 | cur_offset = xyzw_to_wxyz(cur_offset) 210 | new_quat[:, i] = transforms.quaternion_multiply( 211 | cur_offset, original_quat[:, i] 212 | ) 213 | new_rot_6d = transforms.matrix_to_rotation_6d( 214 | transforms.quaternion_to_matrix(new_quat) 215 | ) 216 | 217 | if reversed: 218 | if original_jpos is not None: 219 | new_jpos = torch.flip(new_jpos, dims=[1]) 220 | if original_rot_6d is not None: 221 | new_rot_6d = torch.flip(new_rot_6d, dims=[1]) 222 | 223 | return new_jpos, new_rot_6d 224 | -------------------------------------------------------------------------------- /visualizer/vis/visualize_finger_results.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | # 4 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 5 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 6 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 7 | # 8 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 9 | # on this computer program. You can only use this computer program if you have closed a license agreement 10 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 11 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 12 | # Contact: ps-license@tuebingen.mpg.de 13 | # 14 | import sys 15 | sys.path.append('.') 16 | sys.path.append('..') 17 | 18 | import numpy as np 19 | import torch 20 | import os, glob 21 | import smplx 22 | import argparse 23 | from tqdm import tqdm 24 | os.environ['PYOPENGL_PLATFORM'] = 'egl' 25 | 26 | from visualizer.tools.objectmodel import ObjectModel 27 | from visualizer.tools.meshviewer import Mesh, MeshViewer, points2sphere, colors 28 | from visualizer.tools.utils import parse_npz 29 | from visualizer.tools.utils import params2torch 30 | from visualizer.tools.utils import to_cpu 31 | from visualizer.tools.utils import euler 32 | from visualizer.tools.cfg_parser import Config 33 | 34 | import trimesh 35 | 36 | USE_FLAT_HAND_MEAN = True 37 | 38 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 39 | 40 | def visualize_sequences(cfg): 41 | 42 | result_path = cfg.result_path 43 | 44 | mv = MeshViewer(offscreen=cfg.offscreen) 45 | 46 | # set the camera pose 47 | camera_pose = np.eye(4) 48 | camera_pose[:3, :3] = euler([80, -15, 0], 'xzx') 49 | camera_pose[:3, 3] = np.array([-.5, -3., 1.5]) 50 | mv.update_camera_pose(camera_pose) 51 | 52 | vis_sequence(cfg, result_path, mv) 53 | mv.close_viewer() 54 | 55 | 56 | def vis_sequence(cfg, result_path, mv): 57 | human_meshes = [] 58 | object_meshes = [] 59 | human_meshes_gt = [] 60 | object_meshes_gt = [] 61 | 62 | ori_obj_files = os.listdir(result_path) 63 | ori_obj_files.sort() 64 | human_files = [] 65 | object_files = [] 66 | for tmp_name in ori_obj_files: 67 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 68 | if ".ply" in tmp_name: 69 | if "object" not in tmp_name: 70 | human_files.append(tmp_name) 71 | else: 72 | object_files.append(tmp_name) 73 | 74 | T = len(human_files) 75 | for i in range(T): 76 | human_meshes.append(trimesh.load(os.path.join(result_path, human_files[i]))) 77 | object_meshes.append(trimesh.load(os.path.join(result_path, object_files[i]))) 78 | 79 | idx_str = result_path[result_path.find("idx"):result_path.rfind("/")] 80 | gt_path = result_path.replace(idx_str, idx_str+"_gt") 81 | vis_gt = os.path.exists(gt_path) and len(os.listdir(gt_path)) > 0 82 | if vis_gt: 83 | offset = np.array([1, 0, 0]) 84 | ori_obj_files = os.listdir(gt_path) 85 | ori_obj_files.sort() 86 | human_files = [] 87 | object_files = [] 88 | for tmp_name in ori_obj_files: 89 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 90 | if ".ply" in tmp_name: 91 | if "object" not in tmp_name: 92 | human_files.append(tmp_name) 93 | else: 94 | object_files.append(tmp_name) 95 | 96 | T = len(human_files) 97 | for i in range(T): 98 | human_meshes_gt.append(trimesh.load(os.path.join(gt_path, human_files[i]))) 99 | object_meshes_gt.append(trimesh.load(os.path.join(gt_path, object_files[i]))) 100 | 101 | skip_frame = 1 102 | o_meshes = [] 103 | s_meshes = [] 104 | o_meshes_gt = [] 105 | s_meshes_gt = [] 106 | for frame in range(0, T, skip_frame): 107 | o_mesh = Mesh(vertices=object_meshes[frame].vertices, faces=object_meshes[frame].faces, vc=colors['yellow']) 108 | s_mesh = Mesh(vertices=human_meshes[frame].vertices, faces=human_meshes[frame].faces, vc=colors['pink'], smooth=True) 109 | 110 | # o_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=abs(object_meshes[frame].vertices[:, 0] - (-0.22)) < 0.01) 111 | # s_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=abs(human_meshes[frame].vertices[:, 0] - (-0.22)) < 0.01) 112 | # s_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=abs(human_meshes[frame].vertices[:, 1] - (-0.522)) < 0.01) 113 | # s_mesh.set_vertex_colors(vc=colors['red'], vertex_ids=abs(human_meshes[frame].vertices[:, 2] - (1)) < 0.01) 114 | 115 | o_meshes.append(o_mesh) 116 | s_meshes.append(s_mesh) 117 | if vis_gt: 118 | o_mesh_gt = Mesh(vertices=object_meshes_gt[frame].vertices + offset, faces=object_meshes_gt[frame].faces, vc=colors['green']) 119 | s_mesh_gt = Mesh(vertices=human_meshes_gt[frame].vertices + offset, faces=human_meshes_gt[frame].faces, vc=colors['pink'], smooth=True) 120 | o_meshes_gt.append(o_mesh_gt) 121 | s_meshes_gt.append(s_mesh_gt) 122 | 123 | if not cfg.offscreen: 124 | while True: 125 | import time 126 | for frame in range(0, T, skip_frame): 127 | start_time = time.time() 128 | o_mesh = o_meshes[frame] 129 | s_mesh = s_meshes[frame] 130 | 131 | if vis_gt: 132 | o_mesh_gt = o_meshes_gt[frame] 133 | s_mesh_gt = s_meshes_gt[frame] 134 | mv.set_static_meshes([o_mesh, s_mesh, o_mesh_gt, s_mesh_gt]) 135 | else: 136 | mv.set_static_meshes([o_mesh, s_mesh]) 137 | while time.time() - start_time < 0.03: 138 | pass 139 | else: 140 | import imageio 141 | 142 | img_path = os.path.join(result_path, "vis") 143 | if not os.path.exists(img_path): 144 | os.makedirs(img_path) 145 | 146 | path1, path2 = result_path.split("/")[-2], result_path.split("/")[-1] 147 | save_dir = os.path.join(result_path, "../../vis", path1) 148 | if not os.path.exists(save_dir): 149 | os.makedirs(save_dir) 150 | save_dir = os.path.join(save_dir, path2) 151 | if not os.path.exists(save_dir): 152 | os.makedirs(save_dir) 153 | 154 | for frame in tqdm(range(0, T, skip_frame)): 155 | o_mesh = o_meshes[frame] 156 | s_mesh = s_meshes[frame] 157 | 158 | if vis_gt: 159 | o_mesh_gt = o_meshes_gt[frame] 160 | s_mesh_gt = s_meshes_gt[frame] 161 | mv.set_static_meshes([o_mesh, s_mesh, o_mesh_gt, s_mesh_gt]) 162 | else: 163 | mv.set_static_meshes([o_mesh, s_mesh]) 164 | 165 | camera_pose = np.eye(4) 166 | camera_pose[:3, :3] = euler([80, -15, 0], 'xzx') 167 | camera_pose[:2, 3] = np.mean(o_mesh.vertices, axis=0)[:2] + np.array([0, -2.]) 168 | camera_pose[2, 3] = 1.3 169 | mv.update_camera_pose(camera_pose) 170 | 171 | mv.save_snapshot(os.path.join(img_path, "%05d.png"%frame)) 172 | 173 | video_name = os.path.join(save_dir, 'output.mp4') 174 | 175 | images = [img for img in os.listdir(img_path) if img.endswith(".png")] 176 | images.sort() 177 | im_arr = [] 178 | 179 | for image in images: 180 | path = os.path.join(img_path, image) 181 | im = imageio.v2.imread(path) 182 | im_arr.append(im) 183 | im_arr = np.asarray(im_arr) 184 | imageio.mimwrite(video_name, im_arr, fps=30, quality=8) 185 | 186 | print("video saved to %s"%video_name) 187 | 188 | if __name__ == '__main__': 189 | 190 | 191 | parser = argparse.ArgumentParser(description='GRAB-visualize') 192 | 193 | parser.add_argument('--result-path', required=True, type=str, 194 | help='The path to the downloaded grab data') 195 | 196 | parser.add_argument('--model-path', required=True, type=str, 197 | help='The path to the folder containing smplx models') 198 | 199 | parser.add_argument("--offscreen", action="store_true") 200 | 201 | args = parser.parse_args() 202 | 203 | result_path = args.result_path 204 | model_path = args.model_path 205 | 206 | # grab_path = 'PATH_TO_DOWNLOADED_GRAB_DATA/grab' 207 | # model_path = 'PATH_TO_DOWNLOADED_MODELS_FROM_SMPLX_WEBSITE/' 208 | 209 | cfg = { 210 | 'result_path': result_path, 211 | 'model_path': model_path, 212 | 'offscreen': args.offscreen, 213 | } 214 | 215 | cfg = Config(**cfg) 216 | visualize_sequences(cfg) 217 | 218 | -------------------------------------------------------------------------------- /visualizer/vis/visualize_navigation_results.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 4 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 5 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 6 | # 7 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 8 | # on this computer program. You can only use this computer program if you have closed a license agreement 9 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 10 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 11 | # Contact: ps-license@tuebingen.mpg.de 12 | # 13 | import sys 14 | 15 | sys.path.append(".") 16 | sys.path.append("..") 17 | 18 | import argparse 19 | import os 20 | 21 | import numpy as np 22 | import torch 23 | from tqdm import tqdm 24 | 25 | os.environ["PYOPENGL_PLATFORM"] = "egl" 26 | 27 | import trimesh 28 | 29 | from visualizer.tools.cfg_parser import Config 30 | from visualizer.tools.meshviewer import Mesh, MeshViewer, colors, points2sphere 31 | from visualizer.tools.utils import euler, params2torch, parse_npz, to_cpu 32 | 33 | USE_FLAT_HAND_MEAN = True 34 | 35 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 36 | 37 | 38 | def visualize_sequences(cfg): 39 | result_path = cfg.result_path 40 | 41 | mv = MeshViewer(offscreen=cfg.offscreen) 42 | 43 | # set the camera pose 44 | camera_pose = np.eye(4) 45 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 46 | camera_pose[:3, 3] = np.array([-0.5, -5.0, 1.5]) 47 | mv.update_camera_pose(camera_pose) 48 | 49 | vis_sequence(cfg, result_path, mv, cfg.interaction_epoch, cfg.s_idx) 50 | mv.close_viewer() 51 | 52 | 53 | def vis_sequence(cfg, result_path, mv, interaction_epoch, s_idx): 54 | # 1 pass, process all data 55 | ball_meshes = [] 56 | human_meshes = [] 57 | 58 | b_meshes = [] 59 | s_meshes = [] 60 | b_meshes_gt = [] 61 | s_meshes_gt = [] 62 | 63 | ball_path = os.path.join( 64 | result_path, "ball_objs_step_{}_bs_idx_0".format(interaction_epoch) 65 | ) 66 | subject_path = os.path.join( 67 | result_path, "objs_step_{}_bs_idx_0".format(interaction_epoch) 68 | ) 69 | 70 | ori_ball_files = os.listdir(ball_path) 71 | ori_ball_files.sort() 72 | 73 | for tmp_name in ori_ball_files: 74 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 75 | if ".ply" in tmp_name: 76 | if "start" in tmp_name: 77 | continue 78 | ball_meshes.append(trimesh.load(os.path.join(ball_path, tmp_name))) 79 | b_meshes.append( 80 | Mesh( 81 | vertices=ball_meshes[-1].vertices, 82 | faces=ball_meshes[-1].faces, 83 | vc=colors["green"], 84 | ) 85 | ) 86 | ori_obj_files = os.listdir(subject_path) 87 | ori_obj_files.sort() 88 | for tmp_name in ori_obj_files: 89 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 90 | if ".ply" in tmp_name: 91 | if "object" not in tmp_name: 92 | human_meshes.append(trimesh.load(os.path.join(subject_path, tmp_name))) 93 | s_meshes.append( 94 | Mesh( 95 | vertices=human_meshes[-1].vertices, 96 | faces=human_meshes[-1].faces, 97 | vc=colors["pink"], 98 | smooth=True, 99 | ) 100 | ) 101 | 102 | # gt 103 | ball_path_gt = os.path.join( 104 | result_path, "ball_objs_step_{}_bs_idx_0_gt".format(interaction_epoch) 105 | ) 106 | subject_path_gt = os.path.join( 107 | result_path, "objs_step_{}_bs_idx_0_gt".format(interaction_epoch) 108 | ) 109 | vis_gt = os.path.exists(subject_path_gt) and len(os.listdir(subject_path_gt)) > 0 110 | if vis_gt: 111 | offset = np.array([2, 0, 0]) 112 | ori_ball_files_gt = os.listdir(ball_path_gt) 113 | ori_ball_files_gt.sort() 114 | for tmp_name in ori_ball_files_gt: 115 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 116 | if ".ply" in tmp_name: 117 | if "start" in tmp_name: 118 | continue 119 | mesh = trimesh.load(os.path.join(ball_path_gt, tmp_name)) 120 | new_v = mesh.vertices + offset 121 | b_meshes_gt.append( 122 | Mesh(vertices=new_v, faces=mesh.faces, vc=colors["green"]) 123 | ) 124 | ori_obj_files_gt = os.listdir(subject_path_gt) 125 | ori_obj_files_gt.sort() 126 | for tmp_name in ori_obj_files_gt: 127 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 128 | if ".ply" in tmp_name: 129 | if "object" not in tmp_name: 130 | mesh = trimesh.load(os.path.join(subject_path_gt, tmp_name)) 131 | new_v = mesh.vertices + offset 132 | s_meshes_gt.append( 133 | Mesh( 134 | vertices=new_v, 135 | faces=mesh.faces, 136 | vc=colors["yellow"], 137 | smooth=True, 138 | ) 139 | ) 140 | 141 | # 2 pass, render the data 142 | skip_frame = 1 143 | if not cfg.offscreen: 144 | while True: 145 | import time 146 | 147 | T = len(s_meshes) 148 | for frame in range(0, T, skip_frame): 149 | start_time = time.time() 150 | 151 | s_mesh = s_meshes[frame] 152 | meshes = [s_mesh] 153 | 154 | meshes.extend(b_meshes) 155 | if vis_gt: 156 | s_mesh_gt = s_meshes_gt[frame] 157 | meshes.extend([s_mesh_gt]) 158 | meshes.extend(b_meshes_gt) 159 | 160 | mv.set_static_meshes(meshes) 161 | while time.time() - start_time < 0.03: 162 | pass 163 | else: 164 | import imageio 165 | 166 | img_paths = [] 167 | 168 | path1, path2 = result_path.split("/")[-2], result_path.split("/")[-1] 169 | save_dir = os.path.join(result_path, "../..", "vis") 170 | if not os.path.exists(save_dir): 171 | os.makedirs(save_dir) 172 | save_dir = os.path.join(save_dir, path1) 173 | if not os.path.exists(save_dir): 174 | os.makedirs(save_dir) 175 | 176 | img_path = os.path.join(result_path, "vis") 177 | if not os.path.exists(img_path): 178 | os.makedirs(img_path) 179 | 180 | T = len(s_meshes) 181 | for frame in tqdm(range(0, T, skip_frame), leave=False): 182 | s_mesh = s_meshes[frame] 183 | meshes = [s_mesh] 184 | 185 | meshes.extend(b_meshes) 186 | if vis_gt: 187 | s_mesh_gt = s_meshes_gt[frame] 188 | meshes.extend([s_mesh_gt]) 189 | meshes.extend(b_meshes_gt) 190 | 191 | mv.set_static_meshes(meshes) 192 | 193 | camera_pose = np.eye(4) 194 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 195 | camera_pose[:2, 3] = np.mean(s_mesh.vertices, axis=0)[:2] + np.array( 196 | [0.5, -3.0] 197 | ) 198 | camera_pose[2, 3] = 1.3 199 | mv.update_camera_pose(camera_pose) 200 | 201 | mv.save_snapshot(os.path.join(img_path, "%05d.png" % frame)) 202 | img_paths.append(os.path.join(img_path, "%05d.png" % frame)) 203 | 204 | video_name = os.path.join(save_dir, "output_{}.mp4".format(s_idx)) 205 | 206 | im_arr = [] 207 | 208 | for image in img_paths: 209 | im = imageio.v2.imread(image) 210 | im_arr.append(im) 211 | im_arr = np.asarray(im_arr) 212 | imageio.mimwrite(video_name, im_arr, fps=30, quality=8) 213 | 214 | print("video saved to %s" % video_name) 215 | 216 | 217 | if __name__ == "__main__": 218 | parser = argparse.ArgumentParser(description="visualize") 219 | 220 | parser.add_argument( 221 | "--result-path", required=True, type=str, help="The path to the results" 222 | ) 223 | parser.add_argument("--s-idx", required=True, type=str) 224 | parser.add_argument("--interaction-epoch", required=True, type=int) 225 | parser.add_argument( 226 | "--model-path", 227 | required=True, 228 | type=str, 229 | help="The path to the folder containing smplx models", 230 | ) 231 | parser.add_argument("--offscreen", action="store_true") 232 | 233 | args = parser.parse_args() 234 | 235 | result_path = args.result_path 236 | model_path = args.model_path 237 | 238 | cfg = { 239 | "result_path": result_path, 240 | "model_path": model_path, 241 | "offscreen": args.offscreen, 242 | "s_idx": args.s_idx, 243 | "interaction_epoch": args.interaction_epoch, 244 | } 245 | 246 | cfg = Config(**cfg) 247 | visualize_sequences(cfg) 248 | -------------------------------------------------------------------------------- /manip/utils/visualize/tools/meshviewer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2019 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems and the Max Planck Institute for Biological 14 | # Cybernetics. All rights reserved. 15 | # 16 | # Contact: ps-license@tuebingen.mpg.de 17 | # 18 | 19 | import numpy as np 20 | import trimesh 21 | import pyrender 22 | from pyrender.light import DirectionalLight 23 | from pyrender.node import Node 24 | from PIL import Image 25 | from .utils import euler 26 | 27 | 28 | class Mesh(trimesh.Trimesh): 29 | 30 | def __init__(self, 31 | filename=None, 32 | vertices=None, 33 | faces=None, 34 | vc=None, 35 | fc=None, 36 | vscale=None, 37 | process = False, 38 | visual = None, 39 | wireframe=False, 40 | smooth = False, 41 | **kwargs): 42 | 43 | self.wireframe = wireframe 44 | self.smooth = smooth 45 | 46 | if filename is not None: 47 | mesh = trimesh.load(filename, process = process) 48 | vertices = mesh.vertices 49 | faces= mesh.faces 50 | visual = mesh.visual 51 | if vscale is not None: 52 | vertices = vertices*vscale 53 | 54 | if faces is None: 55 | mesh = points2sphere(vertices) 56 | vertices = mesh.vertices 57 | faces = mesh.faces 58 | visual = mesh.visual 59 | 60 | super(Mesh, self).__init__(vertices=vertices, faces=faces, process=process, visual=visual) 61 | 62 | if vc is not None: 63 | self.set_vertex_colors(vc) 64 | if fc is not None: 65 | self.set_face_colors(fc) 66 | 67 | def rot_verts(self, vertices, rxyz): 68 | return np.array(vertices * rxyz.T) 69 | 70 | def colors_like(self,color, array, ids): 71 | 72 | color = np.array(color) 73 | 74 | if color.max() <= 1.: 75 | color = color * 255 76 | color = color.astype(np.int8) 77 | 78 | n_color = color.shape[0] 79 | n_ids = ids.shape[0] 80 | 81 | new_color = np.array(array) 82 | if n_color <= 4: 83 | new_color[ids, :n_color] = np.repeat(color[np.newaxis], n_ids, axis=0) 84 | else: 85 | new_color[ids, :] = color 86 | 87 | return new_color 88 | 89 | def set_vertex_colors(self,vc, vertex_ids = None): 90 | 91 | all_ids = np.arange(self.vertices.shape[0]) 92 | if vertex_ids is None: 93 | vertex_ids = all_ids 94 | 95 | vertex_ids = all_ids[vertex_ids] 96 | new_vc = self.colors_like(vc, self.visual.vertex_colors, vertex_ids) 97 | self.visual.vertex_colors[:] = new_vc 98 | 99 | def set_face_colors(self,fc, face_ids = None): 100 | 101 | if face_ids is None: 102 | face_ids = np.arange(self.faces.shape[0]) 103 | 104 | new_fc = self.colors_like(fc, self.visual.face_colors, face_ids) 105 | self.visual.face_colors[:] = new_fc 106 | 107 | @staticmethod 108 | def concatenate_meshes(meshes): 109 | return trimesh.util.concatenate(meshes) 110 | 111 | 112 | class MeshViewer(object): 113 | 114 | def __init__(self, 115 | width=1200, 116 | height=800, 117 | bg_color = [0.0, 0.0, 0.0, 1.0], 118 | offscreen = False, 119 | registered_keys=None): 120 | super(MeshViewer, self).__init__() 121 | 122 | if registered_keys is None: 123 | registered_keys = dict() 124 | 125 | self.bg_color = bg_color 126 | self.offscreen = offscreen 127 | self.scene = pyrender.Scene(bg_color=bg_color, 128 | ambient_light=(0.3, 0.3, 0.3), 129 | name = 'scene') 130 | 131 | self.aspect_ratio = float(width) / height 132 | pc = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=self.aspect_ratio) 133 | camera_pose = np.eye(4) 134 | camera_pose[:3,:3] = euler([80,-15,0], 'xzx') 135 | camera_pose[:3, 3] = np.array([-.5, -2., 1.5]) 136 | 137 | self.cam = pyrender.Node(name = 'camera', camera=pc, matrix=camera_pose) 138 | 139 | self.scene.add_node(self.cam) 140 | 141 | if self.offscreen: 142 | light = Node(light=DirectionalLight(color=np.ones(3), intensity=3.0), 143 | matrix=camera_pose) 144 | self.scene.add_node(light) 145 | self.viewer = pyrender.OffscreenRenderer(width, height) 146 | else: 147 | self.viewer = pyrender.Viewer(self.scene, 148 | use_raymond_lighting=True, 149 | viewport_size=(width, height), 150 | cull_faces=False, 151 | run_in_thread=True, 152 | registered_keys=registered_keys) 153 | 154 | for i, node in enumerate(self.scene.get_nodes()): 155 | if node.name is None: 156 | node.name = 'Req%d'%i 157 | 158 | 159 | def is_active(self): 160 | return self.viewer.is_active 161 | 162 | def close_viewer(self): 163 | if self.viewer.is_active: 164 | self.viewer.close_external() 165 | 166 | def set_background_color(self, bg_color=[1., 1., 1.]): 167 | self.scene.bg_color = bg_color 168 | 169 | def to_pymesh(self, mesh): 170 | 171 | wireframe = mesh.wireframe if hasattr(mesh, 'wireframe') else False 172 | smooth = mesh.smooth if hasattr(mesh, 'smooth') else False 173 | return pyrender.Mesh.from_trimesh(mesh, wireframe=wireframe, smooth=smooth) 174 | 175 | def update_camera_pose(self, pose): 176 | if self.offscreen: 177 | self.scene.set_pose(self.cam, pose=pose) 178 | else: 179 | self.viewer._default_camera_pose[:] = pose 180 | 181 | def _create_raymond_lights(self): 182 | thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) 183 | phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) 184 | 185 | nodes = [] 186 | 187 | for phi, theta in zip(phis, thetas): 188 | xp = np.sin(theta) * np.cos(phi) 189 | yp = np.sin(theta) * np.sin(phi) 190 | zp = np.cos(theta) 191 | 192 | z = np.array([xp, yp, zp]) 193 | z = z / np.linalg.norm(z) 194 | x = np.array([-z[1], z[0], 0.0]) 195 | if np.linalg.norm(x) == 0: 196 | x = np.array([1.0, 0.0, 0.0]) 197 | x = x / np.linalg.norm(x) 198 | y = np.cross(z, x) 199 | 200 | matrix = np.eye(4) 201 | matrix[:3, :3] = np.c_[x, y, z] 202 | nodes.append( 203 | pyrender.Node( 204 | light=pyrender.DirectionalLight(color=np.ones(3), 205 | intensity=1.0), 206 | matrix=matrix 207 | )) 208 | 209 | return nodes 210 | 211 | 212 | def set_meshes(self, meshes =[], set_type = 'static'): 213 | 214 | if not self.offscreen: 215 | self.viewer.render_lock.acquire() 216 | 217 | for node in self.scene.get_nodes(): 218 | if node.name is None: 219 | continue 220 | if 'static' in set_type and 'mesh' in node.name: 221 | self.scene.remove_node(node) 222 | elif 'dynamic' in node.name: 223 | self.scene.remove_node(node) 224 | 225 | for i, mesh in enumerate(meshes): 226 | mesh = self.to_pymesh(mesh) 227 | self.scene.add(mesh, name='%s_mesh_%d'%(set_type,i)) 228 | 229 | if not self.offscreen: 230 | self.viewer.render_lock.release() 231 | 232 | def set_static_meshes(self, meshes =[]): 233 | self.set_meshes(meshes=meshes, set_type='static') 234 | 235 | def set_dynamic_meshes(self, meshes =[]): 236 | self.set_meshes(meshes=meshes, set_type='dynamic') 237 | 238 | def save_snapshot(self, save_path): 239 | if not self.offscreen: 240 | print('We do not support rendering in Interactive mode!') 241 | return 242 | color, depth = self.viewer.render(self.scene) 243 | img = Image.fromarray(color) 244 | img.save(save_path) 245 | 246 | 247 | def points2sphere(points, radius = .001, vc = [0., 0., 1.], count = [5,5]): 248 | 249 | points = points.reshape(-1,3) 250 | n_points = points.shape[0] 251 | 252 | spheres = [] 253 | for p in range(n_points): 254 | sphs = trimesh.creation.uv_sphere(radius=radius, count = count) 255 | sphs.apply_translation(points[p]) 256 | sphs = Mesh(vertices=sphs.vertices, faces=sphs.faces, vc=vc) 257 | 258 | spheres.append(sphs) 259 | 260 | spheres = Mesh.concatenate_meshes(spheres) 261 | return spheres 262 | 263 | colors = { 264 | 'pink': [1.00, 0.75, 0.80], 265 | 'purple': [0.63, 0.13, 0.94], 266 | 'red': [1.0, 0.0, 0.0], 267 | 'green': [.0, 1., .0], 268 | 'yellow': [1., 1., 0], 269 | 'brown': [1.00, 0.25, 0.25], 270 | 'blue': [.0, .0, 1.], 271 | 'white': [1., 1., 1.], 272 | 'orange': [1.00, 0.65, 0.00], 273 | 'grey': [0.75, 0.75, 0.75], 274 | 'black': [0., 0., 0.], 275 | } -------------------------------------------------------------------------------- /visualizer/tools/meshviewer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2019 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems and the Max Planck Institute for Biological 14 | # Cybernetics. All rights reserved. 15 | # 16 | # Contact: ps-license@tuebingen.mpg.de 17 | # 18 | 19 | import numpy as np 20 | import trimesh 21 | import pyrender 22 | from pyrender.light import DirectionalLight 23 | from pyrender.node import Node 24 | from PIL import Image 25 | from .utils import euler 26 | 27 | 28 | class Mesh(trimesh.Trimesh): 29 | 30 | def __init__(self, 31 | filename=None, 32 | vertices=None, 33 | faces=None, 34 | vc=None, 35 | fc=None, 36 | vscale=None, 37 | process = False, 38 | visual = None, 39 | wireframe=False, 40 | smooth = False, 41 | **kwargs): 42 | 43 | self.wireframe = wireframe 44 | self.smooth = smooth 45 | 46 | if filename is not None: 47 | mesh = trimesh.load(filename, process = process) 48 | vertices = mesh.vertices 49 | faces= mesh.faces 50 | visual = mesh.visual 51 | if vscale is not None: 52 | vertices = vertices*vscale 53 | 54 | if faces is None: 55 | mesh = points2sphere(vertices) 56 | vertices = mesh.vertices 57 | faces = mesh.faces 58 | visual = mesh.visual 59 | 60 | super(Mesh, self).__init__(vertices=vertices, faces=faces, process=process, visual=visual) 61 | 62 | if vc is not None: 63 | self.set_vertex_colors(vc) 64 | if fc is not None: 65 | self.set_face_colors(fc) 66 | 67 | def rot_verts(self, vertices, rxyz): 68 | return np.array(vertices * rxyz.T) 69 | 70 | def colors_like(self,color, array, ids): 71 | 72 | color = np.array(color) 73 | 74 | if color.max() <= 1.: 75 | color = color * 255 76 | color = color.astype(np.int8) 77 | 78 | n_color = color.shape[0] 79 | n_ids = ids.shape[0] 80 | 81 | new_color = np.array(array) 82 | if n_color <= 4: 83 | new_color[ids, :n_color] = np.repeat(color[np.newaxis], n_ids, axis=0) 84 | else: 85 | new_color[ids, :] = color 86 | 87 | return new_color 88 | 89 | def set_vertex_colors(self,vc, vertex_ids = None): 90 | 91 | all_ids = np.arange(self.vertices.shape[0]) 92 | if vertex_ids is None: 93 | vertex_ids = all_ids 94 | 95 | vertex_ids = all_ids[vertex_ids] 96 | new_vc = self.colors_like(vc, self.visual.vertex_colors, vertex_ids) 97 | self.visual.vertex_colors[:] = new_vc 98 | 99 | def set_face_colors(self,fc, face_ids = None): 100 | 101 | if face_ids is None: 102 | face_ids = np.arange(self.faces.shape[0]) 103 | 104 | new_fc = self.colors_like(fc, self.visual.face_colors, face_ids) 105 | self.visual.face_colors[:] = new_fc 106 | 107 | @staticmethod 108 | def concatenate_meshes(meshes): 109 | return trimesh.util.concatenate(meshes) 110 | 111 | 112 | class MeshViewer(object): 113 | 114 | def __init__(self, 115 | width=1200, 116 | height=800, 117 | bg_color = [0.0, 0.0, 0.0, 1.0], 118 | offscreen = False, 119 | registered_keys=None): 120 | super(MeshViewer, self).__init__() 121 | 122 | if registered_keys is None: 123 | registered_keys = dict() 124 | 125 | self.bg_color = bg_color 126 | self.offscreen = offscreen 127 | self.scene = pyrender.Scene(bg_color=bg_color, 128 | ambient_light=(0.3, 0.3, 0.3), 129 | name = 'scene') 130 | 131 | self.aspect_ratio = float(width) / height 132 | pc = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=self.aspect_ratio) 133 | camera_pose = np.eye(4) 134 | camera_pose[:3,:3] = euler([80,-15,0], 'xzx') 135 | camera_pose[:3, 3] = np.array([-.5, -2., 1.5]) 136 | 137 | self.cam = pyrender.Node(name = 'camera', camera=pc, matrix=camera_pose) 138 | 139 | self.scene.add_node(self.cam) 140 | 141 | if self.offscreen: 142 | light = Node(light=DirectionalLight(color=np.ones(3), intensity=3.0), 143 | matrix=camera_pose) 144 | self.scene.add_node(light) 145 | self.viewer = pyrender.OffscreenRenderer(width, height) 146 | else: 147 | self.viewer = pyrender.Viewer(self.scene, 148 | use_raymond_lighting=True, 149 | viewport_size=(width, height), 150 | cull_faces=False, 151 | run_in_thread=True, 152 | registered_keys=registered_keys) 153 | 154 | for i, node in enumerate(self.scene.get_nodes()): 155 | if node.name is None: 156 | node.name = 'Req%d'%i 157 | 158 | 159 | def is_active(self): 160 | return self.viewer.is_active 161 | 162 | def close_viewer(self): 163 | if not self.offscreen and self.viewer.is_active: 164 | self.viewer.close_external() 165 | 166 | def set_background_color(self, bg_color=[1., 1., 1.]): 167 | self.scene.bg_color = bg_color 168 | 169 | def to_pymesh(self, mesh): 170 | 171 | wireframe = mesh.wireframe if hasattr(mesh, 'wireframe') else False 172 | smooth = mesh.smooth if hasattr(mesh, 'smooth') else False 173 | return pyrender.Mesh.from_trimesh(mesh, wireframe=wireframe, smooth=smooth) 174 | 175 | def update_camera_pose(self, pose): 176 | if self.offscreen: 177 | self.scene.set_pose(self.cam, pose=pose) 178 | else: 179 | self.viewer._default_camera_pose[:] = pose 180 | 181 | def _create_raymond_lights(self): 182 | thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) 183 | phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) 184 | 185 | nodes = [] 186 | 187 | for phi, theta in zip(phis, thetas): 188 | xp = np.sin(theta) * np.cos(phi) 189 | yp = np.sin(theta) * np.sin(phi) 190 | zp = np.cos(theta) 191 | 192 | z = np.array([xp, yp, zp]) 193 | z = z / np.linalg.norm(z) 194 | x = np.array([-z[1], z[0], 0.0]) 195 | if np.linalg.norm(x) == 0: 196 | x = np.array([1.0, 0.0, 0.0]) 197 | x = x / np.linalg.norm(x) 198 | y = np.cross(z, x) 199 | 200 | matrix = np.eye(4) 201 | matrix[:3, :3] = np.c_[x, y, z] 202 | nodes.append( 203 | pyrender.Node( 204 | light=pyrender.DirectionalLight(color=np.ones(3), 205 | intensity=1.0), 206 | matrix=matrix 207 | )) 208 | 209 | return nodes 210 | 211 | 212 | def set_meshes(self, meshes =[], set_type = 'static'): 213 | 214 | if not self.offscreen: 215 | self.viewer.render_lock.acquire() 216 | 217 | for node in self.scene.get_nodes(): 218 | if node.name is None: 219 | continue 220 | if 'static' in set_type and 'mesh' in node.name: 221 | self.scene.remove_node(node) 222 | elif 'dynamic' in node.name: 223 | self.scene.remove_node(node) 224 | 225 | for i, mesh in enumerate(meshes): 226 | mesh = self.to_pymesh(mesh) 227 | self.scene.add(mesh, name='%s_mesh_%d'%(set_type,i)) 228 | 229 | if not self.offscreen: 230 | self.viewer.render_lock.release() 231 | 232 | def set_static_meshes(self, meshes =[]): 233 | self.set_meshes(meshes=meshes, set_type='static') 234 | 235 | def set_dynamic_meshes(self, meshes =[]): 236 | self.set_meshes(meshes=meshes, set_type='dynamic') 237 | 238 | def save_snapshot(self, save_path): 239 | if not self.offscreen: 240 | print('We do not support rendering in Interactive mode!') 241 | return 242 | color, depth = self.viewer.render(self.scene) 243 | img = Image.fromarray(color) 244 | img.save(save_path) 245 | 246 | 247 | def points2sphere(points, radius = .001, vc = [0., 0., 1.], count = [5,5]): 248 | 249 | points = points.reshape(-1,3) 250 | n_points = points.shape[0] 251 | 252 | spheres = [] 253 | for p in range(n_points): 254 | sphs = trimesh.creation.uv_sphere(radius=radius, count = count) 255 | sphs.apply_translation(points[p]) 256 | sphs = Mesh(vertices=sphs.vertices, faces=sphs.faces, vc=vc) 257 | 258 | spheres.append(sphs) 259 | 260 | spheres = Mesh.concatenate_meshes(spheres) 261 | return spheres 262 | 263 | colors = { 264 | 'pink': [1.00, 0.75, 0.80], 265 | 'purple': [0.63, 0.13, 0.94], 266 | 'red': [1.0, 0.0, 0.0], 267 | 'green': [.0, 1., .0], 268 | 'yellow': [1., 1., 0], 269 | 'brown': [1.00, 0.25, 0.25], 270 | 'blue': [.0, .0, 1.], 271 | 'white': [1., 1., 1.], 272 | 'orange': [1.00, 0.65, 0.00], 273 | 'grey': [0.75, 0.75, 0.75], 274 | 'black': [0., 0., 0.], 275 | } -------------------------------------------------------------------------------- /visualizer/vis/visualize_interaction_results.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 4 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 5 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 6 | # 7 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 8 | # on this computer program. You can only use this computer program if you have closed a license agreement 9 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 10 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 11 | # Contact: ps-license@tuebingen.mpg.de 12 | # 13 | import sys 14 | 15 | sys.path.append(".") 16 | sys.path.append("..") 17 | 18 | import argparse 19 | import glob 20 | import os 21 | 22 | import numpy as np 23 | import smplx 24 | import torch 25 | from tqdm import tqdm 26 | 27 | os.environ["PYOPENGL_PLATFORM"] = "egl" 28 | 29 | import trimesh 30 | 31 | from visualizer.tools.cfg_parser import Config 32 | from visualizer.tools.meshviewer import Mesh, MeshViewer, colors, points2sphere 33 | from visualizer.tools.objectmodel import ObjectModel 34 | from visualizer.tools.utils import euler, params2torch, parse_npz, to_cpu 35 | 36 | USE_FLAT_HAND_MEAN = True 37 | 38 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 39 | 40 | 41 | def visualize_sequences(cfg): 42 | result_path = cfg.result_path 43 | 44 | mv = MeshViewer(offscreen=cfg.offscreen) 45 | 46 | # set the camera pose 47 | camera_pose = np.eye(4) 48 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 49 | camera_pose[:3, 3] = np.array([-0.5, -5.0, 1.5]) 50 | mv.update_camera_pose(camera_pose) 51 | 52 | vis_sequence(cfg, result_path, mv, cfg.interaction_epoch, cfg.s_idx) 53 | mv.close_viewer() 54 | 55 | 56 | def vis_sequence(cfg, result_path, mv, interaction_epoch, s_idx): 57 | # 1 pass, process all data 58 | ball_meshes = [] 59 | human_meshes = [] 60 | 61 | b_meshes = [] 62 | s_meshes = [] 63 | o_meshes = [] 64 | b_meshes_gt = [] 65 | s_meshes_gt = [] 66 | o_meshes_gt = [] 67 | 68 | ball_path = os.path.join( 69 | result_path, "ball_objs_step_{}_bs_idx_0".format(interaction_epoch) 70 | ) 71 | subject_path = os.path.join( 72 | result_path, "objs_step_{}_bs_idx_0".format(interaction_epoch) 73 | ) 74 | 75 | ori_ball_files = os.listdir(ball_path) 76 | ori_ball_files.sort() 77 | 78 | for tmp_name in ori_ball_files: 79 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 80 | if ".ply" in tmp_name: 81 | if "start" in tmp_name: 82 | continue 83 | ball_meshes.append(trimesh.load(os.path.join(ball_path, tmp_name))) 84 | b_meshes.append( 85 | Mesh( 86 | vertices=ball_meshes[-1].vertices, 87 | faces=ball_meshes[-1].faces, 88 | vc=colors["green"], 89 | ) 90 | ) 91 | ori_obj_files = os.listdir(subject_path) 92 | ori_obj_files.sort() 93 | for tmp_name in ori_obj_files: 94 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 95 | if ".ply" in tmp_name: 96 | if "object" not in tmp_name: 97 | human_meshes.append(trimesh.load(os.path.join(subject_path, tmp_name))) 98 | s_meshes.append( 99 | Mesh( 100 | vertices=human_meshes[-1].vertices, 101 | faces=human_meshes[-1].faces, 102 | vc=colors["pink"], 103 | smooth=True, 104 | ) 105 | ) 106 | else: 107 | o_mesh = trimesh.load(os.path.join(subject_path, tmp_name)) 108 | o_meshes.append( 109 | Mesh( 110 | vertices=o_mesh.vertices, 111 | faces=o_mesh.faces, 112 | vc=colors["pink"], 113 | smooth=True, 114 | ) 115 | ) 116 | 117 | # gt 118 | ball_path_gt = os.path.join( 119 | result_path, "ball_objs_step_{}_bs_idx_0_gt".format(interaction_epoch) 120 | ) 121 | subject_path_gt = os.path.join( 122 | result_path, "objs_step_{}_bs_idx_0_gt".format(interaction_epoch) 123 | ) 124 | vis_gt = os.path.exists(subject_path_gt) and len(os.listdir(subject_path_gt)) > 0 125 | if vis_gt: 126 | offset = np.array([2, 0, 0]) 127 | ori_ball_files_gt = os.listdir(ball_path_gt) 128 | ori_ball_files_gt.sort() 129 | for tmp_name in ori_ball_files_gt: 130 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 131 | if ".ply" in tmp_name: 132 | if "start" in tmp_name: 133 | continue 134 | mesh = trimesh.load(os.path.join(ball_path_gt, tmp_name)) 135 | new_v = mesh.vertices + offset 136 | b_meshes_gt.append( 137 | Mesh(vertices=new_v, faces=mesh.faces, vc=colors["green"]) 138 | ) 139 | ori_obj_files_gt = os.listdir(subject_path_gt) 140 | ori_obj_files_gt.sort() 141 | for tmp_name in ori_obj_files_gt: 142 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 143 | if ".ply" in tmp_name: 144 | if "object" not in tmp_name: 145 | mesh = trimesh.load(os.path.join(subject_path_gt, tmp_name)) 146 | new_v = mesh.vertices + offset 147 | s_meshes_gt.append( 148 | Mesh( 149 | vertices=new_v, 150 | faces=mesh.faces, 151 | vc=colors["yellow"], 152 | smooth=True, 153 | ) 154 | ) 155 | else: 156 | mesh = trimesh.load(os.path.join(subject_path_gt, tmp_name)) 157 | new_v = mesh.vertices + offset 158 | o_meshes_gt.append( 159 | Mesh( 160 | vertices=new_v, 161 | faces=mesh.faces, 162 | vc=colors["yellow"], 163 | smooth=True, 164 | ) 165 | ) 166 | 167 | # 2 pass, render the data 168 | skip_frame = 1 169 | if not cfg.offscreen: 170 | while True: 171 | import time 172 | 173 | T = len(s_meshes) 174 | for frame in range(0, T, skip_frame): 175 | start_time = time.time() 176 | 177 | s_mesh = s_meshes[frame] 178 | o_mesh = o_meshes[frame] 179 | meshes = [s_mesh, o_mesh] 180 | 181 | meshes.extend(b_meshes) 182 | if vis_gt: 183 | s_mesh_gt = s_meshes_gt[frame] 184 | o_mesh_gt = o_meshes_gt[frame] 185 | meshes.extend([s_mesh_gt, o_mesh_gt]) 186 | meshes.extend(b_meshes_gt) 187 | 188 | mv.set_static_meshes(meshes) 189 | while time.time() - start_time < 0.03: 190 | pass 191 | else: 192 | import imageio 193 | 194 | img_paths = [] 195 | 196 | path1, path2 = result_path.split("/")[-2], result_path.split("/")[-1] 197 | save_dir = os.path.join(result_path, "../..", "vis") 198 | if not os.path.exists(save_dir): 199 | os.makedirs(save_dir) 200 | save_dir = os.path.join(save_dir, path1) 201 | if not os.path.exists(save_dir): 202 | os.makedirs(save_dir) 203 | 204 | img_path = os.path.join(result_path, "vis") 205 | if not os.path.exists(img_path): 206 | os.makedirs(img_path) 207 | 208 | T = len(s_meshes) 209 | for frame in tqdm(range(0, T, skip_frame), leave=False): 210 | s_mesh = s_meshes[frame] 211 | o_mesh = o_meshes[frame] 212 | meshes = [s_mesh, o_mesh] 213 | 214 | meshes.extend(b_meshes) 215 | if vis_gt: 216 | s_mesh_gt = s_meshes_gt[frame] 217 | o_mesh_gt = o_meshes_gt[frame] 218 | meshes.extend([s_mesh_gt, o_mesh_gt]) 219 | meshes.extend(b_meshes_gt) 220 | 221 | mv.set_static_meshes(meshes) 222 | 223 | camera_pose = np.eye(4) 224 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 225 | camera_pose[:2, 3] = np.mean(s_mesh.vertices, axis=0)[:2] + np.array( 226 | [0.5, -3.0] 227 | ) 228 | camera_pose[2, 3] = 1.3 229 | mv.update_camera_pose(camera_pose) 230 | 231 | mv.save_snapshot(os.path.join(img_path, "%05d.png" % frame)) 232 | img_paths.append(os.path.join(img_path, "%05d.png" % frame)) 233 | 234 | video_name = os.path.join(save_dir, "output_{}.mp4".format(s_idx)) 235 | 236 | im_arr = [] 237 | 238 | for image in img_paths: 239 | im = imageio.v2.imread(image) 240 | im_arr.append(im) 241 | im_arr = np.asarray(im_arr) 242 | imageio.mimwrite(video_name, im_arr, fps=30, quality=8) 243 | 244 | print("video saved to %s" % video_name) 245 | 246 | 247 | if __name__ == "__main__": 248 | parser = argparse.ArgumentParser(description="visualize") 249 | 250 | parser.add_argument( 251 | "--result-path", required=True, type=str, help="The path to the results" 252 | ) 253 | parser.add_argument("--s-idx", required=True, type=str) 254 | parser.add_argument("--interaction-epoch", required=True, type=int) 255 | parser.add_argument( 256 | "--model-path", 257 | required=True, 258 | type=str, 259 | help="The path to the folder containing smplx models", 260 | ) 261 | parser.add_argument("--offscreen", action="store_true") 262 | 263 | args = parser.parse_args() 264 | 265 | result_path = args.result_path 266 | model_path = args.model_path 267 | 268 | cfg = { 269 | "result_path": result_path, 270 | "model_path": model_path, 271 | "offscreen": args.offscreen, 272 | "s_idx": args.s_idx, 273 | "interaction_epoch": args.interaction_epoch, 274 | } 275 | 276 | cfg = Config(**cfg) 277 | visualize_sequences(cfg) 278 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: hoifhli_env 2 | channels: 3 | - pytorch3d 4 | - bottler 5 | - iopath 6 | - fvcore 7 | - pytorch 8 | - conda-forge 9 | - defaults 10 | dependencies: 11 | - _libgcc_mutex=0.1=conda_forge 12 | - _openmp_mutex=4.5=2_gnu 13 | - asttokens=2.4.1=pyhd8ed1ab_0 14 | - backcall=0.2.0=pyh9f0ad1d_0 15 | - blas=1.0=mkl 16 | - bzip2=1.0.8=h7b6447c_0 17 | - ca-certificates=2024.8.30=hbcca054_0 18 | - certifi=2024.8.30=pyhd8ed1ab_0 19 | - colorama=0.4.6=pyhd8ed1ab_0 20 | - cudatoolkit=11.3.1=h2bc3f7f_2 21 | - executing=2.0.1=pyhd8ed1ab_0 22 | - ffmpeg=4.3=hf484d3e_0 23 | - freetype=2.12.1=h4a9f257_0 24 | - fvcore=0.1.5.post20210915=py38 25 | - giflib=5.2.1=h5eee18b_3 26 | - gmp=6.2.1=h295c915_3 27 | - gnutls=3.6.15=he1e5248_0 28 | - importlib-metadata=7.0.0=pyha770c72_0 29 | - importlib_metadata=7.0.0=hd8ed1ab_0 30 | - intel-openmp=2023.1.0=hdb19cb5_46306 31 | - iopath=0.1.9=py38 32 | - jedi=0.19.1=pyhd8ed1ab_0 33 | - jpeg=9e=h5eee18b_1 34 | - jupyter_client=8.6.0=pyhd8ed1ab_0 35 | - jupyter_core=5.5.0=py38h578d9bd_0 36 | - lame=3.100=h7b6447c_0 37 | - lcms2=2.12=h3be6417_0 38 | - ld_impl_linux-64=2.38=h1181459_1 39 | - lerc=3.0=h295c915_0 40 | - libdeflate=1.17=h5eee18b_1 41 | - libffi=3.4.4=h6a678d5_0 42 | - libgcc=14.2.0=h77fa898_1 43 | - libgcc-ng=14.2.0=h69a702a_1 44 | - libgomp=14.2.0=h77fa898_1 45 | - libiconv=1.16=h7f8727e_2 46 | - libidn2=2.3.4=h5eee18b_0 47 | - libpng=1.6.39=h5eee18b_0 48 | - libsodium=1.0.18=h36c2ea0_1 49 | - libspatialindex=1.9.3=h2531618_0 50 | - libstdcxx-ng=11.2.0=h1234567_1 51 | - libtasn1=4.19.0=h5eee18b_0 52 | - libtiff=4.5.1=h6a678d5_0 53 | - libunistring=0.9.10=h27cfd23_0 54 | - libuv=1.44.2=h5eee18b_0 55 | - libwebp=1.3.2=h11a3e52_0 56 | - libwebp-base=1.3.2=h5eee18b_0 57 | - lz4-c=1.9.4=h6a678d5_0 58 | - matplotlib-inline=0.1.6=pyhd8ed1ab_0 59 | - mkl=2023.1.0=h213fc3f_46344 60 | - mkl-service=2.4.0=py38h5eee18b_1 61 | - mkl_fft=1.3.8=py38h5eee18b_0 62 | - mkl_random=1.2.4=py38hdb19cb5_0 63 | - ncurses=6.4=h6a678d5_0 64 | - nest-asyncio=1.5.8=pyhd8ed1ab_0 65 | - nettle=3.7.3=hbbd107a_1 66 | - nvidiacub=1.10.0=0 67 | - openh264=2.1.1=h4ff587b_0 68 | - openjpeg=2.4.0=h3ad879b_0 69 | - openssl=3.3.2=hb9d3cd8_0 70 | - packaging=23.2=pyhd8ed1ab_0 71 | - parso=0.8.3=pyhd8ed1ab_0 72 | - pickleshare=0.7.5=py_1003 73 | - platformdirs=4.1.0=pyhd8ed1ab_0 74 | - portalocker=2.8.2=py38h578d9bd_1 75 | - prompt-toolkit=3.0.41=pyha770c72_0 76 | - prompt_toolkit=3.0.41=hd8ed1ab_0 77 | - psutil=5.9.5=py38h01eb140_1 78 | - ptyprocess=0.7.0=pyhd3deb0d_0 79 | - pure_eval=0.2.2=pyhd8ed1ab_0 80 | - pygments=2.17.2=pyhd8ed1ab_0 81 | - python=3.8.18=h955ad1f_0 82 | - python-dateutil=2.8.2=pyhd8ed1ab_0 83 | - python_abi=3.8=2_cp38 84 | - pytorch-mutex=1.0=cuda 85 | - pytorch3d=0.7.4=py38_cu113_pyt1110 86 | - pyyaml=6.0.1=py38h01eb140_1 87 | - readline=8.2=h5eee18b_0 88 | - rtree=1.0.1=py38h06a4308_0 89 | - six=1.16.0=pyh6c4a22f_0 90 | - sqlite=3.41.2=h5eee18b_0 91 | - stack_data=0.6.2=pyhd8ed1ab_0 92 | - tabulate=0.9.0=pyhd8ed1ab_1 93 | - tbb=2021.8.0=hdb19cb5_0 94 | - termcolor=2.3.0=pyhd8ed1ab_0 95 | - tk=8.6.12=h1ccaba5_0 96 | - torchaudio=0.11.0=py38_cu113 97 | - tqdm=4.66.1=pyhd8ed1ab_0 98 | - traitlets=5.14.0=pyhd8ed1ab_0 99 | - transforms3d=0.4.1=pyhd8ed1ab_0 100 | - typing_extensions=4.12.2=pyha770c72_0 101 | - wcwidth=0.2.12=pyhd8ed1ab_0 102 | - xz=5.4.5=h5eee18b_0 103 | - yacs=0.1.8=pyhd8ed1ab_0 104 | - yaml=0.2.5=h7f98852_2 105 | - zeromq=4.3.4=h9c3ff4c_1 106 | - zipp=3.17.0=pyhd8ed1ab_0 107 | - zlib=1.2.13=h5eee18b_0 108 | - zstd=1.5.5=hc292b87_0 109 | - pip: 110 | - absl-py==2.1.0 111 | - addict==2.4.0 112 | - annotated-types==0.7.0 113 | - antlr4-python3-runtime==4.9.3 114 | - anyio==4.1.0 115 | - argcomplete==3.5.1 116 | - argon2-cffi==23.1.0 117 | - argon2-cffi-bindings==21.2.0 118 | - argunparse==0.1.4 119 | - arrow==1.3.0 120 | - async-lru==2.0.4 121 | - attrs==23.1.0 122 | - babel==2.13.1 123 | - beautifulsoup4==4.12.2 124 | - bleach==6.1.0 125 | - blinker==1.7.0 126 | - boilerplates==1.1.1 127 | # - bps-torch==0.1 128 | - cachetools==5.3.2 129 | # - chamfer-distance==0.1 130 | - charset-normalizer==3.4.0 131 | # - chumpy==0.67.5 132 | - click==8.1.7 133 | - git+https://github.com/openai/CLIP.git 134 | # - git+https://github.com/vchoutas/torch-sdf.git 135 | - cmake==3.24.0 136 | - coacd==1.0.0 137 | - colorlog==6.8.2 138 | - colour==0.1.5 139 | - comm==0.2.0 140 | - common==0.1.2 141 | - configargparse==1.7 142 | - contourpy==1.1.1 143 | - cycler==0.12.1 144 | - dash==2.15.0 145 | - dash-core-components==2.0.0 146 | - dash-html-components==2.0.0 147 | - dash-table==5.0.0 148 | - debugpy==1.8.0 149 | - decorator==4.4.2 150 | - defusedxml==0.7.1 151 | - distro==1.9.0 152 | - docker-pycreds==0.4.0 153 | - dotmap==1.3.30 154 | - easydict==1.11 155 | - einops==0.4.1 156 | - ema-pytorch==0.0.10 157 | - et-xmlfile==1.1.0 158 | - exceptiongroup==1.2.0 159 | - fastjsonschema==2.19.0 160 | - filelock==3.16.1 161 | - flameprof==0.4 162 | - flask==3.0.2 163 | - fonttools==4.46.0 164 | - fqdn==1.5.1 165 | - freetype-py==2.4.0 166 | - fsspec==2024.9.0 167 | - ftfy==6.1.3 168 | - gitdb==4.0.11 169 | - gitpython==3.1.40 170 | - google-auth==2.27.0 171 | - google-auth-oauthlib==1.0.0 172 | - grpcio==1.60.0 173 | - h11==0.14.0 174 | - httpcore==1.0.5 175 | - httpx==0.27.0 176 | - idna==3.10 177 | - imageio==2.19.3 178 | - imageio-ffmpeg==0.4.7 179 | - importlib-resources==6.1.1 180 | - ipdb==0.13.13 181 | - ipykernel==6.27.1 182 | - ipython==8.12.3 183 | - ipywidgets==8.1.1 184 | - isoduration==20.11.0 185 | - itsdangerous==2.1.2 186 | - jinja2==3.1.4 187 | - jiter==0.7.0 188 | - joblib==1.3.2 189 | - json5==0.9.14 190 | - jsonpointer==2.4 191 | - jsonschema==4.20.0 192 | - jsonschema-specifications==2023.11.2 193 | - jupyter==1.0.0 194 | - jupyter-console==6.6.3 195 | - jupyter-events==0.9.0 196 | - jupyter-lsp==2.2.1 197 | - jupyter-server==2.11.2 198 | - jupyter-server-terminals==0.4.4 199 | - jupyterlab==4.0.9 200 | - jupyterlab-pygments==0.3.0 201 | - jupyterlab-server==2.25.2 202 | - jupyterlab-widgets==3.0.9 203 | - kiwisolver==1.4.5 204 | - lazy-loader==0.3 205 | - line-profiler==4.1.3 206 | - loguru==0.7.2 207 | - lxml==5.1.0 208 | # - git+https://github.com/hassony2/manopth.git 209 | - markdown==3.5.2 210 | - markupsafe>=2.1.1 211 | - matplotlib==3.2.2 212 | - meshlib==2.4.4.153 213 | - mistune==3.0.2 214 | - moviepy==1.0.3 215 | - mpmath==1.3.0 216 | - nbclient==0.9.0 217 | - nbconvert==7.12.0 218 | - nbformat==5.9.2 219 | - networkx==3.1 220 | - ninja==1.11.1.1 221 | - notebook==7.0.6 222 | - notebook-shim==0.2.3 223 | - numpy==1.24.4 224 | - nvidia-cublas-cu12==12.1.3.1 225 | - nvidia-cuda-cupti-cu12==12.1.105 226 | - nvidia-cuda-nvrtc-cu12==12.1.105 227 | - nvidia-cuda-runtime-cu12==12.1.105 228 | - nvidia-cudnn-cu12==9.1.0.70 229 | - nvidia-cufft-cu12==11.0.2.54 230 | - nvidia-curand-cu12==10.3.2.106 231 | - nvidia-cusolver-cu12==11.4.5.107 232 | - nvidia-cusparse-cu12==12.1.0.106 233 | - nvidia-nccl-cu12==2.20.5 234 | - nvidia-nvjitlink-cu12==12.6.77 235 | - nvidia-nvtx-cu12==12.1.105 236 | - oauthlib==3.2.2 237 | - omegaconf==2.3.0 238 | - open3d==0.18.0 239 | - openai==1.54.3 240 | - opencv-python==4.5.1.48 241 | - openpyxl==3.1.5 242 | - overrides==7.4.0 243 | - pandas==2.0.3 244 | - pandocfilters==1.5.0 245 | - pathfinding==1.0.9 246 | - pathtools==0.1.2 247 | - pexpect==4.9.0 248 | - pillow==10.2.0 249 | - pip==24.3.1 250 | - pkgutil-resolve-name==1.3.10 251 | - plotly==5.18.0 252 | - prettytable==3.11.0 253 | - proglog==0.1.10 254 | - prometheus-client==0.19.0 255 | - promise==2.3 256 | - protobuf==3.20.3 257 | # - psbody-mesh==0.4 258 | - pyarrow==15.0.2 259 | - pyasn1==0.5.1 260 | - pyasn1-modules==0.3.0 261 | - pydantic==2.9.2 262 | - pydantic-core==2.23.4 263 | - pyglet==2.0.10 264 | - pymeshlab==2023.12.post2 265 | - pyopengl==3.1.0 266 | - pyopengl-accelerate==3.1.7 267 | - pyparsing==3.1.1 268 | - pyquaternion==0.9.9 269 | - pyrender==0.1.43 270 | - pysdf==0.1.9 271 | - python-json-logger==2.0.7 272 | - pytz==2023.3.post1 273 | - pywavelets==1.4.1 274 | - pyzmq==25.1.2 275 | - qtconsole==5.5.1 276 | - qtpy==2.4.1 277 | - referencing==0.31.1 278 | - regex==2023.12.25 279 | - requests==2.32.3 280 | - requests-oauthlib==1.3.1 281 | - rerun-sdk==0.14.1 282 | - retrying==1.3.4 283 | - rfc3339-validator==0.1.4 284 | - rfc3986-validator==0.1.1 285 | - rpds-py==0.13.2 286 | - rsa==4.9 287 | - scenepic==1.1.1 288 | - scikit-image==0.20.0 289 | - scikit-learn==1.3.2 290 | - scipy==1.9.1 291 | - semver==3.0.2 292 | - send2trash==1.8.2 293 | - sentry-sdk==1.38.0 294 | - setproctitle==1.3.3 295 | - setuptools==75.2.0 296 | - shapely==2.0.4 297 | - shortuuid==1.0.11 298 | - smmap==5.0.1 299 | - sniffio==1.3.0 300 | - soupsieve==2.5 301 | - stack-data==0.6.3 302 | - sympy==1.13.3 303 | - tenacity==8.2.3 304 | - tensorboard==2.14.0 305 | - tensorboard-data-server==0.7.2 306 | - tensorboardx==2.6.2.2 307 | - terminado==0.18.0 308 | - threadpoolctl==3.2.0 309 | - tifffile==2023.7.10 310 | - tinycss2==1.2.1 311 | - tomli==2.0.1 312 | - torch==1.11.0 313 | - torchgeometry==0.1.2 314 | - torchvision==0.12.0 315 | - tornado==6.4 316 | - transformations==2022.9.26 317 | - trimesh==3.9.5 318 | - triton==3.0.0 319 | - types-python-dateutil==2.8.19.14 320 | - tzdata==2024.1 321 | - uri-template==1.3.0 322 | - urllib3==2.2.3 323 | - version-query==1.5.5 324 | - wandb==0.12.21 325 | - webcolors==1.13 326 | - webencodings==0.5.1 327 | - websocket-client==1.7.0 328 | - werkzeug==3.0.1 329 | - wheel==0.44.0 330 | - widgetsnbextension==4.0.9 331 | -------------------------------------------------------------------------------- /visualizer/vis/visualize_long_sequence_results.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG), 4 | # acting on behalf of its Max Planck Institute for Intelligent Systems and the 5 | # Max Planck Institute for Biological Cybernetics. All rights reserved. 6 | # 7 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights 8 | # on this computer program. You can only use this computer program if you have closed a license agreement 9 | # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right. 10 | # Any use of the computer program without a valid license is prohibited and liable to prosecution. 11 | # Contact: ps-license@tuebingen.mpg.de 12 | # 13 | import sys 14 | 15 | sys.path.append(".") 16 | sys.path.append("..") 17 | 18 | import argparse 19 | import glob 20 | import os 21 | 22 | import numpy as np 23 | import smplx 24 | import torch 25 | from tqdm import tqdm 26 | 27 | os.environ["PYOPENGL_PLATFORM"] = "egl" 28 | 29 | import trimesh 30 | 31 | from visualizer.tools.cfg_parser import Config 32 | from visualizer.tools.meshviewer import Mesh, MeshViewer, colors, points2sphere 33 | from visualizer.tools.objectmodel import ObjectModel 34 | from visualizer.tools.utils import euler, params2torch, parse_npz, to_cpu 35 | 36 | USE_FLAT_HAND_MEAN = True 37 | 38 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 39 | 40 | 41 | def visualize_sequences(cfg): 42 | result_path = cfg.result_path 43 | 44 | mv = MeshViewer(offscreen=cfg.offscreen) 45 | 46 | # set the camera pose 47 | camera_pose = np.eye(4) 48 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 49 | camera_pose[:3, 3] = np.array([-0.5, -3.0, 1.5]) 50 | mv.update_camera_pose(camera_pose) 51 | 52 | vis_sequence(cfg, result_path, mv, cfg.interaction_epoch, cfg.s_idx, cfg.video_idx) 53 | mv.close_viewer() 54 | 55 | 56 | def vis_sequence(cfg, result_path_all, mv, interaction_epoch, s_idx, video_idx): 57 | # process the initial object mesh 58 | initial_obj_meshes = [] 59 | initial_obj_paths = cfg.initial_obj_path.split("&") 60 | for initial_obj_path in initial_obj_paths: 61 | if os.path.exists(initial_obj_path): 62 | initial_obj_mesh = trimesh.load(initial_obj_path) 63 | initial_obj_meshes.append( 64 | Mesh( 65 | vertices=initial_obj_mesh.vertices, 66 | faces=initial_obj_mesh.faces, 67 | vc=colors["green"], 68 | smooth=True, 69 | ) 70 | ) 71 | 72 | result_paths = result_path_all.split("&") 73 | # 1 pass, process all data 74 | ball_meshes = [] 75 | human_meshes = [] 76 | object_meshes = [] 77 | 78 | b_meshes = [] 79 | o_meshes = [] 80 | s_meshes = [] 81 | for result_path in result_paths: 82 | if "navi" in result_path: 83 | ball_path = os.path.join(result_path, "ball_objs_bs_idx_0_vis_no_scene") 84 | subject_path = os.path.join(result_path, "objs_bs_idx_0_vis_no_scene") 85 | else: 86 | ball_path = os.path.join( 87 | result_path, 88 | "ball_objs_step_{}_bs_idx_0_vis_no_scene".format(interaction_epoch), 89 | ) 90 | subject_path = os.path.join( 91 | result_path, 92 | "objs_step_{}_bs_idx_0_vis_no_scene".format(interaction_epoch), 93 | ) 94 | sub_ball_meshes = [] 95 | sub_human_meshes = [] 96 | sub_object_meshes = [] 97 | 98 | sub_b_meshes = [] 99 | sub_o_meshes = [] 100 | sub_s_meshes = [] 101 | ori_ball_files = os.listdir(ball_path) 102 | ori_ball_files.sort() 103 | for tmp_name in ori_ball_files: 104 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 105 | if ".ply" in tmp_name: 106 | if "start" in tmp_name: 107 | continue 108 | sub_ball_meshes.append(trimesh.load(os.path.join(ball_path, tmp_name))) 109 | sub_b_meshes.append( 110 | Mesh( 111 | vertices=sub_ball_meshes[-1].vertices, 112 | faces=sub_ball_meshes[-1].faces, 113 | vc=colors["green"], 114 | ) 115 | ) 116 | ori_obj_files = os.listdir(subject_path) 117 | ori_obj_files.sort() 118 | for tmp_name in ori_obj_files: 119 | # if ".obj" in tmp_name or ".ply" in tmp_name and "object" not in tmp_name: 120 | if ".ply" in tmp_name: 121 | if "object" not in tmp_name: 122 | sub_human_meshes.append( 123 | trimesh.load(os.path.join(subject_path, tmp_name)) 124 | ) 125 | sub_human_mesh = Mesh( 126 | vertices=sub_human_meshes[-1].vertices, 127 | faces=sub_human_meshes[-1].faces, 128 | vc=colors["pink"], 129 | smooth=True, 130 | ) 131 | sub_human_mesh.set_vertex_colors( 132 | vc=colors["red"], 133 | vertex_ids=sub_human_meshes[-1].vertices[:, 2] < 0, 134 | ) 135 | sub_s_meshes.append(sub_human_mesh) 136 | else: 137 | sub_object_meshes.append( 138 | trimesh.load(os.path.join(subject_path, tmp_name)) 139 | ) 140 | sub_o_meshes.append( 141 | Mesh( 142 | vertices=sub_object_meshes[-1].vertices, 143 | faces=sub_object_meshes[-1].faces, 144 | vc=colors["yellow"], 145 | ) 146 | ) 147 | assert len(sub_object_meshes) == 0 or len(sub_human_meshes) == len( 148 | sub_object_meshes 149 | ) 150 | 151 | ball_meshes.append(sub_ball_meshes) 152 | human_meshes.append(sub_human_meshes) 153 | object_meshes.append(sub_object_meshes) 154 | 155 | b_meshes.append(sub_b_meshes) 156 | o_meshes.append(sub_o_meshes) 157 | s_meshes.append(sub_s_meshes) 158 | 159 | # 2 pass, render the data 160 | skip_frame = 1 161 | if not cfg.offscreen: 162 | while True: 163 | import time 164 | 165 | for sub_b_meshes, sub_o_meshes, sub_s_meshes in zip( 166 | b_meshes, o_meshes, s_meshes 167 | ): 168 | T = len(sub_s_meshes) 169 | for frame in range(0, T, skip_frame): 170 | start_time = time.time() 171 | 172 | s_mesh = sub_s_meshes[frame] 173 | meshes = [s_mesh] 174 | 175 | if len(sub_o_meshes) > 0: 176 | o_mesh = sub_o_meshes[frame] 177 | meshes = [o_mesh, s_mesh] 178 | 179 | meshes.extend(initial_obj_meshes) 180 | # meshes.extend(sub_b_meshes) 181 | 182 | mv.set_static_meshes(meshes) 183 | while time.time() - start_time < 0.03: 184 | pass 185 | else: 186 | import imageio 187 | 188 | save_dir = os.path.join("{}".format(cfg.save_dir_name)) 189 | if not os.path.exists(save_dir): 190 | os.makedirs(save_dir) 191 | save_dir = os.path.join(save_dir, s_idx + "_" + cfg.use_guidance) 192 | if not os.path.exists(save_dir): 193 | os.makedirs(save_dir) 194 | 195 | img_paths = [] 196 | for result_path, sub_b_meshes, sub_o_meshes, sub_s_meshes in tqdm( 197 | zip(result_paths, b_meshes, o_meshes, s_meshes) 198 | ): 199 | img_path = os.path.join(result_path, "vis") 200 | if not os.path.exists(img_path): 201 | os.makedirs(img_path) 202 | 203 | T = len(sub_s_meshes) 204 | for frame in tqdm(range(0, T, skip_frame), leave=False): 205 | s_mesh = sub_s_meshes[frame] 206 | meshes = [s_mesh] 207 | 208 | if len(sub_o_meshes) > 0: 209 | o_mesh = sub_o_meshes[frame] 210 | meshes = [o_mesh, s_mesh] 211 | 212 | meshes.extend(sub_b_meshes) 213 | meshes.extend(initial_obj_meshes) 214 | 215 | mv.set_static_meshes(meshes) 216 | 217 | camera_pose = np.eye(4) 218 | camera_pose[:3, :3] = euler([80, -15, 0], "xzx") 219 | camera_pose[:2, 3] = np.mean(s_mesh.vertices, axis=0)[:2] + np.array( 220 | [0, -2.0] 221 | ) 222 | camera_pose[2, 3] = 1.3 223 | mv.update_camera_pose(camera_pose) 224 | 225 | mv.save_snapshot(os.path.join(img_path, "%05d.png" % frame)) 226 | img_paths.append(os.path.join(img_path, "%05d.png" % frame)) 227 | 228 | video_name = os.path.join(save_dir, "output_{}.mp4".format(video_idx)) 229 | 230 | im_arr = [] 231 | 232 | for image in tqdm(img_paths, desc="loading images"): 233 | im = imageio.v2.imread(image) 234 | im_arr.append(im) 235 | im_arr = np.asarray(im_arr) 236 | imageio.mimwrite(video_name, im_arr, fps=30, quality=8) 237 | 238 | print("video saved to %s" % video_name) 239 | 240 | 241 | if __name__ == "__main__": 242 | parser = argparse.ArgumentParser(description="GRAB-visualize") 243 | 244 | parser.add_argument( 245 | "--result-path", required=True, type=str, help="The path to the chois_results" 246 | ) 247 | parser.add_argument( 248 | "--initial-obj-path", 249 | type=str, 250 | ) 251 | parser.add_argument( 252 | "--save-dir-name", 253 | type=str, 254 | ) 255 | 256 | parser.add_argument("--s-idx", required=True, type=str) 257 | parser.add_argument("--video-idx", required=True, type=str) 258 | parser.add_argument("--interaction-epoch", required=True, type=int) 259 | parser.add_argument("--use_guidance", required=True, type=str) 260 | 261 | parser.add_argument( 262 | "--model-path", 263 | required=True, 264 | type=str, 265 | help="The path to the folder containing smplx models", 266 | ) 267 | 268 | parser.add_argument("--offscreen", action="store_true") 269 | 270 | args = parser.parse_args() 271 | 272 | result_path = args.result_path 273 | model_path = args.model_path 274 | 275 | # grab_path = 'PATH_TO_DOWNLOADED_GRAB_DATA/grab' 276 | # model_path = 'PATH_TO_DOWNLOADED_MODELS_FROM_SMPLX_WEBSITE/' 277 | 278 | cfg = { 279 | "result_path": result_path, 280 | "initial_obj_path": args.initial_obj_path, 281 | "model_path": model_path, 282 | "offscreen": args.offscreen, 283 | "s_idx": args.s_idx, 284 | "video_idx": args.video_idx, 285 | "interaction_epoch": args.interaction_epoch, 286 | "use_guidance": args.use_guidance, 287 | "save_dir_name": args.save_dir_name, 288 | } 289 | 290 | cfg = Config(**cfg) 291 | visualize_sequences(cfg) 292 | -------------------------------------------------------------------------------- /grasp_generation/utils/object_model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/PKU-EPIC/DexGraspNet 3 | """ 4 | 5 | import os 6 | import pickle 7 | 8 | import numpy as np 9 | import plotly.graph_objects as go 10 | import pytorch3d.ops 11 | import pytorch3d.structures 12 | import torch 13 | import trimesh as tm 14 | from torchsdf import compute_sdf, index_vertices_by_faces 15 | 16 | 17 | class ObjectModel: 18 | def __init__( 19 | self, data_root_path, batch_size_each, num_samples=2000, device="cuda" 20 | ): 21 | """ 22 | Create a Object Model 23 | 24 | Parameters 25 | ---------- 26 | data_root_path: str 27 | directory to object meshes 28 | batch_size_each: int 29 | batch size for each objects 30 | num_samples: int 31 | numbers of object surface points, sampled with fps 32 | device: str | torch.Device 33 | device for torch tensors 34 | """ 35 | 36 | self.device = device 37 | self.batch_size_each = batch_size_each 38 | self.data_root_path = data_root_path 39 | self.num_samples = num_samples 40 | 41 | self.object_code_list = None 42 | self.object_scale_tensor = None 43 | self.object_mesh_list = None 44 | self.object_face_verts_list = None 45 | self.scale_choice = torch.tensor( 46 | [1.0, 1.0, 1.0], dtype=torch.float, device=self.device 47 | ) # NOTE: for omomo, it's always 1 48 | 49 | self.cache_folder = os.path.join(data_root_path, "cache") 50 | if not os.path.exists(self.cache_folder): 51 | os.makedirs(self.cache_folder) 52 | 53 | def initialize( 54 | self, object_code_list, small_object=False, wrist_pos=None, is_train=True 55 | ): 56 | """ 57 | Initialize Object Model with list of objects 58 | 59 | Choose scales, load meshes, sample surface points 60 | 61 | Parameters 62 | ---------- 63 | object_code_list: list | str 64 | list of object codes 65 | """ 66 | point_k = 1 if small_object else 10 67 | if not isinstance(object_code_list, list): 68 | object_code_list = [object_code_list] 69 | self.object_code_list = object_code_list 70 | self.object_scale_tensor = [] 71 | self.object_mesh_list = [] 72 | self.object_face_verts_list = [] 73 | self.surface_points_tensor = [] 74 | for object_code in object_code_list: 75 | self.object_scale_tensor.append( 76 | self.scale_choice[ 77 | torch.randint( 78 | 0, 79 | self.scale_choice.shape[0], 80 | (self.batch_size_each,), 81 | device=self.device, 82 | ) 83 | ] 84 | ) 85 | # Load from cache. 86 | if os.path.exists(os.path.join(self.cache_folder, f"{object_code}.pkl")): 87 | data = pickle.load( 88 | open(os.path.join(self.cache_folder, f"{object_code}.pkl"), "rb") 89 | ) 90 | self.object_mesh_list.append(data["object_mesh_list"]) 91 | self.object_face_verts_list.append(data["object_face_verts_list"]) 92 | if is_train and self.num_samples != 0: 93 | self.surface_points_tensor.append(data["surface_points_tensor"]) 94 | print(f"Loaded {object_code} from cache.") 95 | continue 96 | 97 | self.object_mesh_list.append( 98 | tm.load( 99 | os.path.join(self.data_root_path, f"{object_code}.obj"), 100 | force="mesh", 101 | process=False, 102 | ) 103 | ) 104 | object_verts = torch.Tensor(self.object_mesh_list[-1].vertices).to( 105 | self.device 106 | ) 107 | object_faces = ( 108 | torch.Tensor(self.object_mesh_list[-1].faces).long().to(self.device) 109 | ) 110 | self.object_face_verts_list.append( 111 | index_vertices_by_faces(object_verts, object_faces) 112 | ) 113 | if is_train and self.num_samples != 0: 114 | vertices = torch.tensor( 115 | self.object_mesh_list[-1].vertices, 116 | dtype=torch.float, 117 | device=self.device, 118 | ) 119 | faces = torch.tensor( 120 | self.object_mesh_list[-1].faces, 121 | dtype=torch.float, 122 | device=self.device, 123 | ) 124 | mesh = pytorch3d.structures.Meshes( 125 | vertices.unsqueeze(0), faces.unsqueeze(0) 126 | ) 127 | dense_point_cloud = pytorch3d.ops.sample_points_from_meshes( 128 | mesh, num_samples=100 * self.num_samples 129 | ) 130 | if wrist_pos is not None: 131 | distances = ( 132 | dense_point_cloud * self.scale_choice[0] - wrist_pos 133 | ).norm(dim=2) 134 | dense_point_cloud = dense_point_cloud[distances < 0.2][None] 135 | while dense_point_cloud.shape[1] < self.num_samples: 136 | dense_point_cloud_ = torch.cat( 137 | [ 138 | dense_point_cloud, 139 | pytorch3d.ops.sample_points_from_meshes( 140 | mesh, num_samples=100 * self.num_samples 141 | ), 142 | ], 143 | dim=1, 144 | ) 145 | distances = ( 146 | dense_point_cloud_ * self.scale_choice[0] - wrist_pos 147 | ).norm(dim=2) 148 | dense_point_cloud = torch.cat( 149 | ( 150 | dense_point_cloud, 151 | dense_point_cloud_[distances < 0.2][None], 152 | ), 153 | dim=1, 154 | ) 155 | surface_points = pytorch3d.ops.sample_farthest_points( 156 | dense_point_cloud, K=self.num_samples * point_k 157 | )[0][0] 158 | surface_points.to(dtype=float, device=self.device) 159 | self.surface_points_tensor.append(surface_points) 160 | 161 | # Save to cache. 162 | data = { 163 | "object_mesh_list": self.object_mesh_list[-1], 164 | "object_face_verts_list": self.object_face_verts_list[-1], 165 | "surface_points_tensor": self.surface_points_tensor[-1] 166 | if is_train and self.num_samples != 0 167 | else None, 168 | } 169 | pickle.dump( 170 | data, open(os.path.join(self.cache_folder, f"{object_code}.pkl"), "wb") 171 | ) 172 | 173 | self.object_scale_tensor = torch.stack(self.object_scale_tensor, dim=0) 174 | if is_train and self.num_samples != 0: 175 | self.surface_points_tensor = torch.stack( 176 | self.surface_points_tensor, dim=0 177 | ).repeat_interleave( 178 | self.batch_size_each, dim=0 179 | ) # (n_objects * batch_size_each, num_samples, 3) 180 | 181 | def cal_distance(self, x, with_closest_points=False): 182 | """ 183 | Calculate signed distances from hand contact points to object meshes and return contact normals 184 | 185 | Interiors are positive, exteriors are negative 186 | 187 | Use our modified Kaolin package 188 | 189 | Parameters 190 | ---------- 191 | x: (B, `n_contact`, 3) torch.Tensor 192 | hand contact points 193 | with_closest_points: bool 194 | whether to return closest points on object meshes 195 | 196 | Returns 197 | ------- 198 | distance: (B, `n_contact`) torch.Tensor 199 | signed distances from hand contact points to object meshes, inside is positive 200 | normals: (B, `n_contact`, 3) torch.Tensor 201 | contact normal vectors defined by gradient 202 | closest_points: (B, `n_contact`, 3) torch.Tensor 203 | contact points on object meshes, returned only when `with_closest_points is True` 204 | """ 205 | _, n_points, _ = x.shape 206 | x = x.reshape(-1, self.batch_size_each * n_points, 3) 207 | distance = [] 208 | normals = [] 209 | closest_points = [] 210 | scale = self.object_scale_tensor.repeat_interleave(n_points, dim=1) 211 | x = x / scale.unsqueeze(2) 212 | for i in range(len(self.object_mesh_list)): 213 | face_verts = self.object_face_verts_list[i] 214 | dis, dis_signs, normal, _ = compute_sdf(x[i], face_verts) 215 | if with_closest_points: 216 | closest_points.append(x[i] - dis.sqrt().unsqueeze(1) * normal) 217 | dis = torch.sqrt(dis + 1e-8) 218 | dis = dis * (-dis_signs) 219 | distance.append(dis) 220 | normals.append(normal * dis_signs.unsqueeze(1)) 221 | distance = torch.stack(distance) 222 | normals = torch.stack(normals) 223 | distance = distance * scale 224 | distance = distance.reshape(-1, n_points) 225 | normals = normals.reshape(-1, n_points, 3) 226 | if with_closest_points: 227 | closest_points = (torch.stack(closest_points) * scale.unsqueeze(2)).reshape( 228 | -1, n_points, 3 229 | ) 230 | return distance, normals, closest_points 231 | return distance, normals 232 | 233 | def get_plotly_data( 234 | self, i, color="lightgreen", opacity=0.5, pose=None, with_surface_points=False 235 | ): 236 | """ 237 | Get visualization data for plotly.graph_objects 238 | 239 | Parameters 240 | ---------- 241 | i: int 242 | index of data 243 | color: str 244 | color of mesh 245 | opacity: float 246 | opacity 247 | pose: (4, 4) matrix 248 | homogeneous transformation matrix 249 | 250 | Returns 251 | ------- 252 | data: list 253 | list of plotly.graph_object visualization data 254 | """ 255 | model_index = i // self.batch_size_each 256 | model_scale = ( 257 | self.object_scale_tensor[model_index, i % self.batch_size_each] 258 | .detach() 259 | .cpu() 260 | .numpy() 261 | ) 262 | mesh = self.object_mesh_list[model_index] 263 | vertices = mesh.vertices * model_scale 264 | if pose is not None: 265 | pose = np.array(pose, dtype=np.float32) 266 | vertices = vertices @ pose[:3, :3].T + pose[:3, 3] 267 | data = [ 268 | go.Mesh3d( 269 | x=vertices[:, 0], 270 | y=vertices[:, 1], 271 | z=vertices[:, 2], 272 | i=mesh.faces[:, 0], 273 | j=mesh.faces[:, 1], 274 | k=mesh.faces[:, 2], 275 | color=color, 276 | opacity=opacity, 277 | ) 278 | ] 279 | if with_surface_points: 280 | surface_points = self.surface_points_tensor[i].detach().cpu() * model_scale 281 | data.append( 282 | go.Scatter3d( 283 | x=surface_points[:, 0], 284 | y=surface_points[:, 1], 285 | z=surface_points[:, 2], 286 | mode="markers", 287 | marker=dict(color="red", size=3), 288 | ) 289 | ) 290 | return data 291 | 292 | def get_surface_points_near_wrist(self, idx, wrist_pos): 293 | object_scale = self.object_scale_tensor.flatten().unsqueeze(1).unsqueeze(2)[idx] 294 | surface_points = self.surface_points_tensor[idx] * object_scale 295 | distances = (surface_points - wrist_pos).norm(dim=-1) 296 | surface_points = surface_points[distances < 0.2] 297 | 298 | return wrist_pos - surface_points.mean(0) 299 | -------------------------------------------------------------------------------- /argument_parser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | def parse_opt(): 5 | parser = argparse.ArgumentParser() 6 | 7 | parser.add_argument("--vis_wdir", type=str, default="default_folder") 8 | 9 | parser.add_argument("--navi_save_dir", type=str, default="navi_release") 10 | parser.add_argument("--rnet_save_dir", type=str, default="rnet_release") 11 | parser.add_argument("--cnet_save_dir", type=str, default="cnet_release") 12 | 13 | parser.add_argument("--project", default="runs/train", help="project/name") 14 | parser.add_argument("--wandb_pj_name", type=str, default="", help="project name") 15 | parser.add_argument("--entity", default="zhenkirito123", help="W&B entity") 16 | parser.add_argument("--exp_name", default="", help="save to project/name") 17 | parser.add_argument("--device", default="0", help="cuda device") 18 | 19 | parser.add_argument("--window", type=int, default=120, help="horizon") 20 | 21 | parser.add_argument("--batch_size", type=int, default=32, help="batch size") 22 | parser.add_argument( 23 | "--learning_rate", type=float, default=2e-4, help="generator_learning_rate" 24 | ) 25 | 26 | parser.add_argument("--checkpoint", type=str, default="", help="checkpoint") 27 | 28 | parser.add_argument("--data_root_folder", type=str, default="", help="checkpoint") 29 | 30 | parser.add_argument( 31 | "--n_dec_layers_nav", type=int, default=4, help="the number of decoder layers" 32 | ) 33 | parser.add_argument( 34 | "--n_head_nav", 35 | type=int, 36 | default=4, 37 | help="the number of heads in self-attention", 38 | ) 39 | parser.add_argument( 40 | "--d_k_nav", type=int, default=256, help="the dimension of keys in transformer" 41 | ) 42 | parser.add_argument( 43 | "--d_v_nav", 44 | type=int, 45 | default=256, 46 | help="the dimension of values in transformer", 47 | ) 48 | parser.add_argument( 49 | "--d_model_nav", 50 | type=int, 51 | default=512, 52 | help="the dimension of intermediate representation in transformer", 53 | ) 54 | 55 | # For testing sampled results 56 | parser.add_argument("--test_sample_res", action="store_true") 57 | 58 | # For testing sampled results on training dataset 59 | parser.add_argument("--test_on_train", action="store_true") 60 | 61 | # For loss type 62 | parser.add_argument("--use_l2_loss", action="store_true") 63 | 64 | parser.add_argument("--compute_metrics", action="store_true") 65 | 66 | # For interaction model setting 67 | # For adding full body prediction. 68 | parser.add_argument( 69 | "--n_dec_layers", type=int, default=4, help="the number of decoder layers" 70 | ) 71 | parser.add_argument( 72 | "--n_head", type=int, default=4, help="the number of heads in self-attention" 73 | ) 74 | parser.add_argument( 75 | "--d_k", type=int, default=256, help="the dimension of keys in transformer" 76 | ) 77 | parser.add_argument( 78 | "--d_v", type=int, default=256, help="the dimension of values in transformer" 79 | ) 80 | parser.add_argument( 81 | "--d_model", 82 | type=int, 83 | default=512, 84 | help="the dimension of intermediate representation in transformer", 85 | ) 86 | 87 | # Add language conditions. 88 | parser.add_argument("--add_contact_label", action="store_true") 89 | 90 | # Train and test on different objects. 91 | parser.add_argument("--use_object_split", action="store_true") 92 | 93 | # For adding start and end object position (xyz) and rotation (6D rotation). 94 | parser.add_argument("--add_start_end_object_pos_rot", action="store_true") 95 | 96 | # For adding start and end object position (xyz). 97 | parser.add_argument("--add_start_end_object_pos", action="store_true") 98 | 99 | # For adding start and end object position at z plane (xy). 100 | parser.add_argument("--add_start_end_object_pos_xy", action="store_true") 101 | 102 | # Random sample waypoints instead of fixed intervals. 103 | parser.add_argument("--use_random_waypoints", action="store_true") 104 | 105 | # Input the first human pose, maybe can connect the windows better. 106 | parser.add_argument("--remove_target_z", action="store_true") 107 | 108 | # Input the first human pose, maybe can connect the windows better. 109 | parser.add_argument("--use_guidance_in_denoising", action="store_true") 110 | 111 | parser.add_argument("--use_optimization_in_denoising", action="store_true") 112 | 113 | # Add rest offsets for body shape information. 114 | parser.add_argument("--add_rest_human_skeleton", action="store_true") 115 | 116 | # Add rest offsets for body shape information. 117 | parser.add_argument("--use_first_frame_bps", action="store_true") 118 | 119 | parser.add_argument("--input_full_human_pose", action="store_true") 120 | 121 | parser.add_argument("--use_two_stage_pipeline", action="store_true") 122 | 123 | parser.add_argument("--use_unified_interaction_model", action="store_true") 124 | 125 | # Visualize the results from different noise levels. 126 | parser.add_argument("--return_diff_level_res", action="store_true") 127 | 128 | parser.add_argument( 129 | "--test_object_names", 130 | type=str, 131 | nargs="+", 132 | default=["largebox"], 133 | help="object names for long sequence generation testing", 134 | ) 135 | 136 | parser.add_argument("--use_long_planned_path", action="store_true") 137 | 138 | parser.add_argument("--use_long_planned_path_only_navi", action="store_true") 139 | 140 | # For testing sampled results w planned path 141 | parser.add_argument("--use_planned_path", action="store_true") 142 | 143 | parser.add_argument( 144 | "--loss_w_feet", 145 | type=float, 146 | default=1, 147 | help="the loss weight for feet contact loss", 148 | ) 149 | 150 | parser.add_argument( 151 | "--loss_w_fk", type=float, default=1, help="the loss weight for fk loss" 152 | ) 153 | 154 | parser.add_argument( 155 | "--loss_w_obj_pts", type=float, default=1, help="the loss weight for fk loss" 156 | ) 157 | 158 | # For interaction model setting. 159 | parser.add_argument("--add_interaction_root_xy_ori", action="store_true",) 160 | parser.add_argument("--add_interaction_feet_contact", action="store_true",) 161 | parser.add_argument("--add_wrist_relative", action="store_true") 162 | parser.add_argument("--add_object_static", action="store_true") 163 | 164 | parser.add_argument("--use_noisy_traj", action="store_true") 165 | 166 | ############################ finger motion related settings ############################ 167 | parser.add_argument("--add_finger_motion", action="store_true") 168 | 169 | parser.add_argument( 170 | "--finger_project", 171 | default="./experiments", 172 | help="output folder for weights and visualizations", 173 | ) 174 | parser.add_argument( 175 | "--finger_wandb_pj_name", 176 | type=str, 177 | default="omomo_only_finger", 178 | help="wandb_proj_name", 179 | ) 180 | parser.add_argument("--finger_entity", default="zhenkirito123", help="W&B entity") 181 | parser.add_argument( 182 | "--finger_exp_name", 183 | default="fnet_release", 184 | help="save to project/name", 185 | ) 186 | parser.add_argument("--finger_device", default="0", help="cuda device") 187 | 188 | parser.add_argument("--finger_window", type=int, default=30, help="horizon") 189 | 190 | parser.add_argument("--finger_batch_size", type=int, default=32, help="batch size") 191 | parser.add_argument( 192 | "--finger_learning_rate", 193 | type=float, 194 | default=2e-4, 195 | help="generator_learning_rate", 196 | ) 197 | 198 | parser.add_argument("--finger_checkpoint", type=str, default="", help="checkpoint") 199 | 200 | parser.add_argument( 201 | "--finger_n_dec_layers", 202 | type=int, 203 | default=4, 204 | help="the number of decoder layers", 205 | ) 206 | parser.add_argument( 207 | "--finger_n_head", 208 | type=int, 209 | default=4, 210 | help="the number of heads in self-attention", 211 | ) 212 | parser.add_argument( 213 | "--finger_d_k", 214 | type=int, 215 | default=256, 216 | help="the dimension of keys in transformer", 217 | ) 218 | parser.add_argument( 219 | "--finger_d_v", 220 | type=int, 221 | default=256, 222 | help="the dimension of values in transformer", 223 | ) 224 | parser.add_argument( 225 | "--finger_d_model", 226 | type=int, 227 | default=512, 228 | help="the dimension of intermediate representation in transformer", 229 | ) 230 | 231 | parser.add_argument("--finger_generate_refine_dataset", action="store_true") 232 | parser.add_argument("--finger_generate_quant_eval", action="store_true") 233 | parser.add_argument("--finger_generate_vis_eval", action="store_true") 234 | parser.add_argument("--finger_milestone", type=str, default="17") 235 | # For testing sampled results 236 | parser.add_argument("--finger_test_sample_res", action="store_true") 237 | 238 | # For testing sampled results on training dataset 239 | parser.add_argument("--finger_test_sample_res_on_train", action="store_true") 240 | 241 | # For loss type 242 | parser.add_argument("--finger_use_l2_loss", action="store_true") 243 | 244 | # For training diffusion model without condition 245 | parser.add_argument("--finger_remove_condition", action="store_true") 246 | 247 | # For normalizing condition 248 | parser.add_argument("--finger_normalize_condition", action="store_true") 249 | 250 | # For adding first human pose as condition (shared by current trainer and FullBody trainer) 251 | parser.add_argument("--finger_add_start_human_pose", action="store_true") 252 | 253 | # FullBody trainer config: For adding hand pose (translation+rotation) as condition 254 | parser.add_argument("--finger_add_hand_pose", action="store_true") 255 | 256 | # FullBody trainer config: For adding hand pose (translation only) as condition 257 | parser.add_argument("--finger_add_hand_trans_only", action="store_true") 258 | 259 | # FullBody trainer config: For adding hand and foot trans as condition 260 | parser.add_argument("--finger_add_hand_foot_trans", action="store_true") 261 | 262 | # For canonicalizing the first pose's facing direction 263 | parser.add_argument("--finger_cano_init_pose", action="store_true") 264 | 265 | # For predicting hand position only 266 | parser.add_argument("--finger_pred_hand_jpos_only_from_obj", action="store_true") 267 | 268 | # For predicting hand position only 269 | parser.add_argument("--finger_pred_palm_jpos_from_obj", action="store_true") 270 | 271 | # For predicting hand position only 272 | parser.add_argument("--finger_pred_hand_jpos_and_rot_from_obj", action="store_true") 273 | 274 | # For running the whole pipeline. 275 | parser.add_argument("--finger_run_whole_pipeline", action="store_true") 276 | 277 | parser.add_argument("--finger_add_object_bps", action="store_true") 278 | 279 | parser.add_argument("--finger_add_hand_processing", action="store_true") 280 | 281 | parser.add_argument("--finger_for_quant_eval", action="store_true") 282 | 283 | parser.add_argument("--finger_use_gt_hand_for_eval", action="store_true") 284 | 285 | parser.add_argument("--finger_use_object_split", action="store_true") 286 | 287 | parser.add_argument("--finger_add_palm_jpos_only", action="store_true") 288 | 289 | parser.add_argument("--finger_use_blender_data", action="store_true") 290 | 291 | parser.add_argument("--finger_use_wandb", action="store_true") 292 | 293 | parser.add_argument("--finger_use_arctic", action="store_true") 294 | parser.add_argument("--finger_omomo_obj", type=str, default="") 295 | 296 | parser.add_argument("--finger_train_ambient_sensor", action="store_true") 297 | parser.add_argument("--finger_train_proximity_sensor", action="store_true") 298 | # parser.add_argument("--finger_train_contact_label", action="store_true") 299 | 300 | parser.add_argument("--finger_wrist_obj_traj_condition", action="store_true") 301 | parser.add_argument("--finger_ambient_sensor_condition", action="store_true") 302 | parser.add_argument("--finger_contact_label_condition", action="store_true") 303 | 304 | # Set a few args to True by default. 305 | parser.set_defaults( 306 | finger_use_joints24=True, 307 | finger_use_omomo=True, 308 | finger_train_both_sensor=True, 309 | finger_proximity_sensor_condition=True, 310 | finger_ref_pose_condition=True, 311 | ) 312 | 313 | opt = parser.parse_args() 314 | return opt 315 | -------------------------------------------------------------------------------- /physics_tracking/models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from typing import Optional, Union 4 | 5 | 6 | class RunningMeanStd(torch.nn.Module): 7 | def __init__(self, dim: int, clamp: float=0): 8 | super().__init__() 9 | self.epsilon = 1e-5 10 | self.clamp = clamp 11 | self.register_buffer("mean", torch.zeros(dim, dtype=torch.float64)) 12 | self.register_buffer("var", torch.ones(dim, dtype=torch.float64)) 13 | self.register_buffer("count", torch.ones((), dtype=torch.float64)) 14 | 15 | def forward(self, x, unnorm=False): 16 | mean = self.mean.to(torch.float32) 17 | var = self.var.to(torch.float32)+self.epsilon 18 | if unnorm: 19 | if self.clamp: 20 | x = torch.clamp(x, min=-self.clamp, max=self.clamp) 21 | return mean + torch.sqrt(var) * x 22 | x = (x - mean) * torch.rsqrt(var) 23 | if self.clamp: 24 | return torch.clamp(x, min=-self.clamp, max=self.clamp) 25 | return x 26 | 27 | @torch.no_grad() 28 | def update(self, x): 29 | x = x.view(-1, x.size(-1)) 30 | var, mean = torch.var_mean(x, dim=0, unbiased=True) 31 | count = x.size(0) 32 | count_ = count + self.count 33 | delta = mean - self.mean 34 | m = self.var * self.count + var * count + delta**2 * self.count * count / count_ 35 | self.mean.copy_(self.mean+delta*count/count_) 36 | self.var.copy_(m / count_) 37 | self.count.copy_(count_) 38 | 39 | def reset_counter(self): 40 | self.count.fill_(1) 41 | 42 | class DiagonalPopArt(torch.nn.Module): 43 | def __init__(self, dim: int, weight: torch.Tensor, bias: torch.Tensor, momentum:float=0.1): 44 | super().__init__() 45 | self.epsilon = 1e-5 46 | 47 | self.momentum = momentum 48 | self.register_buffer("m", torch.zeros((dim,), dtype=torch.float64)) 49 | self.register_buffer("v", torch.full((dim,), self.epsilon, dtype=torch.float64)) 50 | self.register_buffer("debias", torch.zeros(1, dtype=torch.float64)) 51 | 52 | self.weight = weight 53 | self.bias = bias 54 | 55 | def forward(self, x, unnorm=False): 56 | debias = self.debias.clip(min=self.epsilon) 57 | mean = self.m/debias 58 | var = (self.v - self.m.square()).div_(debias) 59 | if unnorm: 60 | std = torch.sqrt(var) 61 | return (mean + std * x).to(x.dtype) 62 | x = ((x - mean) * torch.rsqrt(var)).to(x.dtype) 63 | return x 64 | 65 | @torch.no_grad() 66 | def update(self, x): 67 | x = x.view(-1, x.size(-1)) 68 | running_m = torch.mean(x, dim=0) 69 | running_v = torch.mean(x.square(), dim=0) 70 | new_m = self.m.mul(1-self.momentum).add_(running_m, alpha=self.momentum) 71 | new_v = self.v.mul(1-self.momentum).add_(running_v, alpha=self.momentum) 72 | std = (self.v - self.m.square()).sqrt_() 73 | new_std_inv = (new_v - new_m.square()).rsqrt_() 74 | 75 | scale = std.mul_(new_std_inv) 76 | shift = (self.m - new_m).mul_(new_std_inv) 77 | 78 | self.bias.data.mul_(scale).add_(shift) 79 | self.weight.data.mul_(scale.unsqueeze_(-1)) 80 | 81 | self.debias.data.mul_(1-self.momentum).add_(1.0*self.momentum) 82 | self.m.data.copy_(new_m) 83 | self.v.data.copy_(new_v) 84 | 85 | 86 | class Discriminator(torch.nn.Module): 87 | def __init__(self, disc_dim, latent_dim=256): 88 | super().__init__() 89 | self.rnn = torch.nn.GRU(disc_dim, latent_dim, batch_first=True) 90 | self.mlp = torch.nn.Sequential( 91 | torch.nn.Linear(latent_dim, 256), 92 | torch.nn.ReLU(), 93 | torch.nn.Linear(256, 128), 94 | torch.nn.ReLU(), 95 | torch.nn.Linear(128, 32) 96 | ) 97 | if self.rnn is not None: 98 | i = 0 99 | for n, p in self.mlp.named_parameters(): 100 | if "bias" in n: 101 | torch.nn.init.constant_(p, 0.) 102 | elif "weight" in n: 103 | gain = 1 if i == 2 else 2**0.5 104 | torch.nn.init.orthogonal_(p, gain=gain) 105 | i += 1 106 | self.ob_normalizer = RunningMeanStd(disc_dim) 107 | self.all_inst = torch.arange(0) 108 | 109 | def forward(self, s, seq_end_frame, normalize=True): 110 | if normalize: s = self.ob_normalizer(s) 111 | if self.rnn is None: 112 | s = s.view(s.size(0), -1) 113 | else: 114 | n_inst = s.size(0) 115 | if n_inst > self.all_inst.size(0): 116 | self.all_inst = torch.arange(n_inst, 117 | dtype=seq_end_frame.dtype, device=seq_end_frame.device) 118 | s, _ = self.rnn(s) 119 | s = s[(self.all_inst[:n_inst], torch.clip(seq_end_frame, max=s.size(1)-1))] 120 | return self.mlp(s) 121 | 122 | 123 | class ACModel(torch.nn.Module): 124 | 125 | class Critic(torch.nn.Module): 126 | def __init__(self, state_dim, goal_dim, value_dim=1, latent_dim=256): 127 | super().__init__() 128 | self.rnn = None 129 | # self.rnn = torch.nn.GRU(state_dim, latent_dim, batch_first=True) 130 | # self.mlp = torch.nn.Sequential( 131 | # torch.nn.Linear(latent_dim+goal_dim, 1024), 132 | # torch.nn.ReLU6(), 133 | # torch.nn.Linear(1024, 512), 134 | # torch.nn.ReLU6(), 135 | # torch.nn.Linear(512, value_dim) 136 | # ) 137 | self.mlp = torch.nn.Sequential( 138 | torch.nn.Linear(state_dim if self.rnn is None else latent_dim, 2048), 139 | torch.nn.GELU(), 140 | torch.nn.Linear(2048, 2048), 141 | torch.nn.GELU(), 142 | torch.nn.Linear(2048, 1024), 143 | torch.nn.GELU(), 144 | torch.nn.Linear(1024, 1024), 145 | torch.nn.GELU(), 146 | torch.nn.Linear(1024, 512), 147 | torch.nn.GELU(), 148 | torch.nn.Linear(512, value_dim) 149 | ) 150 | # self.embed_goal = torch.nn.Sequential( 151 | # torch.nn.Linear(goal_dim, latent_dim), 152 | # torch.nn.ReLU6(), 153 | # torch.nn.Linear(latent_dim, latent_dim) 154 | # ) 155 | i = 0 156 | for n, p in self.mlp.named_parameters(): 157 | if "bias" in n: 158 | torch.nn.init.constant_(p, 0.) 159 | elif "weight" in n: 160 | torch.nn.init.uniform_(p, -0.0001, 0.0001) 161 | i += 1 162 | self.all_inst = torch.arange(0) 163 | 164 | def forward(self, s, seq_end_frame, g=None): 165 | if self.rnn is None: 166 | s = s.view(s.size(0), -1) 167 | else: 168 | n_inst = s.size(0) 169 | if n_inst > self.all_inst.size(0): 170 | self.all_inst = torch.arange(n_inst, 171 | dtype=seq_end_frame.dtype, device=seq_end_frame.device) 172 | s, _ = self.rnn(s) 173 | s = s[(self.all_inst[:n_inst], torch.clip(seq_end_frame, max=s.size(1)-1))] 174 | if g is not None: 175 | s = s + self.embed_goal(g) 176 | # s = torch.cat((s, g), -1) 177 | return self.mlp(s) 178 | 179 | 180 | class Actor(torch.nn.Module): 181 | def __init__(self, state_dim, act_dim, goal_dim, latent_dim=256, init_mu=None, init_sigma=None): 182 | super().__init__() 183 | self.rnn = None 184 | # self.rnn = torch.nn.GRU(state_dim, latent_dim, batch_first=True) 185 | # self.mlp = torch.nn.Sequential( 186 | # torch.nn.Linear(latent_dim+goal_dim, 1024), 187 | # torch.nn.ReLU6(), 188 | # torch.nn.Linear(1024, 512), 189 | # torch.nn.ReLU6() 190 | # ) 191 | self.mlp = torch.nn.Sequential( 192 | torch.nn.Linear(state_dim if self.rnn is None else latent_dim, 2048), 193 | torch.nn.GELU(), 194 | torch.nn.Linear(2048, 2048), 195 | torch.nn.GELU(), 196 | torch.nn.Linear(2048, 1024), 197 | torch.nn.GELU(), 198 | torch.nn.Linear(1024, 1024), 199 | torch.nn.GELU(), 200 | torch.nn.Linear(1024, 512), 201 | torch.nn.GELU(), 202 | ) 203 | # self.embed_goal = torch.nn.Sequential( 204 | # torch.nn.Linear(goal_dim, latent_dim), 205 | # torch.nn.ReLU6(), 206 | # torch.nn.Linear(latent_dim, latent_dim) 207 | # ) 208 | self.mu = torch.nn.Linear(512, act_dim) 209 | self.log_sigma = torch.nn.Linear(512, act_dim) 210 | with torch.no_grad(): 211 | if init_mu is not None: 212 | if torch.is_tensor(init_mu): 213 | mu = torch.ones_like(self.mu.bias)*init_mu 214 | else: 215 | mu = np.ones(self.mu.bias.shape, dtype=np.float32)*init_mu 216 | mu = torch.from_numpy(mu) 217 | self.mu.bias.data.copy_(mu) 218 | torch.nn.init.uniform_(self.mu.weight, -0.00001, 0.00001) 219 | if init_sigma is None: 220 | torch.nn.init.constant_(self.log_sigma.bias, -3) 221 | torch.nn.init.uniform_(self.log_sigma.weight, -0.0001, 0.0001) 222 | else: 223 | if torch.is_tensor(init_sigma): 224 | log_sigma = (torch.ones_like(self.log_sigma.bias)*init_sigma).log_() 225 | else: 226 | log_sigma = np.log(np.ones(self.log_sigma.bias.shape, dtype=np.float32)*init_sigma) 227 | log_sigma = torch.from_numpy(log_sigma) 228 | self.log_sigma.bias.data.copy_(log_sigma) 229 | torch.nn.init.uniform_(self.log_sigma.weight, -0.00001, 0.00001) 230 | self.all_inst = torch.arange(0) 231 | 232 | def forward(self, s, seq_end_frame, g=None): 233 | s0 = s 234 | if self.rnn is None: 235 | s = s.view(s.size(0), -1) 236 | else: 237 | n_inst = s.size(0) 238 | if n_inst > self.all_inst.size(0): 239 | self.all_inst = torch.arange(n_inst, 240 | dtype=seq_end_frame.dtype, device=seq_end_frame.device) 241 | s, _ = self.rnn(s) 242 | s = s[(self.all_inst[:n_inst], torch.clip(seq_end_frame, max=s.size(1)-1))] 243 | if g is not None: 244 | s = s + self.embed_goal(g) 245 | # s = torch.cat((s, g), -1) 246 | latent = self.mlp(s) 247 | mu = self.mu(latent) 248 | log_sigma = self.log_sigma(latent) 249 | sigma = torch.exp(log_sigma) + 1e-8 250 | try: 251 | return torch.distributions.Normal(mu, sigma) 252 | except: 253 | m = torch.where(torch.isnan(sigma).logical_or_(torch.isinf(sigma))) 254 | print(m) 255 | print(torch.max(s0[m[0][0]]), torch.min(s0[m[0][0]])) 256 | print(torch.max(s[m[0][0]]), torch.min(s[m[0][0]])) 257 | print(torch.max(latent[m[0][0]]), torch.min(latent[m[0][0]])) 258 | raise 259 | 260 | 261 | def __init__(self, state_dim: int, act_dim: int, goal_dim: int=0, value_dim: int=1, 262 | normalize_value: bool=False, 263 | init_mu:Optional[Union[torch.Tensor, float]]=None, init_sigma:Optional[Union[torch.Tensor, float]]=None 264 | ): 265 | super().__init__() 266 | self.state_dim = state_dim 267 | self.goal_dim = goal_dim 268 | self.actor = self.Actor(state_dim, act_dim, self.goal_dim, init_mu=init_mu, init_sigma=init_sigma) 269 | self.critic = self.Critic(state_dim, goal_dim, value_dim) 270 | self.ob_normalizer = RunningMeanStd(state_dim, clamp=5.0) 271 | if normalize_value: 272 | self.value_normalizer = DiagonalPopArt(value_dim, 273 | self.critic.mlp[-1].weight, self.critic.mlp[-1].bias) 274 | else: 275 | self.value_normalizer = None 276 | 277 | def observe(self, obs, norm=True): 278 | if self.goal_dim > 0: 279 | s = obs[:, :-self.goal_dim] 280 | g = obs[:, -self.goal_dim:] 281 | else: 282 | s = obs 283 | g = None 284 | s = s.view(*s.shape[:-1], -1, self.state_dim) 285 | return self.ob_normalizer(s) if norm else s, g 286 | 287 | def eval_(self, s, seq_end_frame, g, unnorm): 288 | v = self.critic(s, seq_end_frame, g) 289 | if unnorm and self.value_normalizer is not None: 290 | v = self.value_normalizer(v, unnorm=True) 291 | return v 292 | 293 | def act(self, obs, seq_end_frame, stochastic=None, unnorm=False): 294 | if stochastic is None: 295 | stochastic = self.training 296 | s, g = self.observe(obs) 297 | pi = self.actor(s, seq_end_frame, g) 298 | if stochastic: 299 | a = pi.sample() 300 | lp = pi.log_prob(a) 301 | if g is not None: 302 | g = g[...,:self.goal_dim] 303 | return a, self.eval_(s, seq_end_frame, g, unnorm), lp 304 | else: 305 | return pi.mean 306 | 307 | def evaluate(self, obs, seq_end_frame, unnorm=False): 308 | s, g = self.observe(obs) 309 | if g is not None: 310 | g = g[...,:self.goal_dim] 311 | return self.eval_(s, seq_end_frame, g, unnorm) 312 | 313 | def forward(self, obs, seq_end_frame, unnorm=False): 314 | s, g = self.observe(obs) 315 | pi = self.actor(s, seq_end_frame, g) 316 | if g is not None: 317 | g = g[...,:self.goal_dim] 318 | return pi, self.eval_(s, seq_end_frame, g, unnorm) 319 | -------------------------------------------------------------------------------- /grasp_generation/utils/initializations.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/PKU-EPIC/DexGraspNet 3 | """ 4 | 5 | import math 6 | 7 | import numpy as np 8 | import pytorch3d.ops 9 | import pytorch3d.structures 10 | import pytorch3d.transforms 11 | import torch 12 | import torch.nn.functional 13 | import transforms3d 14 | import trimesh as tm 15 | 16 | 17 | def initialize_convex_hull( 18 | hand_model, object_model, args, wrist_pos, wrist_rot, move_away=False, no_fc=False 19 | ): 20 | """ 21 | Initialize grasp translation, rotation, thetas, and contact point indices 22 | 23 | Parameters 24 | ---------- 25 | hand_model: hand_model.HandModel 26 | object_model: object_model.ObjectModel 27 | args: Namespace 28 | """ 29 | 30 | device = hand_model.device 31 | n_objects = len(object_model.object_mesh_list) 32 | batch_size_each = object_model.batch_size_each 33 | total_batch_size = n_objects * batch_size_each 34 | 35 | # initialize translation and rotation 36 | 37 | translation = torch.zeros([total_batch_size, 3], dtype=torch.float, device=device) 38 | rotation = torch.zeros([total_batch_size, 3, 3], dtype=torch.float, device=device) 39 | 40 | if len(wrist_pos.shape) == 1: # 3, 3 X 3 41 | hand_model.initial_translation = wrist_pos[None].repeat(total_batch_size, 1) 42 | translation[0] = wrist_pos 43 | rotation[0] = wrist_rot 44 | for i in range(1, total_batch_size): 45 | translation[i] = wrist_pos 46 | axis = torch.randn(3) 47 | axis /= torch.norm(axis) 48 | aa = axis * torch.rand(1) * np.pi / 6 49 | rotation[i] = wrist_rot.matmul( 50 | pytorch3d.transforms.axis_angle_to_matrix(aa).to(device) 51 | ) 52 | if move_away: 53 | if hand_model.left_hand: 54 | offset = rotation.matmul( 55 | torch.tensor([0.1, 0, 0], dtype=torch.float, device=device).reshape( 56 | 1, 3, 1 57 | ) 58 | ).squeeze(-1) 59 | translation = 1.2 * (wrist_pos + offset) - offset 60 | else: 61 | offset = rotation.matmul( 62 | torch.tensor( 63 | [-0.1, 0, 0], dtype=torch.float, device=device 64 | ).reshape(1, 3, 1) 65 | ).squeeze(-1) 66 | translation = 1.2 * (wrist_pos + offset) - offset 67 | elif len(wrist_pos.shape) == 2: # T X 3, T X 3 X 3 68 | hand_model.initial_translation = wrist_pos 69 | translation = wrist_pos 70 | rotation = wrist_rot 71 | else: 72 | raise ValueError("Invalid wrist_pos shape") 73 | 74 | # initialize thetas 75 | # thetas_mu: hand-crafted canonicalized hand articulation 76 | # use normal distribution to jitter the thetas 77 | 78 | thetas_mu = ( 79 | torch.tensor( 80 | [ 81 | 0, 82 | 0, 83 | torch.pi / 6, 84 | 0, 85 | 0, 86 | 0, 87 | 0, 88 | 0, 89 | 0, 90 | 0, 91 | 0, 92 | torch.pi / 6, 93 | 0, 94 | 0, 95 | 0, 96 | 0, 97 | 0, 98 | 0, 99 | 0, 100 | 0, 101 | torch.pi / 6, 102 | 0, 103 | 0, 104 | 0, 105 | 0, 106 | 0, 107 | 0, 108 | 0, 109 | 0, 110 | torch.pi / 6, 111 | 0, 112 | 0, 113 | 0, 114 | 0, 115 | 0, 116 | 0, 117 | *(torch.pi / 2 * torch.tensor([1, 0.5, 0], dtype=torch.float)), 118 | 0, 119 | 0, 120 | 0, 121 | 0, 122 | 0, 123 | 0, 124 | ], 125 | dtype=torch.float, 126 | device=device, 127 | ) 128 | .unsqueeze(0) 129 | .repeat(total_batch_size, 1) 130 | ) 131 | if hand_model.left_hand: 132 | thetas_mu = ( 133 | torch.tensor( 134 | [ 135 | 0, 136 | 0, 137 | -torch.pi / 6, 138 | 0, 139 | 0, 140 | 0, 141 | 0, 142 | 0, 143 | 0, 144 | 0, 145 | 0, 146 | -torch.pi / 6, 147 | 0, 148 | 0, 149 | 0, 150 | 0, 151 | 0, 152 | 0, 153 | 0, 154 | 0, 155 | -torch.pi / 6, 156 | 0, 157 | 0, 158 | 0, 159 | 0, 160 | 0, 161 | 0, 162 | 0, 163 | 0, 164 | -torch.pi / 6, 165 | 0, 166 | 0, 167 | 0, 168 | 0, 169 | 0, 170 | 0, 171 | *(torch.pi / 2 * torch.tensor([1, -0.5, 0], dtype=torch.float)), 172 | 0, 173 | 0, 174 | 0, 175 | 0, 176 | 0, 177 | 0, 178 | ], 179 | dtype=torch.float, 180 | device=device, 181 | ) 182 | .unsqueeze(0) 183 | .repeat(total_batch_size, 1) 184 | ) 185 | thetas_sigma = args.jitter_strength * torch.ones( 186 | [total_batch_size, 45], dtype=torch.float, device=device 187 | ) 188 | thetas = torch.normal(thetas_mu, thetas_sigma) 189 | if no_fc: 190 | thetas *= 1e-2 191 | 192 | rotation = pytorch3d.transforms.quaternion_to_axis_angle( 193 | pytorch3d.transforms.matrix_to_quaternion(rotation) 194 | ) 195 | hand_pose = torch.cat( 196 | [ 197 | translation, 198 | rotation, 199 | thetas, 200 | ], 201 | dim=1, 202 | ) 203 | hand_pose.requires_grad_() 204 | 205 | contact_point_indices = torch.randint( 206 | hand_model.n_contact_candidates, 207 | size=[total_batch_size, args.n_contact], 208 | device=device, 209 | ) 210 | hand_model.set_parameters(hand_pose, contact_point_indices) 211 | 212 | 213 | def initialize_convex_hull_original(hand_model, object_model, args, no_fc=False): 214 | """ 215 | Initialize grasp translation, rotation, thetas, and contact point indices 216 | 217 | Parameters 218 | ---------- 219 | hand_model: hand_model.HandModel 220 | object_model: object_model.ObjectModel 221 | args: Namespace 222 | """ 223 | 224 | device = hand_model.device 225 | n_objects = len(object_model.object_mesh_list) 226 | batch_size_each = object_model.batch_size_each 227 | total_batch_size = n_objects * batch_size_each 228 | 229 | # initialize translation and rotation 230 | 231 | translation = torch.zeros([total_batch_size, 3], dtype=torch.float, device=device) 232 | rotation = torch.zeros([total_batch_size, 3, 3], dtype=torch.float, device=device) 233 | 234 | for i in range(n_objects): 235 | # get inflated convex hull 236 | 237 | mesh_origin = object_model.object_mesh_list[i].convex_hull 238 | vertices = mesh_origin.vertices.copy() 239 | faces = mesh_origin.faces 240 | vertices *= object_model.object_scale_tensor[i].max().item() 241 | mesh_origin = tm.Trimesh(vertices, faces) 242 | mesh_origin.faces = mesh_origin.faces[mesh_origin.remove_degenerate_faces()] 243 | vertices += 0.2 * vertices / np.linalg.norm(vertices, axis=1, keepdims=True) 244 | mesh = tm.Trimesh(vertices=vertices, faces=faces).convex_hull 245 | vertices = torch.tensor(mesh.vertices, dtype=torch.float, device=device) 246 | faces = torch.tensor(mesh.faces, dtype=torch.float, device=device) 247 | mesh_pytorch3d = pytorch3d.structures.Meshes( 248 | vertices.unsqueeze(0), faces.unsqueeze(0) 249 | ) 250 | 251 | # sample points 252 | 253 | dense_point_cloud = pytorch3d.ops.sample_points_from_meshes( 254 | mesh_pytorch3d, num_samples=100 * batch_size_each 255 | ) 256 | p = pytorch3d.ops.sample_farthest_points(dense_point_cloud, K=batch_size_each)[ 257 | 0 258 | ][0] 259 | closest_points, _, _ = mesh_origin.nearest.on_surface(p.detach().cpu().numpy()) 260 | closest_points = torch.tensor(closest_points, dtype=torch.float, device=device) 261 | n = (closest_points - p) / (closest_points - p).norm(dim=1).unsqueeze(1) 262 | 263 | # sample parameters 264 | 265 | distance = args.distance_lower + ( 266 | args.distance_upper - args.distance_lower 267 | ) * torch.rand([batch_size_each], dtype=torch.float, device=device) 268 | deviate_theta = args.theta_lower + ( 269 | args.theta_upper - args.theta_lower 270 | ) * torch.rand([batch_size_each], dtype=torch.float, device=device) 271 | process_theta = ( 272 | 2 273 | * math.pi 274 | * torch.rand([batch_size_each], dtype=torch.float, device=device) 275 | ) 276 | rotate_theta = ( 277 | 2 278 | * math.pi 279 | * torch.rand([batch_size_each], dtype=torch.float, device=device) 280 | ) 281 | 282 | # solve transformation 283 | # rotation_hand: rotate the hand to align its grasping direction with the +z axis 284 | # rotation_local: jitter the hand's orientation in a cone 285 | # rotation_global and translation: transform the hand to a position corresponding to point p sampled from the inflated convex hull 286 | 287 | rotation_local = torch.zeros( 288 | [batch_size_each, 3, 3], dtype=torch.float, device=device 289 | ) 290 | rotation_global = torch.zeros( 291 | [batch_size_each, 3, 3], dtype=torch.float, device=device 292 | ) 293 | for j in range(batch_size_each): 294 | rotation_local[j] = torch.tensor( 295 | transforms3d.euler.euler2mat( 296 | process_theta[j], deviate_theta[j], rotate_theta[j], axes="rzxz" 297 | ), 298 | dtype=torch.float, 299 | device=device, 300 | ) 301 | rotation_global[j] = torch.tensor( 302 | transforms3d.euler.euler2mat( 303 | math.atan2(n[j, 1], n[j, 0]) - math.pi / 2, 304 | -math.acos(n[j, 2]), 305 | 0, 306 | axes="rzxz", 307 | ), 308 | dtype=torch.float, 309 | device=device, 310 | ) 311 | translation[i * batch_size_each : (i + 1) * batch_size_each] = ( 312 | p 313 | - distance.unsqueeze(1) 314 | * ( 315 | rotation_global 316 | @ rotation_local 317 | @ torch.tensor([0, 0, 1], dtype=torch.float, device=device).reshape( 318 | 1, -1, 1 319 | ) 320 | ).squeeze(2) 321 | ) 322 | rotation[i * batch_size_each : (i + 1) * batch_size_each] = ( 323 | rotation_global @ rotation_local 324 | ) 325 | 326 | translation_hand = torch.tensor([-0.1, -0.05, 0], dtype=torch.float, device=device) 327 | rotation_hand = torch.tensor( 328 | transforms3d.euler.euler2mat(-np.pi / 2, -np.pi / 2, np.pi / 6, axes="rzxz"), 329 | dtype=torch.float, 330 | device=device, 331 | ) 332 | 333 | translation = translation + rotation @ translation_hand 334 | rotation = rotation @ rotation_hand 335 | 336 | # initialize thetas 337 | # thetas_mu: hand-crafted canonicalized hand articulation 338 | # use normal distribution to jitter the thetas 339 | 340 | thetas_mu = ( 341 | torch.tensor( 342 | [ 343 | 0, 344 | 0, 345 | torch.pi / 6, 346 | 0, 347 | 0, 348 | 0, 349 | 0, 350 | 0, 351 | 0, 352 | 0, 353 | 0, 354 | torch.pi / 6, 355 | 0, 356 | 0, 357 | 0, 358 | 0, 359 | 0, 360 | 0, 361 | 0, 362 | 0, 363 | torch.pi / 6, 364 | 0, 365 | 0, 366 | 0, 367 | 0, 368 | 0, 369 | 0, 370 | 0, 371 | 0, 372 | torch.pi / 6, 373 | 0, 374 | 0, 375 | 0, 376 | 0, 377 | 0, 378 | 0, 379 | *(torch.pi / 2 * torch.tensor([1, 0.5, 0], dtype=torch.float)), 380 | 0, 381 | 0, 382 | 0, 383 | 0, 384 | 0, 385 | 0, 386 | ], 387 | dtype=torch.float, 388 | device=device, 389 | ) 390 | .unsqueeze(0) 391 | .repeat(total_batch_size, 1) 392 | ) 393 | if hand_model.left_hand: 394 | thetas_mu = ( 395 | torch.tensor( 396 | [ 397 | 0, 398 | 0, 399 | -torch.pi / 6, 400 | 0, 401 | 0, 402 | 0, 403 | 0, 404 | 0, 405 | 0, 406 | 0, 407 | 0, 408 | -torch.pi / 6, 409 | 0, 410 | 0, 411 | 0, 412 | 0, 413 | 0, 414 | 0, 415 | 0, 416 | 0, 417 | -torch.pi / 6, 418 | 0, 419 | 0, 420 | 0, 421 | 0, 422 | 0, 423 | 0, 424 | 0, 425 | 0, 426 | -torch.pi / 6, 427 | 0, 428 | 0, 429 | 0, 430 | 0, 431 | 0, 432 | 0, 433 | *(torch.pi / 2 * torch.tensor([1, -0.5, 0], dtype=torch.float)), 434 | 0, 435 | 0, 436 | 0, 437 | 0, 438 | 0, 439 | 0, 440 | ], 441 | dtype=torch.float, 442 | device=device, 443 | ) 444 | .unsqueeze(0) 445 | .repeat(total_batch_size, 1) 446 | ) 447 | thetas_sigma = args.jitter_strength * torch.ones( 448 | [total_batch_size, 45], dtype=torch.float, device=device 449 | ) 450 | thetas = torch.normal(thetas_mu, thetas_sigma) 451 | if no_fc: 452 | thetas *= 1e-2 453 | 454 | rotation = pytorch3d.transforms.quaternion_to_axis_angle( 455 | pytorch3d.transforms.matrix_to_quaternion(rotation) 456 | ) 457 | hand_pose = torch.cat( 458 | [ 459 | translation, 460 | rotation, 461 | thetas, 462 | ], 463 | dim=1, 464 | ) 465 | hand_pose.requires_grad_() 466 | 467 | # initialize contact point indices 468 | 469 | contact_point_indices = torch.randint( 470 | hand_model.n_contact_candidates, 471 | size=[total_batch_size, args.n_contact], 472 | device=device, 473 | ) 474 | 475 | hand_model.set_parameters(hand_pose, contact_point_indices) 476 | --------------------------------------------------------------------------------