├── .gitignore ├── LICENSE ├── README.md ├── apps ├── IFGeo.py ├── Normal.py ├── animation.py ├── avatarizer.py ├── benchmark.py ├── blender_dance.py ├── infer.py ├── multi_render.py └── sapiens.py ├── assets ├── OOD-outfits.jpg ├── OOD-poses.jpg ├── SHHQ.gif ├── animation.gif ├── back-45.gif ├── blender-demo.gif ├── crowd.gif ├── dataset.png ├── double-90.gif ├── front-45.gif ├── register.png ├── sapiens │ ├── normal-econ.png │ ├── normal-sapiens.png │ ├── recon-econ.png │ └── recon-sapiens.png └── teaser.gif ├── configs └── econ.yaml ├── docker-compose.yaml ├── docs ├── installation-docker.md ├── installation-ubuntu.md ├── installation-windows.md ├── testing.md └── tricks.md ├── environment-windows.yaml ├── environment.yaml ├── examples ├── 304e9c4798a8c3967de7c74c24ef2e38.jpg ├── cloth │ ├── 0a64d9c7ac4a86aa0c29195bc6f55246.jpg │ ├── 1f7c9214b80a02071edfadd5be908d8e.jpg │ ├── 2095f20b1390e14d9312913c61c4b621.png │ ├── 267cffcff3809e0df9eff44c443f07b0.jpg │ ├── 26d2e846349647ff04c536816e0e8ca1.jpg │ ├── 351f52b9d1ddebb70241a092af34c2f3.jpg │ ├── 55cc162cc4fcda1df2236847a52db93a.jpg │ ├── 6465c18fc13b862341c33922c79ab490.jpg │ ├── 663dcd6db19490de0b790da430bd5681.jpg │ ├── 6c0a5def2287d9bfa4a42ee0ce9cb7f9.jpg │ ├── 7c6bb9410ea8debe3aca92e299fe2333.jpg │ ├── 930c782d63e180208e0a55754d607f34.jpg │ ├── baf3c945fa6b4349c59953a97740e70f.jpg │ ├── c7ca6894119f235caba568a7e01684af.jpg │ ├── da135ecd046333c0dc437a383325c90b.jpg │ ├── df90cff51a84dd602024ac3aa03ad182.jpg │ └── e80b36c782ce869caef9abb55b37d464.jpg ├── motions │ └── rasputin_smplx.pkl ├── multi │ ├── 1.png │ ├── 2.jpg │ ├── 3.jpg │ └── 4.jpg └── pose │ ├── 02986d0998ce01aa0aa67a99fbd1e09a.jpg │ ├── 105545f93dcaecd13f2e3f01db92331c.jpg │ ├── 1af2662b5026ef82ed0e8b08b6698017.png │ ├── 3745ee0a7f31fafc3dfd3d8bf246f3b8.jpg │ ├── 4ac9ca7a3e34a365c073317f98525add.jpg │ ├── 4d1ed606c3c0a346c8a75507fc81abff.jpg │ ├── 5617dc56d25918217b81f27c98011ea5.jpg │ ├── 5ef3bc939cf82dbd0c541eba41b517c2.jpg │ ├── 68757076df6c98e9d6ba6ed00870daef.jpg │ ├── 6f0029a9592a11530267b3a51413ae74.jpg │ ├── 7530ae51e811b1878fae23ea243a3a30.jpg │ ├── 780047b55ee80b0dc2468ad16cab2278.jpg │ ├── ab2192beaefb58e872ce55099cbed8fe.jpg │ ├── d5241b4de0cebd2722b05d855a5a9ca6.jpg │ ├── d6fcd37df9983973af08af3f9267cd1e.jpg │ └── d7e876bc5f9e8277d58e30bd83c7452f.jpg ├── fetch_data.sh ├── lib ├── __init__.py ├── common │ ├── BNI.py │ ├── BNI_utils.py │ ├── __init__.py │ ├── cloth_extraction.py │ ├── config.py │ ├── imutils.py │ ├── libmesh │ │ ├── inside_mesh.py │ │ ├── setup.py │ │ ├── triangle_hash.cpp │ │ └── triangle_hash.pyx │ ├── libvoxelize │ │ ├── setup.py │ │ ├── tribox2.h │ │ ├── voxelize.c │ │ └── voxelize.pyx │ ├── local_affine.py │ ├── render.py │ ├── render_utils.py │ ├── seg3d_lossless.py │ ├── seg3d_utils.py │ ├── smpl_vert_segmentation.json │ ├── smplx_vert_segmentation.json │ ├── train_util.py │ └── voxelize.py ├── dataset │ ├── EvalDataset.py │ ├── Evaluator.py │ ├── NormalDataset.py │ ├── NormalModule.py │ ├── PointFeat.py │ ├── TestDataset.py │ ├── __init__.py │ ├── body_model.py │ ├── mesh_util.py │ └── tbfo.ttf ├── net │ ├── BasePIFuNet.py │ ├── Discriminator.py │ ├── FBNet.py │ ├── GANLoss.py │ ├── IFGeoNet.py │ ├── IFGeoNet_nobody.py │ ├── NormalNet.py │ ├── __init__.py │ ├── geometry.py │ ├── net_util.py │ └── voxelize.py ├── pixielib │ ├── __init__.py │ ├── models │ │ ├── FLAME.py │ │ ├── SMPLX.py │ │ ├── __init__.py │ │ ├── encoders.py │ │ ├── hrnet.py │ │ ├── lbs.py │ │ ├── moderators.py │ │ └── resnet.py │ ├── pixie.py │ └── utils │ │ ├── array_cropper.py │ │ ├── config.py │ │ ├── renderer.py │ │ ├── rotation_converter.py │ │ ├── tensor_cropper.py │ │ └── util.py ├── pymafx │ ├── configs │ │ └── pymafx_config.yaml │ ├── core │ │ ├── __init__.py │ │ ├── cfgs.py │ │ ├── constants.py │ │ └── path_config.py │ ├── models │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── hmr.py │ │ ├── hr_module.py │ │ ├── maf_extractor.py │ │ ├── pose_resnet.py │ │ ├── pymaf_net.py │ │ ├── res_module.py │ │ ├── smpl.py │ │ └── transformers │ │ │ ├── __init__.py │ │ │ ├── bert │ │ │ ├── __init__.py │ │ │ ├── bert-base-uncased │ │ │ │ └── config.json │ │ │ ├── e2e_body_network.py │ │ │ ├── e2e_hand_network.py │ │ │ ├── file_utils.py │ │ │ ├── modeling_bert.py │ │ │ ├── modeling_graphormer.py │ │ │ └── modeling_utils.py │ │ │ ├── net_utils.py │ │ │ ├── texformer.py │ │ │ ├── tokenlearner.py │ │ │ └── transformer_basics.py │ └── utils │ │ ├── __init__.py │ │ ├── binvox_rw.py │ │ ├── blob.py │ │ ├── cam_params.py │ │ ├── collections.py │ │ ├── colormap.py │ │ ├── common.py │ │ ├── data_loader.py │ │ ├── demo_utils.py │ │ ├── densepose_methods.py │ │ ├── geometry.py │ │ ├── imutils.py │ │ ├── io.py │ │ ├── iuvmap.py │ │ ├── keypoints.py │ │ ├── mesh_generation.py │ │ ├── part_utils.py │ │ ├── pose_tracker.py │ │ ├── pose_utils.py │ │ ├── renderer.py │ │ ├── sample_mesh.py │ │ ├── saver.py │ │ ├── segms.py │ │ ├── smooth_bbox.py │ │ ├── transforms.py │ │ ├── uv_vis.py │ │ └── vis.py ├── smplx │ ├── .gitignore │ ├── LICENSE │ ├── README.md │ ├── __init__.py │ ├── body_models.py │ ├── joint_names.py │ ├── lbs.py │ ├── utils.py │ ├── vertex_ids.py │ └── vertex_joint_selector.py └── torch_utils │ ├── __init__.py │ ├── custom_ops.py │ ├── misc.py │ ├── ops │ ├── __init__.py │ ├── bias_act.cpp │ ├── bias_act.cu │ ├── bias_act.h │ ├── bias_act.py │ ├── conv2d_gradfix.py │ ├── conv2d_resample.py │ ├── fma.py │ ├── fused_act.py │ ├── fused_bias_act.cpp │ ├── fused_bias_act_kernel.cu │ ├── grid_sample_gradfix.py │ ├── native_ops.py │ ├── upfirdn2d.cpp │ ├── upfirdn2d.cu │ ├── upfirdn2d.h │ └── upfirdn2d.py │ ├── persistence.py │ └── training_stats.py ├── loose.txt ├── pose.txt ├── requirements.txt └── setup.cfg /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | debug/ 3 | log/ 4 | results/* 5 | results 6 | .vscode 7 | !.gitignore 8 | .idea 9 | cluster/ 10 | cluster 11 | *.zip 12 | data/ 13 | data 14 | wandb 15 | build 16 | dist 17 | *egg-info 18 | *.so 19 | run.sh 20 | *.log 21 | gradio_cached_examples/ 22 | *.blend 23 | *.blend1 24 | examples/junxuan -------------------------------------------------------------------------------- /apps/IFGeo.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import numpy as np 18 | import pytorch_lightning as pl 19 | import torch 20 | 21 | from lib.common.seg3d_lossless import Seg3dLossless 22 | from lib.common.train_util import * 23 | 24 | torch.backends.cudnn.benchmark = True 25 | 26 | 27 | class IFGeo(pl.LightningModule): 28 | def __init__(self, cfg): 29 | super(IFGeo, self).__init__() 30 | 31 | self.cfg = cfg 32 | self.batch_size = self.cfg.batch_size 33 | self.lr_G = self.cfg.lr_G 34 | 35 | self.use_sdf = cfg.sdf 36 | self.mcube_res = cfg.mcube_res 37 | self.clean_mesh_flag = cfg.clean_mesh 38 | self.overfit = cfg.overfit 39 | 40 | if cfg.dataset.prior_type == "SMPL": 41 | from lib.net.IFGeoNet import IFGeoNet 42 | self.netG = IFGeoNet(cfg) 43 | else: 44 | from lib.net.IFGeoNet_nobody import IFGeoNet 45 | self.netG = IFGeoNet(cfg) 46 | 47 | self.resolutions = ( 48 | np.logspace( 49 | start=5, 50 | stop=np.log2(self.mcube_res), 51 | base=2, 52 | num=int(np.log2(self.mcube_res) - 4), 53 | endpoint=True, 54 | ) + 1.0 55 | ) 56 | 57 | self.resolutions = self.resolutions.astype(np.int16).tolist() 58 | 59 | self.reconEngine = Seg3dLossless( 60 | query_func=query_func_IF, 61 | b_min=[[-1.0, 1.0, -1.0]], 62 | b_max=[[1.0, -1.0, 1.0]], 63 | resolutions=self.resolutions, 64 | align_corners=True, 65 | balance_value=0.50, 66 | visualize=False, 67 | debug=False, 68 | use_cuda_impl=False, 69 | faster=True, 70 | ) 71 | 72 | self.export_dir = None 73 | self.result_eval = {} 74 | 75 | # Training related 76 | def configure_optimizers(self): 77 | 78 | # set optimizer 79 | weight_decay = self.cfg.weight_decay 80 | momentum = self.cfg.momentum 81 | 82 | optim_params_G = [{"params": self.netG.parameters(), "lr": self.lr_G}] 83 | 84 | if self.cfg.optim == "Adadelta": 85 | 86 | optimizer_G = torch.optim.Adadelta( 87 | optim_params_G, lr=self.lr_G, weight_decay=weight_decay 88 | ) 89 | 90 | elif self.cfg.optim == "Adam": 91 | 92 | optimizer_G = torch.optim.Adam(optim_params_G, lr=self.lr_G, weight_decay=weight_decay) 93 | 94 | elif self.cfg.optim == "RMSprop": 95 | 96 | optimizer_G = torch.optim.RMSprop( 97 | optim_params_G, 98 | lr=self.lr_G, 99 | weight_decay=weight_decay, 100 | momentum=momentum, 101 | ) 102 | 103 | else: 104 | raise NotImplementedError 105 | 106 | # set scheduler 107 | scheduler_G = torch.optim.lr_scheduler.MultiStepLR( 108 | optimizer_G, milestones=self.cfg.schedule, gamma=self.cfg.gamma 109 | ) 110 | 111 | return [optimizer_G], [scheduler_G] 112 | 113 | def training_step(self, batch, batch_idx): 114 | 115 | self.netG.train() 116 | 117 | preds_G = self.netG(batch) 118 | error_G = self.netG.compute_loss(preds_G, batch["labels_geo"]) 119 | 120 | # metrics processing 121 | metrics_log = { 122 | "loss": error_G, 123 | } 124 | 125 | self.log_dict( 126 | metrics_log, prog_bar=True, logger=True, on_step=True, on_epoch=False, sync_dist=True 127 | ) 128 | 129 | return metrics_log 130 | 131 | def training_epoch_end(self, outputs): 132 | 133 | # metrics processing 134 | metrics_log = { 135 | "train/avgloss": batch_mean(outputs, "loss"), 136 | } 137 | 138 | self.log_dict( 139 | metrics_log, 140 | prog_bar=False, 141 | logger=True, 142 | on_step=False, 143 | on_epoch=True, 144 | rank_zero_only=True 145 | ) 146 | 147 | def validation_step(self, batch, batch_idx): 148 | 149 | self.netG.eval() 150 | self.netG.training = False 151 | 152 | preds_G = self.netG(batch) 153 | error_G = self.netG.compute_loss(preds_G, batch["labels_geo"]) 154 | 155 | metrics_log = { 156 | "val/loss": error_G, 157 | } 158 | 159 | self.log_dict( 160 | metrics_log, prog_bar=True, logger=False, on_step=True, on_epoch=False, sync_dist=True 161 | ) 162 | 163 | return metrics_log 164 | 165 | def validation_epoch_end(self, outputs): 166 | 167 | # metrics processing 168 | metrics_log = { 169 | "val/avgloss": batch_mean(outputs, "val/loss"), 170 | } 171 | 172 | self.log_dict( 173 | metrics_log, 174 | prog_bar=False, 175 | logger=True, 176 | on_step=False, 177 | on_epoch=True, 178 | rank_zero_only=True 179 | ) 180 | -------------------------------------------------------------------------------- /apps/animation.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import numpy as np 5 | import torch 6 | from tqdm import tqdm 7 | 8 | from lib.net.geometry import rotation_matrix_to_angle_axis 9 | from lib.smplx.lbs import general_lbs 10 | 11 | # loading cfg file 12 | parser = argparse.ArgumentParser() 13 | parser.add_argument("-n", "--name", type=str, default="") 14 | parser.add_argument("-m", "--motion", type=str, default="rasputin_smplx") 15 | parser.add_argument("-g", "--gpu", type=int, default=0) 16 | args = parser.parse_args() 17 | 18 | device = torch.device(f"cuda:{args.gpu}") 19 | 20 | econ_dict = torch.load(f"./results/econ/cache/{args.name}/econ.pt") 21 | smplx_pkl = np.load(f"./examples/motions/{args.motion}.pkl", allow_pickle=True) 22 | smplx_pose_mat = torch.tensor(smplx_pkl['pred_thetas']) 23 | smplx_transl = smplx_pkl['transl'] 24 | smplx_pose = rotation_matrix_to_angle_axis(smplx_pose_mat.view(-1, 3, 3)).view(-1, 55, 3) 25 | smplx_pose[:, 23:23 + 2] *= 0.0 # remove the pose of eyes 26 | 27 | n_start = 0 28 | n_end = 100 * 25 29 | # n_end = smplx_pose.shape[0] 30 | n_step = 1 31 | 32 | output_dir = f"./results/econ/seq" 33 | os.makedirs(output_dir, exist_ok=True) 34 | 35 | motion_output = {"v_seq": [], "f": None, "normal": None, "rgb": None} 36 | 37 | for oid, fid in enumerate(tqdm(range(n_start, n_end, n_step))): 38 | posed_econ_verts, _ = general_lbs( 39 | pose=smplx_pose.reshape(-1, 55 * 3)[fid:fid + 1].to(device), 40 | v_template=econ_dict["v_template"].to(device), 41 | posedirs=econ_dict["posedirs"].to(device), 42 | J_regressor=econ_dict["J_regressor"].to(device), 43 | parents=econ_dict["parents"].to(device), 44 | lbs_weights=econ_dict["lbs_weights"].to(device), 45 | ) 46 | smplx_verts = posed_econ_verts[0].float().detach().cpu().numpy() 47 | trans_scale = np.array([1.0, 0.1, 0.1]) # to mitigate z-axis jittering 48 | 49 | motion_output["v_seq"].append((smplx_verts + smplx_transl[fid] * trans_scale).astype( 50 | np.float32 51 | )) 52 | 53 | motion_output["v_seq"] = np.stack(motion_output["v_seq"], axis=0) 54 | motion_output["f"] = econ_dict["faces"].astype(np.uint32) 55 | motion_output["normal"] = econ_dict["final_normal"].astype(np.float32) 56 | motion_output["rgb"] = econ_dict["final_rgb"].astype(np.float32) 57 | 58 | np.savez_compressed(f"{output_dir}/{args.name}_motion.npz", **motion_output) 59 | -------------------------------------------------------------------------------- /apps/blender_dance.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | import os 3 | import sys 4 | from tqdm import tqdm 5 | import numpy as np 6 | 7 | argv = sys.argv 8 | argv = argv[argv.index("--") + 1:] # get all args after "--" 9 | render_normal = True if argv[0] == 'normal' else False 10 | avatar_name = argv[1] 11 | duration = int(argv[2]) 12 | 13 | # use-defined parameters 14 | n_start = 0 15 | fps = 25 16 | n_end = fps * duration 17 | n_step = 1 18 | 19 | 20 | class ndarray_pydata(np.ndarray): 21 | def __bool__(self) -> bool: 22 | return len(self) > 0 23 | 24 | 25 | class Character: 26 | def __init__(self, v_seq, faces, colors, name): 27 | 28 | self.v_seq = v_seq 29 | self.faces = faces 30 | self.colors = np.hstack((colors / 255.0, np.ones((colors.shape[0], 1), dtype=np.float32))) 31 | self.material = bpy.data.materials['vertex-color'] 32 | self.name = name 33 | 34 | def load_frame(self, index, delta_trans): 35 | 36 | name = f"{self.name}_{str(index).zfill(4)}" 37 | 38 | # create mesh 39 | mesh = bpy.data.meshes.new(name) 40 | mesh.from_pydata( 41 | self.v_seq[index][:, [0, 2, 1]] * np.array([1., -1., 1.]), [], 42 | self.faces.view(ndarray_pydata) 43 | ) 44 | mesh.vertex_colors.new(name="vcol") 45 | mesh.vertex_colors["vcol"].active = True 46 | mloops = np.zeros((len(mesh.loops)), dtype=np.int) 47 | mesh.loops.foreach_get("vertex_index", mloops) 48 | mesh.vertex_colors["vcol"].data.foreach_set("color", self.colors[mloops].flatten()) 49 | mesh.validate() 50 | 51 | # create object 52 | obj = bpy.data.objects.new(name, mesh) 53 | bpy.context.scene.collection.objects.link(obj) 54 | bpy.ops.object.select_all(action='DESELECT') 55 | obj.select_set(True) 56 | obj.active_material = self.material 57 | obj.delta_location = delta_trans 58 | bpy.context.view_layer.objects.active = obj 59 | bpy.ops.object.shade_smooth() 60 | 61 | return obj, mesh 62 | 63 | def __len__(self): 64 | return self.v_seq.shape[0] 65 | 66 | 67 | root_dir = os.path.join(os.path.dirname(__file__), "..") 68 | avatar_pos = { 69 | avatar_name: np.array([-0.7, -7.5, 0.]), 70 | } 71 | avatar_names = avatar_pos.keys() 72 | 73 | # load blend file 74 | blend_path = f"{root_dir}/econ_empty.blend" 75 | bpy.ops.wm.open_mainfile(filepath=blend_path) 76 | 77 | # rendering settings 78 | bpy.context.scene.render.image_settings.file_format = 'PNG' 79 | bpy.context.scene.eevee.taa_render_samples = 128 80 | 81 | # load all the large motion data 82 | avatar_data = {} 83 | pbar = tqdm(avatar_names) 84 | for key in pbar: 85 | pbar.set_description(f"Loading {key}") 86 | motion_path = f"{root_dir}/results/econ/seq/{key}_motion.npz" 87 | motion = np.load(motion_path, allow_pickle=True) 88 | if render_normal: 89 | avatar_data[key] = Character(motion['v_seq'], motion['f'], motion['normal'], name=key) 90 | else: 91 | avatar_data[key] = Character(motion['v_seq'], motion['f'], motion['rgb'], name=key) 92 | 93 | export_dir = f"{root_dir}/results/econ/render/{argv[0]}" 94 | os.makedirs(export_dir, exist_ok=True) 95 | 96 | # start rendering 97 | for fid in tqdm(range(n_start, n_end, n_step)): 98 | objs = [] 99 | meshes = [] 100 | for key in avatar_names: 101 | obj, mesh = avatar_data[key].load_frame(fid, avatar_pos[key]) 102 | objs.append(obj) 103 | meshes.append(mesh) 104 | 105 | bpy.context.scene.frame_set(fid) 106 | bpy.context.scene.render.filepath = f"{export_dir}/{str(fid).zfill(5)}.png" 107 | bpy.ops.render.render(use_viewport=True, write_still=True) 108 | bpy.ops.object.select_all(action='SELECT') 109 | 110 | for mesh in meshes: 111 | bpy.data.meshes.remove(mesh) 112 | 113 | # Debug: you can open the blend file to check the blender scene 114 | # if fid == 10: 115 | # bpy.ops.wm.save_as_mainfile(filepath=f"{root_dir}/test.blend") 116 | 117 | # combine all the rendering images into a video 118 | from moviepy.editor import ImageSequenceClip 119 | from glob import glob 120 | from os import cpu_count 121 | 122 | mpy_conf = { 123 | "codec": "libx264", 124 | "remove_temp": True, 125 | "preset": "ultrafast", 126 | "ffmpeg_params": ['-crf', '0', '-pix_fmt', 'yuv444p', '-profile:v', 'high444'], 127 | "logger": "bar", 128 | "fps": fps, 129 | "threads": cpu_count(), 130 | } 131 | 132 | video_lst = sorted(glob(f"{root_dir}/results/econ/render/{argv[0]}/*.png")) 133 | video_clip = ImageSequenceClip(video_lst, fps=fps) 134 | video_clip = video_clip.set_duration(duration) 135 | video_clip.write_videofile(f"{root_dir}/results/econ/render/{argv[0]}.mp4", **mpy_conf) 136 | -------------------------------------------------------------------------------- /apps/multi_render.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import torch 4 | 5 | from lib.common.render import Render 6 | 7 | root = "./results/econ/vid" 8 | 9 | # loading cfg file 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument("-n", "--name", type=str, default="") 12 | parser.add_argument("-g", "--gpu", type=int, default=0) 13 | args = parser.parse_args() 14 | 15 | in_tensor = torch.load(f"{root}/{args.name}_in_tensor.pt") 16 | 17 | render = Render(size=512, device=torch.device(f"cuda:{args.gpu}")) 18 | 19 | # visualize the final results in self-rotation mode 20 | verts_lst = in_tensor["body_verts"] + in_tensor["BNI_verts"] 21 | faces_lst = in_tensor["body_faces"] + in_tensor["BNI_faces"] 22 | 23 | # self-rotated video 24 | render.load_meshes(verts_lst, faces_lst) 25 | render.get_rendered_video_multi(in_tensor, f"{root}/{args.name}_cloth.mp4") 26 | -------------------------------------------------------------------------------- /apps/sapiens.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.nn.functional as F 4 | from PIL import Image 5 | from torchvision import transforms 6 | from huggingface_hub import snapshot_download 7 | 8 | 9 | class Config: 10 | ASSETS_DIR = os.path.join( 11 | snapshot_download(repo_id="facebook/sapiens-normal", repo_type="space"), 'assets' 12 | ) 13 | CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints") 14 | CHECKPOINTS = { 15 | "0.3b": "sapiens_0.3b_normal_render_people_epoch_66_torchscript.pt2", 16 | "0.6b": "sapiens_0.6b_normal_render_people_epoch_200_torchscript.pt2", 17 | "1b": "sapiens_1b_normal_render_people_epoch_115_torchscript.pt2", 18 | "2b": "sapiens_2b_normal_render_people_epoch_70_torchscript.pt2", 19 | } 20 | SEG_CHECKPOINTS = { 21 | "fg-bg-1b": "sapiens_1b_seg_foreground_epoch_8_torchscript.pt2", 22 | "no-bg-removal": None, 23 | "part-seg-1b": "sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2", 24 | } 25 | 26 | 27 | class ModelManager: 28 | @staticmethod 29 | def load_model(checkpoint_name, device): 30 | if checkpoint_name is None: 31 | return None 32 | checkpoint_path = os.path.join(Config.CHECKPOINTS_DIR, checkpoint_name) 33 | model = torch.jit.load(checkpoint_path) 34 | model.eval() 35 | model.to(device) 36 | return model 37 | 38 | @staticmethod 39 | @torch.inference_mode() 40 | def run_model(model, input_tensor, height, width): 41 | output = model(input_tensor) 42 | return F.interpolate(output, size=(height, width), mode="bilinear", align_corners=False) 43 | 44 | 45 | class ImageProcessor: 46 | def __init__(self, device): 47 | 48 | self.mean = [123.5 / 255, 116.5 / 255, 103.5 / 255] 49 | self.std = [58.5 / 255, 57.0 / 255, 57.5 / 255] 50 | 51 | self.transform_fn = transforms.Compose([ 52 | transforms.Resize((1024, 768)), 53 | transforms.ToTensor(), 54 | transforms.Normalize(mean=self.mean, std=self.std), 55 | ]) 56 | self.device = device 57 | 58 | def process_image(self, image: Image.Image, normal_model_name: str, seg_model_name: str): 59 | 60 | # Load models here instead of storing them as class attributes 61 | normal_model = ModelManager.load_model(Config.CHECKPOINTS[normal_model_name], self.device) 62 | input_tensor = self.transform_fn(image).unsqueeze(0).to(self.device) 63 | 64 | # Run normal estimation 65 | normal_map = ModelManager.run_model(normal_model, input_tensor, image.height, image.width) 66 | 67 | # Run segmentation 68 | if seg_model_name != "no-bg-removal": 69 | seg_model = ModelManager.load_model(Config.SEG_CHECKPOINTS[seg_model_name], self.device) 70 | seg_output = ModelManager.run_model(seg_model, input_tensor, image.height, image.width) 71 | seg_mask = (seg_output.argmax(dim=1) > 0).unsqueeze(0).repeat(1, 3, 1, 1) 72 | 73 | # Normalize and visualize normal map 74 | normal_map_norm = torch.linalg.norm(normal_map, dim=1, keepdim=True) 75 | normal_map_normalized = normal_map / (normal_map_norm + 1e-5) 76 | normal_map_normalized[seg_mask == 0] = 0.0 77 | normal_map_normalized = normal_map_normalized.to(self.device) 78 | 79 | return normal_map_normalized 80 | -------------------------------------------------------------------------------- /assets/OOD-outfits.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/OOD-outfits.jpg -------------------------------------------------------------------------------- /assets/OOD-poses.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/OOD-poses.jpg -------------------------------------------------------------------------------- /assets/SHHQ.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/SHHQ.gif -------------------------------------------------------------------------------- /assets/animation.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/animation.gif -------------------------------------------------------------------------------- /assets/back-45.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/back-45.gif -------------------------------------------------------------------------------- /assets/blender-demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/blender-demo.gif -------------------------------------------------------------------------------- /assets/crowd.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/crowd.gif -------------------------------------------------------------------------------- /assets/dataset.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/dataset.png -------------------------------------------------------------------------------- /assets/double-90.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/double-90.gif -------------------------------------------------------------------------------- /assets/front-45.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/front-45.gif -------------------------------------------------------------------------------- /assets/register.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/register.png -------------------------------------------------------------------------------- /assets/sapiens/normal-econ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/sapiens/normal-econ.png -------------------------------------------------------------------------------- /assets/sapiens/normal-sapiens.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/sapiens/normal-sapiens.png -------------------------------------------------------------------------------- /assets/sapiens/recon-econ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/sapiens/recon-econ.png -------------------------------------------------------------------------------- /assets/sapiens/recon-sapiens.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/sapiens/recon-sapiens.png -------------------------------------------------------------------------------- /assets/teaser.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/assets/teaser.gif -------------------------------------------------------------------------------- /configs/econ.yaml: -------------------------------------------------------------------------------- 1 | name: econ 2 | ckpt_dir: "./data/ckpt/" 3 | normal_path: "./data/ckpt/normal.ckpt" 4 | ifnet_path: "./data/ckpt/ifnet.ckpt" 5 | results_path: "./results" 6 | 7 | net: 8 | in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3)) 9 | in_geo: (('normal_F',3), ('normal_B',3)) 10 | 11 | test_mode: True 12 | batch_size: 1 13 | 14 | dataset: 15 | prior_type: "SMPL" 16 | 17 | vol_res: 256 18 | mcube_res: 256 19 | clean_mesh: True 20 | cloth_overlap_thres: 0.50 21 | body_overlap_thres: 0.00 22 | force_smpl_optim: True 23 | 24 | 25 | # normal_model could be '1b' if CUDA OOM, see apps/sapiens.py 26 | 27 | sapiens: 28 | use: True 29 | seg_model: "fg-bg-1b" 30 | normal_model: "2b" 31 | 32 | # For crowded / occluded scene 33 | # body_overlap_thres: 0.98 34 | 35 | bni: 36 | k: 4 37 | lambda1: 1e-4 38 | boundary_consist: 1e-6 39 | poisson_depth: 10 40 | use_smpl: ["hand"] 41 | use_ifnet: False 42 | use_poisson: True 43 | hand_thres: 8e-2 44 | face_thres: 6e-2 45 | thickness: 0.02 46 | hps_type: "pixie" 47 | texture_src: "image" 48 | cut_intersection: True 49 | -------------------------------------------------------------------------------- /docker-compose.yaml: -------------------------------------------------------------------------------- 1 | # build Image from Docker Hub 2 | version: "2.4" 3 | services: 4 | econ: 5 | container_name: econ-container 6 | image: teddy12155555/econ:v1 7 | runtime: nvidia 8 | environment: 9 | - NVIDIA_VISIBLE_DEVICES=all 10 | - DISPLAY=${DISPLAY} 11 | stdin_open: true 12 | tty: true 13 | volumes: 14 | - .:/root/code 15 | - /tmp/.X11-unix:/tmp/.X11-unix 16 | ports: 17 | - "8000:8000" 18 | privileged: true 19 | command: "bash" 20 | -------------------------------------------------------------------------------- /docs/installation-docker.md: -------------------------------------------------------------------------------- 1 | ## Getting started 2 | 3 | Start by cloning the repo: 4 | 5 | ```bash 6 | git clone git@github.com:YuliangXiu/ECON.git 7 | cd ECON 8 | ``` 9 | ## Environment 10 | - **GPU Memory > 12GB** 11 | 12 | start with [docker compose](https://docs.docker.com/compose/) 13 | ```bash 14 | # you can change your container name by passing --name "parameter" 15 | docker compose run [--name myecon] econ 16 | ``` 17 | 18 | ## Docker container's shell 19 | ```bash 20 | # activate the pre-build env 21 | cd code 22 | conda activate econ 23 | 24 | # install libmesh & libvoxelize 25 | cd lib/common/libmesh 26 | python setup.py build_ext --inplace 27 | cd ../libvoxelize 28 | python setup.py build_ext --inplace 29 | ``` 30 | 31 | ## Register at [ICON's website](https://icon.is.tue.mpg.de/) 32 | 33 | ![Register](../assets/register.png) 34 | Required: 35 | 36 | - [SMPL](http://smpl.is.tue.mpg.de/): SMPL Model (Male, Female) 37 | - [SMPL-X](http://smpl-x.is.tue.mpg.de/): SMPL-X Model, used for training 38 | - [SMPLIFY](http://smplify.is.tue.mpg.de/): SMPL Model (Neutral) 39 | - [PIXIE](https://icon.is.tue.mpg.de/user.php): PIXIE SMPL-X estimator 40 | 41 | :warning: Click **Register now** on all dependencies, then you can download them all with **ONE** account. 42 | 43 | ## Downloading required models and extra data 44 | 45 | ```bash 46 | cd ~/code 47 | bash fetch_data.sh # requires username and password 48 | ``` 49 | ## :whale2: **todo** 50 | - **Image Environment Infos** 51 | - Ubuntu 18 52 | - CUDA = 11.3 53 | - Python = 3.8 54 | - [X] pre-built image with docker compose 55 | - [ ] docker run command, Dockerfile 56 | - [ ] verify on WSL (Windows) 57 | 58 | ## Citation 59 | 60 | :+1: Please consider citing these awesome HPS approaches: PyMAF-X, PIXIE 61 | 62 | 63 | ``` 64 | @article{pymafx2022, 65 | title={PyMAF-X: Towards Well-aligned Full-body Model Regression from Monocular Images}, 66 | author={Zhang, Hongwen and Tian, Yating and Zhang, Yuxiang and Li, Mengcheng and An, Liang and Sun, Zhenan and Liu, Yebin}, 67 | journal={arXiv preprint arXiv:2207.06400}, 68 | year={2022} 69 | } 70 | 71 | 72 | @inproceedings{PIXIE:2021, 73 | title={Collaborative Regression of Expressive Bodies using Moderation}, 74 | author={Yao Feng and Vasileios Choutas and Timo Bolkart and Dimitrios Tzionas and Michael J. Black}, 75 | booktitle={International Conference on 3D Vision (3DV)}, 76 | year={2021} 77 | } 78 | 79 | 80 | ``` 81 | -------------------------------------------------------------------------------- /docs/installation-ubuntu.md: -------------------------------------------------------------------------------- 1 | ## Getting started 2 | 3 | Start by cloning the repo: 4 | 5 | ```bash 6 | git clone git@github.com:YuliangXiu/ECON.git 7 | cd ECON 8 | ``` 9 | 10 | ## Environment 11 | 12 | - Ubuntu 20 / 18, (Windows as well, see [issue#7](https://github.com/YuliangXiu/ECON/issues/7)) 13 | - **CUDA=11.6, GPU Memory > 12GB** 14 | - Python = 3.8 15 | - PyTorch >= 1.13.0 (official [Get Started](https://pytorch.org/get-started/locally/)) 16 | - Cupy >= 11.3.0 (offcial [Installation](https://docs.cupy.dev/en/stable/install.html#installing-cupy-from-pypi)) 17 | - PyTorch3D = 0.7.2 (official [INSTALL.md](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md), recommend [install-from-local-clone](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#2-install-from-a-local-clone)) 18 | 19 | ```bash 20 | 21 | sudo apt-get install libeigen3-dev ffmpeg 22 | 23 | # install required packages 24 | cd ECON 25 | conda env create -f environment.yaml 26 | conda activate econ 27 | pip install -r requirements.txt 28 | 29 | # the installation(incl. compilation) of PyTorch3D will take ~20min 30 | pip install git+https://github.com/facebookresearch/pytorch3d.git@v0.7.2 31 | 32 | # install libmesh & libvoxelize 33 | cd lib/common/libmesh 34 | python setup.py build_ext --inplace 35 | cd ../libvoxelize 36 | python setup.py build_ext --inplace 37 | ``` 38 | 39 | ## Register at [ICON's website](https://icon.is.tue.mpg.de/) 40 | 41 | ![Register](../assets/register.png) 42 | Required: 43 | 44 | - [SMPL](http://smpl.is.tue.mpg.de/): SMPL Model (Male, Female) 45 | - [SMPL-X](http://smpl-x.is.tue.mpg.de/): SMPL-X Model, used for training 46 | - [SMPLIFY](http://smplify.is.tue.mpg.de/): SMPL Model (Neutral) 47 | - [PIXIE](https://icon.is.tue.mpg.de/user.php): PIXIE SMPL-X estimator 48 | 49 | :warning: Click **Register now** on all dependencies, then you can download them all with **ONE** account. 50 | 51 | ## Downloading required models and extra data 52 | 53 | ```bash 54 | cd ECON 55 | bash fetch_data.sh # requires username and password 56 | ``` 57 | 58 | ## Citation 59 | 60 | :+1: Please consider citing these awesome HPS approaches: PyMAF-X, PIXIE 61 | 62 | 63 | ``` 64 | @article{pymafx2022, 65 | title={PyMAF-X: Towards Well-aligned Full-body Model Regression from Monocular Images}, 66 | author={Zhang, Hongwen and Tian, Yating and Zhang, Yuxiang and Li, Mengcheng and An, Liang and Sun, Zhenan and Liu, Yebin}, 67 | journal={arXiv preprint arXiv:2207.06400}, 68 | year={2022} 69 | } 70 | 71 | 72 | @inproceedings{PIXIE:2021, 73 | title={Collaborative Regression of Expressive Bodies using Moderation}, 74 | author={Yao Feng and Vasileios Choutas and Timo Bolkart and Dimitrios Tzionas and Michael J. Black}, 75 | booktitle={International Conference on 3D Vision (3DV)}, 76 | year={2021} 77 | } 78 | 79 | 80 | ``` 81 | -------------------------------------------------------------------------------- /docs/installation-windows.md: -------------------------------------------------------------------------------- 1 | # Windows installation tutorial 2 | 3 | Another [issue#16](https://github.com/YuliangXiu/ECON/issues/16) shows the whole process to deploy ECON on _Windows_ 4 | 5 | ## Dependencies and Installation 6 | 7 | - Use [Anaconda](https://www.anaconda.com/products/distribution) 8 | - NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads) 9 | - [Wget for Windows](https://eternallybored.org/misc/wget/1.21.3/64/wget.exe) 10 | - Create a new folder on your C drive and rename it "wget" and move the downloaded "wget.exe" over there. 11 | - Add the path to your wget folder to your system environment variables at `Environment Variables > System Variables Path > Edit environment variable` 12 | 13 | ![image](https://user-images.githubusercontent.com/34035011/210986038-39dbb7a1-12ef-4be9-9af4-5f658c6beb65.png) 14 | 15 | - Install [Git for Windows 64-bit](https://git-scm.com/download/win) 16 | - [Visual Studio Community 2022](https://visualstudio.microsoft.com/) (Make sure to check all the boxes as shown in the image below) 17 | 18 | ![image](https://user-images.githubusercontent.com/34035011/210983023-4e5a0024-68f0-4adb-8089-6ff598aec220.PNG) 19 | 20 | ## Getting started 21 | 22 | Start by cloning the repo: 23 | 24 | ```bash 25 | git clone https://github.com/yuliangxiu/ECON.git 26 | cd ECON 27 | ``` 28 | 29 | ## Environment 30 | 31 | - Windows 10 / 11 32 | - **CUDA=11.3** 33 | - Python = 3.8 34 | - PyTorch >= 1.12.1 (official [Get Started](https://pytorch.org/get-started/locally/)) 35 | - Cupy >= 11.3.0 (offcial [Installation](https://docs.cupy.dev/en/stable/install.html#installing-cupy-from-pypi)) 36 | - PyTorch3D = 0.7.1 (official [INSTALL.md](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md), recommend [install-from-local-clone](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#2-install-from-a-local-clone)) 37 | 38 | ```bash 39 | # install required packages 40 | cd ECON 41 | conda env create -f environment-windows.yaml 42 | conda activate econ 43 | 44 | # install pytorch and cupy 45 | pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 46 | pip install -r requirements.txt 47 | pip install cupy-cuda11x 48 | pip install git+https://github.com/facebookresearch/pytorch3d.git@v0.7.1 49 | 50 | # install libmesh & libvoxelize 51 | cd lib/common/libmesh 52 | python setup.py build_ext --inplace 53 | cd ../libvoxelize 54 | python setup.py build_ext --inplace 55 | ``` 56 | 57 | [Issue#69: Discussion of additional argument `--compiler=msvc` in `python setup.py build_ext --inplace`](https://github.com/YuliangXiu/ECON/issues/69) 58 | 59 |
60 | 61 | ## Register at [ICON's website](https://icon.is.tue.mpg.de/) 62 | 63 | ![Register](../assets/register.png) 64 | Required: 65 | 66 | - [SMPL](http://smpl.is.tue.mpg.de/): SMPL Model (Male, Female) 67 | - [SMPL-X](http://smpl-x.is.tue.mpg.de/): SMPL-X Model, used for training 68 | - [SMPLIFY](http://smplify.is.tue.mpg.de/): SMPL Model (Neutral) 69 | - [PIXIE](https://icon.is.tue.mpg.de/user.php): PIXIE SMPL-X estimator 70 | 71 | :warning: Click **Register now** on all dependencies, then you can download them all with **ONE** account. 72 | 73 | ## Downloading required models and extra data (make sure to install git and wget for windows for this to work) 74 | 75 | ```bash 76 | cd ECON 77 | bash fetch_data.sh # requires username and password 78 | ``` 79 | 80 | ## Citation 81 | 82 | :+1: Please consider citing these awesome HPS approaches: PyMAF-X, PIXIE 83 | 84 | ``` 85 | @article{pymafx2022, 86 | title={PyMAF-X: Towards Well-aligned Full-body Model Regression from Monocular Images}, 87 | author={Zhang, Hongwen and Tian, Yating and Zhang, Yuxiang and Li, Mengcheng and An, Liang and Sun, Zhenan and Liu, Yebin}, 88 | journal={arXiv preprint arXiv:2207.06400}, 89 | year={2022} 90 | } 91 | 92 | 93 | @inproceedings{PIXIE:2021, 94 | title={Collaborative Regression of Expressive Bodies using Moderation}, 95 | author={Yao Feng and Vasileios Choutas and Timo Bolkart and Dimitrios Tzionas and Michael J. Black}, 96 | booktitle={International Conference on 3D Vision (3DV)}, 97 | year={2021} 98 | } 99 | 100 | 101 | ``` 102 | -------------------------------------------------------------------------------- /docs/tricks.md: -------------------------------------------------------------------------------- 1 | ## Technical tricks to improve or accelerate ECON 2 | 3 | ### If the reconstructed geometry is not satisfying, play with the adjustable parameters in _config/econ.yaml_ 4 | 5 | - `use_smpl: ["hand"]` 6 | - [ ]: don't use either hands or face parts from SMPL-X 7 | - ["hand"]: only use the **visible** hands from SMPL-X 8 | - ["hand", "face"]: use both **visible** hands and face from SMPL-X 9 | - `thickness: 2cm` 10 | - could be increased accordingly in case final reconstruction **xx_full.obj** looks flat 11 | - `k: 4` 12 | - could be reduced accordingly in case the surface of **xx_full.obj** has discontinous artifacts 13 | - `hps_type: PIXIE` 14 | - "pixie": more accurate for face and hands 15 | - "pymafx": more robust for challenging poses 16 | - `texture_src: image` 17 | - "image": direct mapping the aligned pixels to final mesh 18 | - "SD": use Stable Diffusion to generate full texture (TODO) 19 | 20 | ### To accelerate the inference, you could 21 | 22 | - `use_ifnet: False` 23 | - True: use IF-Nets+ for mesh completion ( $\text{ECON}_\text{IF}$ - Better quality, **~2min / img**) 24 | - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed, **~1.8min / img**) 25 | 26 | ```bash 27 | # For single-person image-based reconstruction (w/o all visualization steps, 1.5min) 28 | python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./results -novis 29 | ``` 30 | 31 | ### Bending legs 32 | 33 | Related issues: 34 | - ECON: https://github.com/YuliangXiu/ECON/issues/133, https://github.com/YuliangXiu/ECON/issues/5 35 | - ICON: https://github.com/YuliangXiu/ICON/issues/68 36 | - TeCH: https://github.com/huangyangyi/TeCH/issues/14 37 | 38 | Reasons: 39 | 40 | - Most existing human pose estimators (HPS) have some bias towards a mean pose, which has the legs bent 41 | - The pseudo GT of HPS is obtained by fitting SMPL(-X) onto 2D landmarks, 2D landmarks has depth ambiguity, which make legs bent 42 | 43 | Solution: 44 | 45 | - The issue of "bending legs" could be significantly reduced by refining the pose of the SMPL(-X) body using an accurate **front clothed normal map** directly estimated from the input image. For this, we use [Sapiens](https://rawalkhirodkar.github.io/sapiens/). 46 | - The **back clothed normal map** is conditioned on the back body normal map, which is rendered from the SMPL(-X) body model. Therefore, the accuracy of the back clothed normal map benefits from the refined SMPL(-X) body model. 47 | 48 | Potential risk: 49 | 50 | - For extremely challenging poses, ECON's normal estimator might outperform Sapiens. ECON's normal estimation is conditioned on the SMPL(-X) body model, which encodes body articulation information, leaving only the normal disparity between the body and clothing to be estimated. In contrast, Sapiens directly estimates normal maps from the input image without SMPL(-X) conditioning. Thus, Sapiens might struggle with particularly challenging or rare poses. 51 | 52 | 53 | 54 | | Final Reconstruction | 55 | | :---------------------------------------------------: | 56 | | w/ sapiens-normal refinement | 57 | | ![recon-sapiens](../assets/sapiens/recon-sapiens.png) | 58 | | Original ECON normal estimator | 59 | | ![recon-econ](../assets/sapiens/recon-econ.png) | 60 | 61 | 62 | | Normal Estimation | 63 | | :-----------------------------------------------------: | 64 | | w/ sapiens-normal refinement | 65 | | ![normal-sapiens](../assets/sapiens/normal-sapiens.png) | 66 | | Original ECON normal estimator | 67 | | ![normal-econ](../assets/sapiens/normal-econ.png) | 68 | 69 | 70 | 71 | If you use Sapiens in your research, please consider citing it. 72 | 73 | 74 | ```bibtex 75 | @inproceedings{khirodkar2024_sapiens, 76 | title={Sapiens: Foundation for Human Vision Models}, 77 | author={Khirodkar, Rawal and Bagautdinov, Timur and Martinez, Julieta and Zhaoen, Su and James, Austin and Selednik, Peter and Anderson, Stuart and Saito, Shunsuke}, 78 | year={2024}, 79 | booktitle={European Conference on Computer Vision}, 80 | } 81 | ``` 82 | 83 | 84 | 85 | 86 | 87 | -------------------------------------------------------------------------------- /environment-windows.yaml: -------------------------------------------------------------------------------- 1 | name: econ 2 | channels: 3 | - nvidia 4 | - pytorch 5 | - conda-forge 6 | - fvcore 7 | - iopath 8 | - bottler 9 | - defaults 10 | dependencies: 11 | - python=3.8 12 | - pytorch-cuda=11.3 13 | - fvcore 14 | - iopath 15 | - cupy 16 | - cython 17 | - pip 18 | 19 | -------------------------------------------------------------------------------- /environment.yaml: -------------------------------------------------------------------------------- 1 | name: econ 2 | channels: 3 | - pytorch 4 | - nvidia 5 | - conda-forge 6 | - fvcore 7 | - iopath 8 | - bottler 9 | - defaults 10 | dependencies: 11 | - python=3.8 12 | - pytorch-cuda=11.6 13 | - pytorch=1.13.0 14 | - nvidiacub 15 | - torchvision 16 | - fvcore 17 | - iopath 18 | - pyembree 19 | - cupy 20 | - cython 21 | - pip -------------------------------------------------------------------------------- /examples/304e9c4798a8c3967de7c74c24ef2e38.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/304e9c4798a8c3967de7c74c24ef2e38.jpg -------------------------------------------------------------------------------- /examples/cloth/0a64d9c7ac4a86aa0c29195bc6f55246.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/0a64d9c7ac4a86aa0c29195bc6f55246.jpg -------------------------------------------------------------------------------- /examples/cloth/1f7c9214b80a02071edfadd5be908d8e.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/1f7c9214b80a02071edfadd5be908d8e.jpg -------------------------------------------------------------------------------- /examples/cloth/2095f20b1390e14d9312913c61c4b621.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/2095f20b1390e14d9312913c61c4b621.png -------------------------------------------------------------------------------- /examples/cloth/267cffcff3809e0df9eff44c443f07b0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/267cffcff3809e0df9eff44c443f07b0.jpg -------------------------------------------------------------------------------- /examples/cloth/26d2e846349647ff04c536816e0e8ca1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/26d2e846349647ff04c536816e0e8ca1.jpg -------------------------------------------------------------------------------- /examples/cloth/351f52b9d1ddebb70241a092af34c2f3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/351f52b9d1ddebb70241a092af34c2f3.jpg -------------------------------------------------------------------------------- /examples/cloth/55cc162cc4fcda1df2236847a52db93a.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/55cc162cc4fcda1df2236847a52db93a.jpg -------------------------------------------------------------------------------- /examples/cloth/6465c18fc13b862341c33922c79ab490.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/6465c18fc13b862341c33922c79ab490.jpg -------------------------------------------------------------------------------- /examples/cloth/663dcd6db19490de0b790da430bd5681.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/663dcd6db19490de0b790da430bd5681.jpg -------------------------------------------------------------------------------- /examples/cloth/6c0a5def2287d9bfa4a42ee0ce9cb7f9.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/6c0a5def2287d9bfa4a42ee0ce9cb7f9.jpg -------------------------------------------------------------------------------- /examples/cloth/7c6bb9410ea8debe3aca92e299fe2333.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/7c6bb9410ea8debe3aca92e299fe2333.jpg -------------------------------------------------------------------------------- /examples/cloth/930c782d63e180208e0a55754d607f34.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/930c782d63e180208e0a55754d607f34.jpg -------------------------------------------------------------------------------- /examples/cloth/baf3c945fa6b4349c59953a97740e70f.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/baf3c945fa6b4349c59953a97740e70f.jpg -------------------------------------------------------------------------------- /examples/cloth/c7ca6894119f235caba568a7e01684af.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/c7ca6894119f235caba568a7e01684af.jpg -------------------------------------------------------------------------------- /examples/cloth/da135ecd046333c0dc437a383325c90b.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/da135ecd046333c0dc437a383325c90b.jpg -------------------------------------------------------------------------------- /examples/cloth/df90cff51a84dd602024ac3aa03ad182.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/df90cff51a84dd602024ac3aa03ad182.jpg -------------------------------------------------------------------------------- /examples/cloth/e80b36c782ce869caef9abb55b37d464.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/cloth/e80b36c782ce869caef9abb55b37d464.jpg -------------------------------------------------------------------------------- /examples/motions/rasputin_smplx.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/motions/rasputin_smplx.pkl -------------------------------------------------------------------------------- /examples/multi/1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/multi/1.png -------------------------------------------------------------------------------- /examples/multi/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/multi/2.jpg -------------------------------------------------------------------------------- /examples/multi/3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/multi/3.jpg -------------------------------------------------------------------------------- /examples/multi/4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/multi/4.jpg -------------------------------------------------------------------------------- /examples/pose/02986d0998ce01aa0aa67a99fbd1e09a.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/02986d0998ce01aa0aa67a99fbd1e09a.jpg -------------------------------------------------------------------------------- /examples/pose/105545f93dcaecd13f2e3f01db92331c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/105545f93dcaecd13f2e3f01db92331c.jpg -------------------------------------------------------------------------------- /examples/pose/1af2662b5026ef82ed0e8b08b6698017.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/1af2662b5026ef82ed0e8b08b6698017.png -------------------------------------------------------------------------------- /examples/pose/3745ee0a7f31fafc3dfd3d8bf246f3b8.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/3745ee0a7f31fafc3dfd3d8bf246f3b8.jpg -------------------------------------------------------------------------------- /examples/pose/4ac9ca7a3e34a365c073317f98525add.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/4ac9ca7a3e34a365c073317f98525add.jpg -------------------------------------------------------------------------------- /examples/pose/4d1ed606c3c0a346c8a75507fc81abff.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/4d1ed606c3c0a346c8a75507fc81abff.jpg -------------------------------------------------------------------------------- /examples/pose/5617dc56d25918217b81f27c98011ea5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/5617dc56d25918217b81f27c98011ea5.jpg -------------------------------------------------------------------------------- /examples/pose/5ef3bc939cf82dbd0c541eba41b517c2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/5ef3bc939cf82dbd0c541eba41b517c2.jpg -------------------------------------------------------------------------------- /examples/pose/68757076df6c98e9d6ba6ed00870daef.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/68757076df6c98e9d6ba6ed00870daef.jpg -------------------------------------------------------------------------------- /examples/pose/6f0029a9592a11530267b3a51413ae74.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/6f0029a9592a11530267b3a51413ae74.jpg -------------------------------------------------------------------------------- /examples/pose/7530ae51e811b1878fae23ea243a3a30.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/7530ae51e811b1878fae23ea243a3a30.jpg -------------------------------------------------------------------------------- /examples/pose/780047b55ee80b0dc2468ad16cab2278.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/780047b55ee80b0dc2468ad16cab2278.jpg -------------------------------------------------------------------------------- /examples/pose/ab2192beaefb58e872ce55099cbed8fe.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/ab2192beaefb58e872ce55099cbed8fe.jpg -------------------------------------------------------------------------------- /examples/pose/d5241b4de0cebd2722b05d855a5a9ca6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/d5241b4de0cebd2722b05d855a5a9ca6.jpg -------------------------------------------------------------------------------- /examples/pose/d6fcd37df9983973af08af3f9267cd1e.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/d6fcd37df9983973af08af3f9267cd1e.jpg -------------------------------------------------------------------------------- /examples/pose/d7e876bc5f9e8277d58e30bd83c7452f.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/examples/pose/d7e876bc5f9e8277d58e30bd83c7452f.jpg -------------------------------------------------------------------------------- /fetch_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | urle () { [[ "${1}" ]] || return 1; local LANG=C i x; for (( i = 0; i < ${#1}; i++ )); do x="${1:i:1}"; [[ "${x}" == [a-zA-Z0-9.~-] ]] && echo -n "${x}" || printf '%%%02X' "'${x}"; done; echo; } 3 | 4 | mkdir -p data/smpl_related/models 5 | 6 | # username and password input 7 | echo -e "\nYou need to register at https://icon.is.tue.mpg.de/, according to Installation Instruction." 8 | read -p "Username (ICON):" username 9 | read -p "Password (ICON):" password 10 | username=$(urle $username) 11 | password=$(urle $password) 12 | 13 | # SMPL (Male, Female) 14 | echo -e "\nDownloading SMPL..." 15 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smpl&sfile=SMPL_python_v.1.0.0.zip&resume=1' -O './data/smpl_related/models/SMPL_python_v.1.0.0.zip' --no-check-certificate --continue 16 | unzip data/smpl_related/models/SMPL_python_v.1.0.0.zip -d data/smpl_related/models 17 | mv data/smpl_related/models/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_FEMALE.pkl 18 | mv data/smpl_related/models/smpl/models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_MALE.pkl 19 | cd data/smpl_related/models 20 | rm -rf *.zip __MACOSX smpl/models smpl/smpl_webuser 21 | cd ../../.. 22 | 23 | # SMPL (Neutral, from SMPLIFY) 24 | echo -e "\nDownloading SMPLify..." 25 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smplify&sfile=mpips_smplify_public_v2.zip&resume=1' -O './data/smpl_related/models/mpips_smplify_public_v2.zip' --no-check-certificate --continue 26 | unzip data/smpl_related/models/mpips_smplify_public_v2.zip -d data/smpl_related/models 27 | mv data/smpl_related/models/smplify_public/code/models/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl data/smpl_related/models/smpl/SMPL_NEUTRAL.pkl 28 | cd data/smpl_related/models 29 | rm -rf *.zip smplify_public 30 | cd ../../.. 31 | 32 | # SMPL-X 33 | echo -e "\nDownloading SMPL-X..." 34 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=smplx&sfile=models_smplx_v1_1.zip&resume=1' -O './data/smpl_related/models/models_smplx_v1_1.zip' --no-check-certificate --continue 35 | unzip data/smpl_related/models/models_smplx_v1_1.zip -d data/smpl_related 36 | rm -f data/smpl_related/models/models_smplx_v1_1.zip 37 | 38 | # ECON 39 | echo -e "\nDownloading ECON..." 40 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=econ_data.zip&resume=1' -O './data/econ_data.zip' --no-check-certificate --continue 41 | cd data && unzip econ_data.zip 42 | mv smpl_data smpl_related/ 43 | rm -f econ_data.zip 44 | cd .. 45 | 46 | mkdir -p data/HPS 47 | 48 | # PIXIE 49 | echo -e "\nDownloading PIXIE..." 50 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=HPS/pixie_data.zip&resume=1' -O './data/HPS/pixie_data.zip' --no-check-certificate --continue 51 | cd data/HPS && unzip pixie_data.zip 52 | rm -f pixie_data.zip 53 | cd ../.. 54 | 55 | # PyMAF-X 56 | echo -e "\nDownloading PyMAF-X..." 57 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=HPS/pymafx_data.zip&resume=1' -O './data/HPS/pymafx_data.zip' --no-check-certificate --continue 58 | cd data/HPS && unzip pymafx_data.zip 59 | rm -f pymafx_data.zip 60 | cd ../.. 61 | -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/__init__.py -------------------------------------------------------------------------------- /lib/common/BNI.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import trimesh 3 | 4 | from lib.common.BNI_utils import ( 5 | depth_inverse_transform, 6 | double_side_bilateral_normal_integration, 7 | verts_inverse_transform, 8 | ) 9 | 10 | 11 | class BNI: 12 | def __init__(self, dir_path, name, BNI_dict, cfg, device): 13 | 14 | self.scale = 256.0 15 | self.cfg = cfg 16 | self.name = name 17 | 18 | self.normal_front = BNI_dict["normal_F"] 19 | self.normal_back = BNI_dict["normal_B"] 20 | self.mask = BNI_dict["mask"] 21 | 22 | self.depth_front = BNI_dict["depth_F"] 23 | self.depth_back = BNI_dict["depth_B"] 24 | self.depth_mask = BNI_dict["depth_mask"] 25 | 26 | # hparam: 27 | # k --> smaller, keep continuity 28 | # lambda --> larger, more depth-awareness 29 | 30 | self.k = self.cfg['k'] 31 | self.lambda1 = self.cfg['lambda1'] 32 | self.boundary_consist = self.cfg['boundary_consist'] 33 | self.cut_intersection = self.cfg['cut_intersection'] 34 | 35 | self.F_B_surface = None 36 | self.F_B_trimesh = None 37 | self.F_depth = None 38 | self.B_depth = None 39 | 40 | self.device = device 41 | self.export_dir = dir_path 42 | 43 | # code: https://github.com/hoshino042/bilateral_normal_integration 44 | # paper: Bilateral Normal Integration 45 | 46 | def extract_surface(self, verbose=True): 47 | 48 | bni_result = double_side_bilateral_normal_integration( 49 | normal_front=self.normal_front, 50 | normal_back=self.normal_back, 51 | normal_mask=self.mask, 52 | depth_front=self.depth_front * self.scale, 53 | depth_back=self.depth_back * self.scale, 54 | depth_mask=self.depth_mask, 55 | k=self.k, 56 | lambda_normal_back=1.0, 57 | lambda_depth_front=self.lambda1, 58 | lambda_depth_back=self.lambda1, 59 | lambda_boundary_consistency=self.boundary_consist, 60 | cut_intersection=self.cut_intersection, 61 | ) 62 | 63 | F_verts = verts_inverse_transform(bni_result["F_verts"], self.scale) 64 | B_verts = verts_inverse_transform(bni_result["B_verts"], self.scale) 65 | 66 | self.F_depth = depth_inverse_transform(bni_result["F_depth"], self.scale) 67 | self.B_depth = depth_inverse_transform(bni_result["B_depth"], self.scale) 68 | 69 | F_B_verts = torch.cat((F_verts, B_verts), dim=0) 70 | F_B_faces = torch.cat( 71 | (bni_result["F_faces"], bni_result["B_faces"] + bni_result["F_faces"].max() + 1), dim=0 72 | ) 73 | 74 | self.F_B_trimesh = trimesh.Trimesh( 75 | F_B_verts.float(), F_B_faces.long(), process=False, maintain_order=True 76 | ) 77 | 78 | # self.F_trimesh = trimesh.Trimesh( 79 | # F_verts.float(), bni_result["F_faces"].long(), process=False, maintain_order=True 80 | # ) 81 | 82 | # self.B_trimesh = trimesh.Trimesh( 83 | # B_verts.float(), bni_result["B_faces"].long(), process=False, maintain_order=True 84 | # ) 85 | 86 | 87 | if __name__ == "__main__": 88 | 89 | import os.path as osp 90 | 91 | import numpy as np 92 | from tqdm import tqdm 93 | 94 | root = "/home/yxiu/Code/ECON/results/examples/BNI" 95 | npy_file = f"{root}/304e9c4798a8c3967de7c74c24ef2e38.npy" 96 | bni_dict = np.load(npy_file, allow_pickle=True).item() 97 | 98 | default_cfg = {'k': 2, 'lambda1': 1e-4, 'boundary_consist': 1e-6} 99 | 100 | # for k in [1, 2, 4, 10, 100]: 101 | # default_cfg['k'] = k 102 | # for k in [1e-8, 1e-4, 1e-2, 1e-1, 1]: 103 | # default_cfg['lambda1'] = k 104 | # for k in [1e-4, 1e-2, 0]: 105 | # default_cfg['boundary_consist'] = k 106 | 107 | bni_object = BNI( 108 | osp.dirname(npy_file), osp.basename(npy_file), bni_dict, default_cfg, 109 | torch.device('cuda:0') 110 | ) 111 | 112 | bni_object.extract_surface() 113 | bni_object.F_trimesh.export(osp.join(osp.dirname(npy_file), "F.obj")) 114 | bni_object.B_trimesh.export(osp.join(osp.dirname(npy_file), "B.obj")) 115 | bni_object.F_B_trimesh.export(osp.join(osp.dirname(npy_file), "BNI.obj")) 116 | -------------------------------------------------------------------------------- /lib/common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/common/__init__.py -------------------------------------------------------------------------------- /lib/common/libmesh/setup.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | from setuptools import setup 3 | from Cython.Build import cythonize 4 | 5 | setup(name='libmesh', ext_modules=cythonize("*.pyx"), include_dirs=[numpy.get_include()]) 6 | -------------------------------------------------------------------------------- /lib/common/libmesh/triangle_hash.pyx: -------------------------------------------------------------------------------- 1 | 2 | # distutils: language=c++ 3 | import numpy as np 4 | 5 | cimport cython 6 | cimport numpy as np 7 | from libc.math cimport ceil, floor 8 | from libcpp.vector cimport vector 9 | 10 | 11 | cdef class TriangleHash: 12 | cdef vector[vector[int]] spatial_hash 13 | cdef int resolution 14 | 15 | def __cinit__(self, double[:, :, :] triangles, int resolution): 16 | self.spatial_hash.resize(resolution * resolution) 17 | self.resolution = resolution 18 | self._build_hash(triangles) 19 | 20 | @cython.boundscheck(False) # Deactivate bounds checking 21 | @cython.wraparound(False) # Deactivate negative indexing. 22 | cdef int _build_hash(self, double[:, :, :] triangles): 23 | assert(triangles.shape[1] == 3) 24 | assert(triangles.shape[2] == 2) 25 | 26 | cdef int n_tri = triangles.shape[0] 27 | cdef int bbox_min[2] 28 | cdef int bbox_max[2] 29 | 30 | cdef int i_tri, j, x, y 31 | cdef int spatial_idx 32 | 33 | for i_tri in range(n_tri): 34 | # Compute bounding box 35 | for j in range(2): 36 | bbox_min[j] = min( 37 | triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] 38 | ) 39 | bbox_max[j] = max( 40 | triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] 41 | ) 42 | bbox_min[j] = min(max(bbox_min[j], 0), self.resolution - 1) 43 | bbox_max[j] = min(max(bbox_max[j], 0), self.resolution - 1) 44 | 45 | # Find all voxels where bounding box intersects 46 | for x in range(bbox_min[0], bbox_max[0] + 1): 47 | for y in range(bbox_min[1], bbox_max[1] + 1): 48 | spatial_idx = self.resolution * x + y 49 | self.spatial_hash[spatial_idx].push_back(i_tri) 50 | 51 | @cython.boundscheck(False) # Deactivate bounds checking 52 | @cython.wraparound(False) # Deactivate negative indexing. 53 | cpdef query(self, double[:, :] points): 54 | assert(points.shape[1] == 2) 55 | cdef int n_points = points.shape[0] 56 | 57 | cdef vector[int] points_indices 58 | cdef vector[int] tri_indices 59 | # cdef int[:] points_indices_np 60 | # cdef int[:] tri_indices_np 61 | 62 | cdef int i_point, k, x, y 63 | cdef int spatial_idx 64 | 65 | for i_point in range(n_points): 66 | x = int(points[i_point, 0]) 67 | y = int(points[i_point, 1]) 68 | if not (0 <= x < self.resolution and 0 <= y < self.resolution): 69 | continue 70 | 71 | spatial_idx = self.resolution * x + y 72 | for i_tri in self.spatial_hash[spatial_idx]: 73 | points_indices.push_back(i_point) 74 | tri_indices.push_back(i_tri) 75 | 76 | points_indices_np = np.zeros(points_indices.size(), dtype=np.int32) 77 | tri_indices_np = np.zeros(tri_indices.size(), dtype=np.int32) 78 | 79 | cdef int[:] points_indices_view = points_indices_np 80 | cdef int[:] tri_indices_view = tri_indices_np 81 | 82 | for k in range(points_indices.size()): 83 | points_indices_view[k] = points_indices[k] 84 | 85 | for k in range(tri_indices.size()): 86 | tri_indices_view[k] = tri_indices[k] 87 | 88 | return points_indices_np, tri_indices_np 89 | -------------------------------------------------------------------------------- /lib/common/libvoxelize/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from Cython.Build import cythonize 3 | 4 | setup(name='libvoxelize', ext_modules=cythonize("*.pyx")) 5 | -------------------------------------------------------------------------------- /lib/common/libvoxelize/voxelize.pyx: -------------------------------------------------------------------------------- 1 | cimport cython 2 | from cython.view cimport array as cvarray 3 | from libc.math cimport ceil, floor 4 | 5 | 6 | cdef extern from "tribox2.h": 7 | int triBoxOverlap(float boxcenter[3], float boxhalfsize[3], 8 | float tri0[3], float tri1[3], float tri2[3]) 9 | 10 | 11 | @cython.boundscheck(False) # Deactivate bounds checking 12 | @cython.wraparound(False) # Deactivate negative indexing. 13 | cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): 14 | assert(faces.shape[1] == 3) 15 | assert(faces.shape[2] == 3) 16 | 17 | n_faces = faces.shape[0] 18 | cdef int i 19 | for i in range(n_faces): 20 | voxelize_triangle_(occ, faces[i]) 21 | 22 | 23 | @cython.boundscheck(False) # Deactivate bounds checking 24 | @cython.wraparound(False) # Deactivate negative indexing. 25 | cpdef int voxelize_triangle_(bint[:, :, :] occupancies, float[:, ::1] triverts): 26 | cdef int bbox_min[3] 27 | cdef int bbox_max[3] 28 | cdef int i, j, k 29 | cdef float boxhalfsize[3] 30 | cdef float boxcenter[3] 31 | cdef bint intersection 32 | 33 | boxhalfsize[:] = (0.5, 0.5, 0.5) 34 | 35 | for i in range(3): 36 | bbox_min[i] = ( 37 | min(triverts[0, i], triverts[1, i], triverts[2, i]) 38 | ) 39 | bbox_min[i] = min(max(bbox_min[i], 0), occupancies.shape[i] - 1) 40 | 41 | for i in range(3): 42 | bbox_max[i] = ( 43 | max(triverts[0, i], triverts[1, i], triverts[2, i]) 44 | ) 45 | bbox_max[i] = min(max(bbox_max[i], 0), occupancies.shape[i] - 1) 46 | 47 | for i in range(bbox_min[0], bbox_max[0] + 1): 48 | for j in range(bbox_min[1], bbox_max[1] + 1): 49 | for k in range(bbox_min[2], bbox_max[2] + 1): 50 | boxcenter[:] = (i + 0.5, j + 0.5, k + 0.5) 51 | intersection = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], 52 | &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) 53 | occupancies[i, j, k] |= intersection 54 | 55 | 56 | @cython.boundscheck(False) # Deactivate bounds checking 57 | @cython.wraparound(False) # Deactivate negative indexing. 58 | cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): 59 | assert(boxcenter.shape[0] == 3) 60 | assert(boxhalfsize.shape[0] == 3) 61 | assert(triverts.shape[0] == triverts.shape[1] == 3) 62 | 63 | # print(triverts) 64 | # Call functions 65 | cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], 66 | &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) 67 | return result 68 | -------------------------------------------------------------------------------- /lib/common/local_affine.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 by Haozhe Wu, Tsinghua University, Department of Computer Science and Technology. 2 | # All rights reserved. 3 | # This file is part of the pytorch-nicp, 4 | # and is released under the "MIT License Agreement". Please see the LICENSE 5 | # file that should have been included as part of this package. 6 | 7 | import torch 8 | import torch.nn as nn 9 | import trimesh 10 | from pytorch3d.loss import chamfer_distance 11 | from pytorch3d.structures import Meshes 12 | from tqdm import tqdm 13 | 14 | from lib.common.train_util import init_loss 15 | from lib.dataset.mesh_util import update_mesh_shape_prior_losses 16 | 17 | 18 | # reference: https://github.com/wuhaozhe/pytorch-nicp 19 | class LocalAffine(nn.Module): 20 | def __init__(self, num_points, batch_size=1, edges=None): 21 | ''' 22 | specify the number of points, the number of points should be constant across the batch 23 | and the edges torch.Longtensor() with shape N * 2 24 | the local affine operator supports batch operation 25 | batch size must be constant 26 | add additional pooling on top of w matrix 27 | ''' 28 | super(LocalAffine, self).__init__() 29 | self.A = nn.Parameter( 30 | torch.eye(3).unsqueeze(0).unsqueeze(0).repeat(batch_size, num_points, 1, 1) 31 | ) 32 | self.b = nn.Parameter( 33 | torch.zeros(3).unsqueeze(0).unsqueeze(0).unsqueeze(3).repeat( 34 | batch_size, num_points, 1, 1 35 | ) 36 | ) 37 | self.edges = edges 38 | self.num_points = num_points 39 | 40 | def stiffness(self): 41 | ''' 42 | calculate the stiffness of local affine transformation 43 | f norm get infinity gradient when w is zero matrix, 44 | ''' 45 | if self.edges is None: 46 | raise Exception("edges cannot be none when calculate stiff") 47 | affine_weight = torch.cat((self.A, self.b), dim=3) 48 | w1 = torch.index_select(affine_weight, dim=1, index=self.edges[:, 0]) 49 | w2 = torch.index_select(affine_weight, dim=1, index=self.edges[:, 1]) 50 | w_diff = (w1 - w2)**2 51 | w_rigid = (torch.linalg.det(self.A) - 1.0)**2 52 | return w_diff, w_rigid 53 | 54 | def forward(self, x): 55 | ''' 56 | x should have shape of B * N * 3 * 1 57 | ''' 58 | x = x.unsqueeze(3) 59 | out_x = torch.matmul(self.A, x) 60 | out_x = out_x + self.b 61 | out_x.squeeze_(3) 62 | stiffness, rigid = self.stiffness() 63 | 64 | return out_x, stiffness, rigid 65 | 66 | 67 | def trimesh2meshes(mesh): 68 | ''' 69 | convert trimesh mesh to pytorch3d mesh 70 | ''' 71 | verts = torch.from_numpy(mesh.vertices).float() 72 | faces = torch.from_numpy(mesh.faces).long() 73 | mesh = Meshes(verts.unsqueeze(0), faces.unsqueeze(0)) 74 | return mesh 75 | 76 | 77 | def register(target_mesh, src_mesh, device, verbose=True): 78 | 79 | # define local_affine deform verts 80 | tgt_mesh = trimesh2meshes(target_mesh).to(device) 81 | src_verts = src_mesh.verts_padded().clone() 82 | 83 | local_affine_model = LocalAffine( 84 | src_mesh.verts_padded().shape[1], 85 | src_mesh.verts_padded().shape[0], src_mesh.edges_packed() 86 | ).to(device) 87 | 88 | optimizer_cloth = torch.optim.Adam([{'params': local_affine_model.parameters()}], 89 | lr=1e-2, 90 | amsgrad=True) 91 | scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau( 92 | optimizer_cloth, 93 | mode="min", 94 | factor=0.1, 95 | verbose=0, 96 | min_lr=1e-5, 97 | patience=5, 98 | ) 99 | 100 | losses = init_loss() 101 | 102 | if verbose: 103 | loop_cloth = tqdm(range(100)) 104 | else: 105 | loop_cloth = range(100) 106 | 107 | for i in loop_cloth: 108 | 109 | optimizer_cloth.zero_grad() 110 | 111 | deformed_verts, stiffness, rigid = local_affine_model(x=src_verts) 112 | src_mesh = src_mesh.update_padded(deformed_verts) 113 | 114 | # losses for laplacian, edge, normal consistency 115 | update_mesh_shape_prior_losses(src_mesh, losses) 116 | 117 | losses["cloth"]["value"] = chamfer_distance( 118 | x=src_mesh.verts_padded(), y=tgt_mesh.verts_padded() 119 | )[0] 120 | losses["stiff"]["value"] = torch.mean(stiffness) 121 | losses["rigid"]["value"] = torch.mean(rigid) 122 | 123 | # Weighted sum of the losses 124 | cloth_loss = torch.tensor(0.0, requires_grad=True).to(device) 125 | pbar_desc = "Register SMPL-X -> d-BiNI -- " 126 | 127 | for k in losses.keys(): 128 | if losses[k]["weight"] > 0.0 and losses[k]["value"] != 0.0: 129 | cloth_loss = cloth_loss + \ 130 | losses[k]["value"] * losses[k]["weight"] 131 | pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.3f} | " 132 | 133 | if verbose: 134 | pbar_desc += f"TOTAL: {cloth_loss:.3f}" 135 | loop_cloth.set_description(pbar_desc) 136 | 137 | # update params 138 | cloth_loss.backward(retain_graph=True) 139 | optimizer_cloth.step() 140 | scheduler_cloth.step(cloth_loss) 141 | 142 | final = trimesh.Trimesh( 143 | src_mesh.verts_packed().detach().squeeze(0).cpu(), 144 | src_mesh.faces_packed().detach().squeeze(0).cpu(), 145 | process=False, 146 | maintains_order=True 147 | ) 148 | 149 | return final 150 | -------------------------------------------------------------------------------- /lib/common/train_util.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import pytorch_lightning as pl 18 | import torch 19 | from termcolor import colored 20 | 21 | from ..dataset.mesh_util import * 22 | from ..net.geometry import orthogonal 23 | 24 | 25 | class Format: 26 | end = '\033[0m' 27 | start = '\033[4m' 28 | 29 | 30 | def init_loss(): 31 | 32 | losses = { 33 | # Cloth: chamfer distance 34 | "cloth": {"weight": 1e3, "value": 0.0}, 35 | # Stiffness: [RT]_v1 - [RT]_v2 (v1-edge-v2) 36 | "stiff": {"weight": 1e5, "value": 0.0}, 37 | # Cloth: det(R) = 1 38 | "rigid": {"weight": 1e5, "value": 0.0}, 39 | # Cloth: edge length 40 | "edge": {"weight": 0, "value": 0.0}, 41 | # Cloth: normal consistency 42 | "nc": {"weight": 0, "value": 0.0}, 43 | # Cloth: laplacian smoonth 44 | "lapla": {"weight": 1e2, "value": 0.0}, 45 | # Body: Normal_pred - Normal_smpl 46 | "normal": {"weight": 1e0, "value": 0.0}, 47 | # Body: Silhouette_pred - Silhouette_smpl 48 | "silhouette": {"weight": 1e0, "value": 0.0}, 49 | # Joint: reprojected joints difference 50 | "joint": {"weight": 1e0, "value": 0.0}, 51 | } 52 | 53 | return losses 54 | 55 | 56 | class SubTrainer(pl.Trainer): 57 | def save_checkpoint(self, filepath, weights_only=False): 58 | """Save model/training states as a checkpoint file through state-dump and file-write. 59 | Args: 60 | filepath: write-target file's path 61 | weights_only: saving model weights only 62 | """ 63 | _checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only) 64 | 65 | del_keys = [] 66 | for key in _checkpoint["state_dict"].keys(): 67 | for ignore_key in ["normal_filter", "voxelization", "reconEngine"]: 68 | if ignore_key in key: 69 | del_keys.append(key) 70 | for key in del_keys: 71 | del _checkpoint["state_dict"][key] 72 | 73 | pl.utilities.cloud_io.atomic_save(_checkpoint, filepath) 74 | 75 | 76 | def query_func(opt, netG, features, points, proj_matrix=None): 77 | """ 78 | - points: size of (bz, N, 3) 79 | - proj_matrix: size of (bz, 4, 4) 80 | return: size of (bz, 1, N) 81 | """ 82 | assert len(points) == 1 83 | samples = points.repeat(opt.num_views, 1, 1) 84 | samples = samples.permute(0, 2, 1) # [bz, 3, N] 85 | 86 | # view specific query 87 | if proj_matrix is not None: 88 | samples = orthogonal(samples, proj_matrix) 89 | 90 | calib_tensor = torch.stack([torch.eye(4).float()], dim=0).type_as(samples) 91 | 92 | preds = netG.query( 93 | features=features, 94 | points=samples, 95 | calibs=calib_tensor, 96 | regressor=netG.if_regressor, 97 | ) 98 | 99 | if type(preds) is list: 100 | preds = preds[0] 101 | 102 | return preds 103 | 104 | 105 | def query_func_IF(batch, netG, points): 106 | """ 107 | - points: size of (bz, N, 3) 108 | return: size of (bz, 1, N) 109 | """ 110 | 111 | batch["samples_geo"] = points 112 | batch["calib"] = torch.stack([torch.eye(4).float()], dim=0).type_as(points) 113 | 114 | preds = netG(batch) 115 | 116 | return preds.unsqueeze(1) 117 | 118 | 119 | def batch_mean(res, key): 120 | return torch.stack([ 121 | x[key] if torch.is_tensor(x[key]) else torch.as_tensor(x[key]) for x in res 122 | ]).mean() 123 | 124 | 125 | def accumulate(outputs, rot_num, split): 126 | 127 | hparam_log_dict = {} 128 | 129 | metrics = outputs[0].keys() 130 | datasets = split.keys() 131 | 132 | for dataset in datasets: 133 | for metric in metrics: 134 | keyword = f"{dataset}/{metric}" 135 | if keyword not in hparam_log_dict.keys(): 136 | hparam_log_dict[keyword] = 0 137 | for idx in range(split[dataset][0] * rot_num, split[dataset][1] * rot_num): 138 | hparam_log_dict[keyword] += outputs[idx][metric].item() 139 | hparam_log_dict[keyword] /= (split[dataset][1] - split[dataset][0]) * rot_num 140 | 141 | print(colored(hparam_log_dict, "green")) 142 | 143 | return hparam_log_dict 144 | -------------------------------------------------------------------------------- /lib/dataset/NormalModule.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | # pytorch lightning related libs 18 | import pytorch_lightning as pl 19 | from torch.utils.data import DataLoader 20 | 21 | from lib.dataset.NormalDataset import NormalDataset 22 | 23 | 24 | class NormalModule(pl.LightningDataModule): 25 | def __init__(self, cfg): 26 | super(NormalModule, self).__init__() 27 | self.cfg = cfg 28 | 29 | self.batch_size = self.cfg.batch_size 30 | 31 | self.data_size = {} 32 | 33 | def prepare_data(self): 34 | 35 | pass 36 | 37 | def setup(self, stage): 38 | 39 | self.train_dataset = NormalDataset(cfg=self.cfg, split="train") 40 | self.val_dataset = NormalDataset(cfg=self.cfg, split="val") 41 | self.test_dataset = NormalDataset(cfg=self.cfg, split="test") 42 | 43 | self.data_size = { 44 | "train": len(self.train_dataset), 45 | "val": len(self.val_dataset), 46 | } 47 | 48 | def train_dataloader(self): 49 | 50 | train_data_loader = DataLoader( 51 | self.train_dataset, 52 | batch_size=self.batch_size, 53 | shuffle=True, 54 | num_workers=self.cfg.num_threads, 55 | pin_memory=True, 56 | ) 57 | 58 | return train_data_loader 59 | 60 | def val_dataloader(self): 61 | 62 | val_data_loader = DataLoader( 63 | self.val_dataset, 64 | batch_size=self.batch_size, 65 | shuffle=False, 66 | num_workers=self.cfg.num_threads, 67 | pin_memory=True, 68 | ) 69 | 70 | return val_data_loader 71 | 72 | def val_dataloader(self): 73 | 74 | test_data_loader = DataLoader( 75 | self.test_dataset, 76 | batch_size=1, 77 | shuffle=False, 78 | num_workers=self.cfg.num_threads, 79 | pin_memory=True, 80 | ) 81 | 82 | return test_data_loader 83 | -------------------------------------------------------------------------------- /lib/dataset/PointFeat.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from pytorch3d.structures import Meshes, Pointclouds 3 | 4 | from lib.common.render_utils import face_vertices 5 | from lib.dataset.Evaluator import point_mesh_distance 6 | from lib.dataset.mesh_util import SMPLX, barycentric_coordinates_of_projection 7 | 8 | 9 | class PointFeat: 10 | def __init__(self, verts, faces): 11 | 12 | # verts [B, N_vert, 3] 13 | # faces [B, N_face, 3] 14 | # triangles [B, N_face, 3, 3] 15 | 16 | self.Bsize = verts.shape[0] 17 | self.device = verts.device 18 | self.faces = faces 19 | 20 | # SMPL has watertight mesh, but SMPL-X has two eyeballs and open mouth 21 | # 1. remove eye_ball faces from SMPL-X: 9928-9383, 10474-9929 22 | # 2. fill mouth holes with 30 more faces 23 | 24 | if verts.shape[1] == 10475: 25 | faces = faces[:, ~SMPLX().smplx_eyeball_fid_mask] 26 | mouth_faces = ( 27 | torch.as_tensor(SMPLX().smplx_mouth_fid).unsqueeze(0).repeat(self.Bsize, 1, 28 | 1).to(self.device) 29 | ) 30 | self.faces = torch.cat([faces, mouth_faces], dim=1).long() 31 | 32 | self.verts = verts.float() 33 | self.triangles = face_vertices(self.verts, self.faces) 34 | self.mesh = Meshes(self.verts, self.faces).to(self.device) 35 | 36 | def query(self, points): 37 | 38 | points = points.float() 39 | residues, pts_ind = point_mesh_distance(self.mesh, Pointclouds(points), weighted=False) 40 | 41 | closest_triangles = torch.gather( 42 | self.triangles, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3) 43 | ).view(-1, 3, 3) 44 | bary_weights = barycentric_coordinates_of_projection(points.view(-1, 3), closest_triangles) 45 | 46 | feat_normals = face_vertices(self.mesh.verts_normals_padded(), self.faces) 47 | closest_normals = torch.gather( 48 | feat_normals, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3) 49 | ).view(-1, 3, 3) 50 | shoot_verts = ((closest_triangles * bary_weights[:, :, None]).sum(1).unsqueeze(0)) 51 | 52 | pts2shoot_normals = points - shoot_verts 53 | pts2shoot_normals = pts2shoot_normals / torch.norm(pts2shoot_normals, dim=-1, keepdim=True) 54 | 55 | shoot_normals = ((closest_normals * bary_weights[:, :, None]).sum(1).unsqueeze(0)) 56 | shoot_normals = shoot_normals / torch.norm(shoot_normals, dim=-1, keepdim=True) 57 | angles = (pts2shoot_normals * shoot_normals).sum(dim=-1).abs() 58 | 59 | return (torch.sqrt(residues).unsqueeze(0), angles) 60 | -------------------------------------------------------------------------------- /lib/dataset/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/dataset/__init__.py -------------------------------------------------------------------------------- /lib/dataset/tbfo.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/dataset/tbfo.ttf -------------------------------------------------------------------------------- /lib/net/BasePIFuNet.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import pytorch_lightning as pl 18 | import torch.nn as nn 19 | 20 | from .geometry import index, orthogonal, perspective 21 | 22 | 23 | class BasePIFuNet(pl.LightningModule): 24 | def __init__( 25 | self, 26 | projection_mode="orthogonal", 27 | error_term=nn.MSELoss(), 28 | ): 29 | """ 30 | :param projection_mode: 31 | Either orthogonal or perspective. 32 | It will call the corresponding function for projection. 33 | :param error_term: 34 | nn Loss between the predicted [B, Res, N] and the label [B, Res, N] 35 | """ 36 | super(BasePIFuNet, self).__init__() 37 | self.name = "base" 38 | 39 | self.error_term = error_term 40 | 41 | self.index = index 42 | self.projection = orthogonal if projection_mode == "orthogonal" else perspective 43 | 44 | def forward(self, points, images, calibs, transforms=None): 45 | """ 46 | :param points: [B, 3, N] world space coordinates of points 47 | :param images: [B, C, H, W] input images 48 | :param calibs: [B, 3, 4] calibration matrices for each image 49 | :param transforms: Optional [B, 2, 3] image space coordinate transforms 50 | :return: [B, Res, N] predictions for each point 51 | """ 52 | features = self.filter(images) 53 | preds = self.query(features, points, calibs, transforms) 54 | return preds 55 | 56 | def filter(self, images): 57 | """ 58 | Filter the input images 59 | store all intermediate features. 60 | :param images: [B, C, H, W] input images 61 | """ 62 | return None 63 | 64 | def query(self, features, points, calibs, transforms=None): 65 | """ 66 | Given 3D points, query the network predictions for each point. 67 | Image features should be pre-computed before this call. 68 | store all intermediate features. 69 | query() function may behave differently during training/testing. 70 | :param points: [B, 3, N] world space coordinates of points 71 | :param calibs: [B, 3, 4] calibration matrices for each image 72 | :param transforms: Optional [B, 2, 3] image space coordinate transforms 73 | :param labels: Optional [B, Res, N] gt labeling 74 | :return: [B, Res, N] predictions for each point 75 | """ 76 | return None 77 | 78 | def get_error(self, preds, labels): 79 | """ 80 | Get the network loss from the last query 81 | :return: loss term 82 | """ 83 | return self.error_term(preds, labels) 84 | -------------------------------------------------------------------------------- /lib/net/GANLoss.py: -------------------------------------------------------------------------------- 1 | """ The code is based on https://github.com/apple/ml-gsn/ with adaption. """ 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from torch import autograd 7 | 8 | from lib.net.Discriminator import StyleDiscriminator 9 | 10 | 11 | def hinge_loss(fake_pred, real_pred, mode): 12 | if mode == 'd': 13 | # Discriminator update 14 | d_loss_fake = torch.mean(F.relu(1.0 + fake_pred)) 15 | d_loss_real = torch.mean(F.relu(1.0 - real_pred)) 16 | d_loss = d_loss_fake + d_loss_real 17 | elif mode == 'g': 18 | # Generator update 19 | d_loss = -torch.mean(fake_pred) 20 | return d_loss 21 | 22 | 23 | def logistic_loss(fake_pred, real_pred, mode): 24 | if mode == 'd': 25 | # Discriminator update 26 | d_loss_fake = torch.mean(F.softplus(fake_pred)) 27 | d_loss_real = torch.mean(F.softplus(-real_pred)) 28 | d_loss = d_loss_fake + d_loss_real 29 | elif mode == 'g': 30 | # Generator update 31 | d_loss = torch.mean(F.softplus(-fake_pred)) 32 | return d_loss 33 | 34 | 35 | def r1_loss(real_pred, real_img): 36 | (grad_real, ) = autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True) 37 | grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean() 38 | return grad_penalty 39 | 40 | 41 | class GANLoss(nn.Module): 42 | def __init__( 43 | self, 44 | opt, 45 | disc_loss='logistic', 46 | ): 47 | super().__init__() 48 | self.opt = opt.gan 49 | 50 | input_dim = 3 51 | self.discriminator = StyleDiscriminator(input_dim, self.opt.img_res) 52 | 53 | if disc_loss == 'hinge': 54 | self.disc_loss = hinge_loss 55 | elif disc_loss == 'logistic': 56 | self.disc_loss = logistic_loss 57 | 58 | def forward(self, input): 59 | 60 | disc_in_real = input['norm_real'] 61 | disc_in_fake = input['norm_fake'] 62 | 63 | logits_real = self.discriminator(disc_in_real) 64 | logits_fake = self.discriminator(disc_in_fake) 65 | 66 | disc_loss = self.disc_loss(fake_pred=logits_fake, real_pred=logits_real, mode='d') 67 | 68 | log = { 69 | "disc_loss": disc_loss.detach(), 70 | "logits_real": logits_real.mean().detach(), 71 | "logits_fake": logits_fake.mean().detach(), 72 | } 73 | 74 | return disc_loss * self.opt.lambda_gan, log 75 | -------------------------------------------------------------------------------- /lib/net/__init__.py: -------------------------------------------------------------------------------- 1 | from .NormalNet import NormalNet 2 | -------------------------------------------------------------------------------- /lib/pixielib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/pixielib/__init__.py -------------------------------------------------------------------------------- /lib/pixielib/models/FLAME.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 | # Using this computer program means that you agree to the terms 6 | # in the LICENSE file included with this software distribution. 7 | # Any use not explicitly granted by the LICENSE is prohibited. 8 | # 9 | # Copyright©2019 Max-Planck-Gesellschaft zur Förderung 10 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 11 | # for Intelligent Systems. All rights reserved. 12 | # 13 | # For comments or questions, please email us at pixie@tue.mpg.de 14 | # For commercial licensing contact, please contact ps-license@tuebingen.mpg.de 15 | 16 | import pickle 17 | 18 | import numpy as np 19 | import torch 20 | import torch.nn as nn 21 | import torch.nn.functional as F 22 | 23 | 24 | class FLAMETex(nn.Module): 25 | """ 26 | FLAME texture: 27 | https://github.com/TimoBolkart/TF_FLAME/blob/ade0ab152300ec5f0e8555d6765411555c5ed43d/sample_texture.py#L64 28 | FLAME texture converted from BFM: 29 | https://github.com/TimoBolkart/BFM_to_FLAME 30 | """ 31 | def __init__(self, config): 32 | super(FLAMETex, self).__init__() 33 | if config.tex_type == "BFM": 34 | mu_key = "MU" 35 | pc_key = "PC" 36 | n_pc = 199 37 | tex_path = config.tex_path 38 | tex_space = np.load(tex_path) 39 | texture_mean = tex_space[mu_key].reshape(1, -1) 40 | texture_basis = tex_space[pc_key].reshape(-1, n_pc) 41 | 42 | elif config.tex_type == "FLAME": 43 | mu_key = "mean" 44 | pc_key = "tex_dir" 45 | n_pc = 200 46 | tex_path = config.flame_tex_path 47 | tex_space = np.load(tex_path) 48 | texture_mean = tex_space[mu_key].reshape(1, -1) / 255.0 49 | texture_basis = tex_space[pc_key].reshape(-1, n_pc) / 255.0 50 | else: 51 | print("texture type ", config.tex_type, "not exist!") 52 | raise NotImplementedError 53 | 54 | n_tex = config.n_tex 55 | num_components = texture_basis.shape[1] 56 | texture_mean = torch.from_numpy(texture_mean).float()[None, ...] 57 | texture_basis = torch.from_numpy(texture_basis[:, :n_tex]).float()[None, ...] 58 | self.register_buffer("texture_mean", texture_mean) 59 | self.register_buffer("texture_basis", texture_basis) 60 | 61 | def forward(self, texcode=None): 62 | """ 63 | texcode: [batchsize, n_tex] 64 | texture: [bz, 3, 256, 256], range: 0-1 65 | """ 66 | texture = self.texture_mean + (self.texture_basis * texcode[:, None, :]).sum(-1) 67 | texture = texture.reshape(texcode.shape[0], 512, 512, 3).permute(0, 3, 1, 2) 68 | texture = F.interpolate(texture, [256, 256]) 69 | texture = texture[:, [2, 1, 0], :, :] 70 | return texture 71 | 72 | 73 | def texture_flame2smplx(cached_data, flame_texture, smplx_texture): 74 | """Convert flame texture map (face-only) into smplx texture map (includes body texture) 75 | TODO: pytorch version ==> grid sample 76 | """ 77 | if smplx_texture.shape[0] != smplx_texture.shape[1]: 78 | print("SMPL-X texture not squared (%d != %d)" % (smplx_texture[0], smplx_texture[1])) 79 | return 80 | if smplx_texture.shape[0] != cached_data["target_resolution"]: 81 | print( 82 | "SMPL-X texture size does not match cached image resolution (%d != %d)" % 83 | (smplx_texture.shape[0], cached_data["target_resolution"]) 84 | ) 85 | return 86 | x_coords = cached_data["x_coords"] 87 | y_coords = cached_data["y_coords"] 88 | target_pixel_ids = cached_data["target_pixel_ids"] 89 | source_uv_points = cached_data["source_uv_points"] 90 | 91 | source_tex_coords = np.zeros_like((source_uv_points)).astype(int) 92 | source_tex_coords[:, 0] = np.clip( 93 | flame_texture.shape[0] * (1.0 - source_uv_points[:, 1]), 94 | 0.0, 95 | flame_texture.shape[0], 96 | ).astype(int) 97 | source_tex_coords[:, 1] = np.clip( 98 | flame_texture.shape[1] * (source_uv_points[:, 0]), 0.0, flame_texture.shape[1] 99 | ).astype(int) 100 | 101 | smplx_texture[y_coords[target_pixel_ids].astype(int), 102 | x_coords[target_pixel_ids].astype(int), :, ] = flame_texture[source_tex_coords[:, 103 | 0], 104 | source_tex_coords[:, 105 | 1]] 106 | 107 | return smplx_texture 108 | -------------------------------------------------------------------------------- /lib/pixielib/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/pixielib/models/__init__.py -------------------------------------------------------------------------------- /lib/pixielib/models/encoders.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class ResnetEncoder(nn.Module): 8 | def __init__(self, append_layers=None): 9 | super(ResnetEncoder, self).__init__() 10 | from . import resnet 11 | 12 | # feature_size = 2048 13 | self.feature_dim = 2048 14 | self.encoder = resnet.load_ResNet50Model() # out: 2048 15 | # regressor 16 | self.append_layers = append_layers 17 | 18 | def forward(self, inputs): 19 | """inputs: [bz, 3, h, w], range: [0,1]""" 20 | features = self.encoder(inputs) 21 | if self.append_layers: 22 | features = self.last_op(features) 23 | return features 24 | 25 | 26 | class MLP(nn.Module): 27 | def __init__(self, channels=[2048, 1024, 1], last_op=None): 28 | super(MLP, self).__init__() 29 | layers = [] 30 | 31 | for l in range(0, len(channels) - 1): 32 | layers.append(nn.Linear(channels[l], channels[l + 1])) 33 | if l < len(channels) - 2: 34 | layers.append(nn.ReLU()) 35 | if last_op: 36 | layers.append(last_op) 37 | 38 | self.layers = nn.Sequential(*layers) 39 | 40 | def forward(self, inputs): 41 | outs = self.layers(inputs) 42 | return outs 43 | 44 | 45 | class HRNEncoder(nn.Module): 46 | def __init__(self, append_layers=None): 47 | super(HRNEncoder, self).__init__() 48 | from . import hrnet 49 | 50 | self.feature_dim = 2048 51 | self.encoder = hrnet.load_HRNet(pretrained=True) # out: 2048 52 | # regressor 53 | self.append_layers = append_layers 54 | 55 | def forward(self, inputs): 56 | """inputs: [bz, 3, h, w], range: [-1,1]""" 57 | features = self.encoder(inputs)["concat"] 58 | if self.append_layers: 59 | features = self.last_op(features) 60 | return features 61 | -------------------------------------------------------------------------------- /lib/pixielib/models/moderators.py: -------------------------------------------------------------------------------- 1 | """ Moderator 2 | # Input feature: body, part(head, hand) 3 | # output: fused feature, weight 4 | """ 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | # MLP + temperature softmax 11 | # w = SoftMax(w^\prime * temperature) 12 | 13 | 14 | class TempSoftmaxFusion(nn.Module): 15 | def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False): 16 | super(TempSoftmaxFusion, self).__init__() 17 | self.detach_inputs = detach_inputs 18 | self.detach_feature = detach_feature 19 | # weight 20 | layers = [] 21 | for l in range(0, len(channels) - 1): 22 | layers.append(nn.Linear(channels[l], channels[l + 1])) 23 | if l < len(channels) - 2: 24 | layers.append(nn.ReLU()) 25 | self.layers = nn.Sequential(*layers) 26 | # temperature 27 | self.register_parameter("temperature", nn.Parameter(torch.ones(1))) 28 | 29 | def forward(self, x, y, work=True): 30 | """ 31 | x: feature from body 32 | y: feature from part(head/hand) 33 | work: whether to fuse features 34 | """ 35 | if work: 36 | # 1. cat input feature, predict the weights 37 | f_in = torch.cat([x, y], dim=1) 38 | if self.detach_inputs: 39 | f_in = f_in.detach() 40 | f_temp = self.layers(f_in) 41 | f_weight = F.softmax(f_temp * self.temperature, dim=1) 42 | 43 | # 2. feature fusion 44 | if self.detach_feature: 45 | x = x.detach() 46 | y = y.detach() 47 | f_out = f_weight[:, [0]] * x + f_weight[:, [1]] * y 48 | x_out = f_out 49 | y_out = f_out 50 | else: 51 | x_out = x 52 | y_out = y 53 | f_weight = None 54 | return x_out, y_out, f_weight 55 | 56 | 57 | # MLP + Gumbel-Softmax trick 58 | # w = w^{\prime} - w^{\prime}\text{.detach()} + w^{\prime}\text{.gt(0.5)} 59 | 60 | 61 | class GumbelSoftmaxFusion(nn.Module): 62 | def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False): 63 | super(GumbelSoftmaxFusion, self).__init__() 64 | self.detach_inputs = detach_inputs 65 | self.detach_feature = detach_feature 66 | 67 | # weight 68 | layers = [] 69 | for l in range(0, len(channels) - 1): 70 | layers.append(nn.Linear(channels[l], channels[l + 1])) 71 | if l < len(channels) - 2: 72 | layers.append(nn.ReLU()) 73 | layers.append(nn.Softmax()) 74 | self.layers = nn.Sequential(*layers) 75 | 76 | def forward(self, x, y, work=True): 77 | """ 78 | x: feature from body 79 | y: feature from part(head/hand) 80 | work: whether to fuse features 81 | """ 82 | if work: 83 | # 1. cat input feature, predict the weights 84 | f_in = torch.cat([x, y], dim=-1) 85 | if self.detach_inputs: 86 | f_in = f_in.detach() 87 | f_weight = self.layers(f_in) 88 | # weight to be hard 89 | f_weight = f_weight - f_weight.detach() + f_weight.gt(0.5) 90 | 91 | # 2. feature fusion 92 | if self.detach_feature: 93 | x = x.detach() 94 | y = y.detach() 95 | f_out = f_weight[:, [0]] * x + f_weight[:, [1]] * y 96 | x_out = f_out 97 | y_out = f_out 98 | else: 99 | x_out = x 100 | y_out = y 101 | f_weight = None 102 | return x_out, y_out, f_weight 103 | -------------------------------------------------------------------------------- /lib/pixielib/utils/array_cropper.py: -------------------------------------------------------------------------------- 1 | """ 2 | crop 3 | for numpy array 4 | Given image, bbox(center, bboxsize) 5 | return: cropped image, tform(used for transform the keypoint accordingly) 6 | 7 | only support crop to squared images 8 | """ 9 | 10 | import numpy as np 11 | from skimage.transform import estimate_transform, rescale, resize, warp 12 | 13 | 14 | def points2bbox(points, points_scale=None): 15 | # recover range 16 | if points_scale: 17 | points[:, 0] = points[:, 0] * points_scale[1] / 2 + points_scale[1] / 2 18 | points[:, 1] = points[:, 1] * points_scale[0] / 2 + points_scale[0] / 2 19 | 20 | left = np.min(points[:, 0]) 21 | right = np.max(points[:, 0]) 22 | top = np.min(points[:, 1]) 23 | bottom = np.max(points[:, 1]) 24 | size = max(right - left, bottom - top) 25 | # + old_size*0.1]) 26 | center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0]) 27 | return center, size 28 | # translate center 29 | 30 | 31 | def augment_bbox(center, bbox_size, scale=[1.0, 1.0], trans_scale=0.0): 32 | trans_scale = (np.random.rand(2) * 2 - 1) * trans_scale 33 | center = center + trans_scale * bbox_size # 0.5 34 | scale = np.random.rand() * (scale[1] - scale[0]) + scale[0] 35 | size = int(bbox_size * scale) 36 | return center, size 37 | 38 | 39 | def crop_array(image, center, bboxsize, crop_size): 40 | """for single image only 41 | Args: 42 | image (numpy.Array): the reference array of shape HxWXC. 43 | size (Tuple[int, int]): a tuple with the height and width that will be 44 | used to resize the extracted patches. 45 | Returns: 46 | cropped_image 47 | tform: 3x3 affine matrix 48 | """ 49 | # points: top-left, top-right, bottom-right 50 | src_pts = np.array([ 51 | [center[0] - bboxsize / 2, center[1] - bboxsize / 2], 52 | [center[0] + bboxsize / 2, center[1] - bboxsize / 2], 53 | [center[0] + bboxsize / 2, center[1] + bboxsize / 2], 54 | ]) 55 | DST_PTS = np.array([[0, 0], [crop_size - 1, 0], [crop_size - 1, crop_size - 1]]) 56 | 57 | # estimate transformation between points 58 | tform = estimate_transform("similarity", src_pts, DST_PTS) 59 | 60 | # warp images 61 | cropped_image = warp(image, tform.inverse, output_shape=(crop_size, crop_size)) 62 | 63 | return cropped_image, tform.params.T 64 | 65 | 66 | class Cropper(object): 67 | def __init__(self, crop_size, scale=[1, 1], trans_scale=0.0): 68 | self.crop_size = crop_size 69 | self.scale = scale 70 | self.trans_scale = trans_scale 71 | 72 | def crop(self, image, points, points_scale=None): 73 | # points to bbox 74 | center, bbox_size = points2bbox(points, points_scale) 75 | # argument bbox. 76 | center, bbox_size = augment_bbox( 77 | center, bbox_size, scale=self.scale, trans_scale=self.trans_scale 78 | ) 79 | # crop 80 | cropped_image, tform = crop_array(image, center, bbox_size, self.crop_size) 81 | return cropped_image, tform 82 | -------------------------------------------------------------------------------- /lib/pymafx/configs/pymafx_config.yaml: -------------------------------------------------------------------------------- 1 | SOLVER: 2 | MAX_ITER: 500000 3 | TYPE: Adam 4 | BASE_LR: 0.00005 5 | GAMMA: 0.1 6 | STEPS: [0] 7 | EPOCHS: [0] 8 | # DEBUG: False 9 | LOGDIR: '' 10 | DEVICE: cuda 11 | # NUM_WORKERS: 8 12 | SEED_VALUE: -1 13 | LOSS: 14 | KP_2D_W: 300.0 15 | KP_3D_W: 300.0 16 | HF_KP_2D_W: 1000.0 17 | HF_KP_3D_W: 1000.0 18 | GL_HF_KP_2D_W: 30. 19 | FEET_KP_2D_W: 0. 20 | SHAPE_W: 0.06 21 | POSE_W: 60.0 22 | VERT_W: 0.0 23 | VERT_REG_W: 300.0 24 | INDEX_WEIGHTS: 2.0 25 | # Loss weights for surface parts. (24 Parts) 26 | PART_WEIGHTS: 0.3 27 | # Loss weights for UV regression. 28 | POINT_REGRESSION_WEIGHTS: 0.5 29 | TRAIN: 30 | NUM_WORKERS: 8 31 | BATCH_SIZE: 64 32 | LOG_FERQ: 100 33 | SHUFFLE: True 34 | PIN_MEMORY: True 35 | BHF_MODE: 'full_body' 36 | TEST: 37 | BATCH_SIZE: 32 38 | MODEL: 39 | # IWP, Identity rotation and Weak Perspective Camera 40 | USE_IWP_CAM: True 41 | USE_GT_FL: False 42 | PRED_PITCH: False 43 | MESH_MODEL: 'smplx' 44 | ALL_GENDER: False 45 | EVAL_MODE: True 46 | PyMAF: 47 | BACKBONE: 'hr48' 48 | HF_BACKBONE: 'res50' 49 | MAF_ON: True 50 | MLP_DIM: [256, 128, 64, 5] 51 | HF_MLP_DIM: [256, 128, 64, 5] 52 | MLP_VT_DIM: [256, 128, 64, 3] 53 | N_ITER: 3 54 | SUPV_LAST: False 55 | AUX_SUPV_ON: True 56 | HF_AUX_SUPV_ON: False 57 | HF_BOX_CENTER: True 58 | DP_HEATMAP_SIZE: 56 59 | GRID_FEAT: False 60 | USE_CAM_FEAT: True 61 | HF_IMG_SIZE: 224 62 | HF_DET: 'pifpaf' 63 | OPT_WRIST: True 64 | ADAPT_INTEGR: True 65 | PRED_VIS_H: True 66 | HAND_VIS_TH: 0.1 67 | GRID_ALIGN: 68 | USE_ATT: True 69 | USE_FC: False 70 | ATT_FEAT_IDX: 2 71 | ATT_HEAD: 1 72 | ATT_STARTS: 1 73 | RES_MODEL: 74 | DECONV_WITH_BIAS: False 75 | NUM_DECONV_LAYERS: 3 76 | NUM_DECONV_FILTERS: 77 | - 256 78 | - 256 79 | - 256 80 | NUM_DECONV_KERNELS: 81 | - 4 82 | - 4 83 | - 4 84 | POSE_RES_MODEL: 85 | INIT_WEIGHTS: True 86 | NAME: 'pose_resnet' 87 | PRETR_SET: 'imagenet' # 'none' 'imagenet' 'coco' 88 | # PRETRAINED: 'data/pretrained_model/resnet50-19c8e357.pth' 89 | PRETRAINED_IM: 'data/pretrained_model/resnet50-19c8e357.pth' 90 | PRETRAINED_COCO: 'data/pretrained_model/pose_resnet_50_256x192.pth.tar' 91 | EXTRA: 92 | TARGET_TYPE: 'gaussian' 93 | HEATMAP_SIZE: 94 | - 48 95 | - 64 96 | SIGMA: 2 97 | FINAL_CONV_KERNEL: 1 98 | DECONV_WITH_BIAS: False 99 | NUM_DECONV_LAYERS: 3 100 | NUM_DECONV_FILTERS: 101 | - 256 102 | - 256 103 | - 256 104 | NUM_DECONV_KERNELS: 105 | - 4 106 | - 4 107 | - 4 108 | NUM_LAYERS: 50 109 | HR_MODEL: 110 | INIT_WEIGHTS: True 111 | NAME: pose_hrnet 112 | PRETR_SET: 'coco' # 'none' 'imagenet' 'coco' 113 | PRETRAINED_IM: 'data/pretrained_model/hrnet_w48-imgnet-8ef0771d.pth' 114 | PRETRAINED_COCO: 'data/pretrained_model/pose_hrnet_w48_256x192.pth' 115 | TARGET_TYPE: gaussian 116 | IMAGE_SIZE: 117 | - 256 118 | - 256 119 | HEATMAP_SIZE: 120 | - 64 121 | - 64 122 | SIGMA: 2 123 | EXTRA: 124 | PRETRAINED_LAYERS: 125 | - 'conv1' 126 | - 'bn1' 127 | - 'conv2' 128 | - 'bn2' 129 | - 'layer1' 130 | - 'transition1' 131 | - 'stage2' 132 | - 'transition2' 133 | - 'stage3' 134 | - 'transition3' 135 | - 'stage4' 136 | FINAL_CONV_KERNEL: 1 137 | STAGE2: 138 | NUM_MODULES: 1 139 | NUM_BRANCHES: 2 140 | BLOCK: BASIC 141 | NUM_BLOCKS: 142 | - 4 143 | - 4 144 | NUM_CHANNELS: 145 | - 48 146 | - 96 147 | FUSE_METHOD: SUM 148 | STAGE3: 149 | NUM_MODULES: 4 150 | NUM_BRANCHES: 3 151 | BLOCK: BASIC 152 | NUM_BLOCKS: 153 | - 4 154 | - 4 155 | - 4 156 | NUM_CHANNELS: 157 | - 48 158 | - 96 159 | - 192 160 | FUSE_METHOD: SUM 161 | STAGE4: 162 | NUM_MODULES: 3 163 | NUM_BRANCHES: 4 164 | BLOCK: BASIC 165 | NUM_BLOCKS: 166 | - 4 167 | - 4 168 | - 4 169 | - 4 170 | NUM_CHANNELS: 171 | - 48 172 | - 96 173 | - 192 174 | - 384 175 | FUSE_METHOD: SUM 176 | -------------------------------------------------------------------------------- /lib/pymafx/core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/pymafx/core/__init__.py -------------------------------------------------------------------------------- /lib/pymafx/core/cfgs.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import argparse 18 | import json 19 | import os 20 | import random 21 | import string 22 | from datetime import datetime 23 | 24 | from yacs.config import CfgNode as CN 25 | 26 | # Configuration variables 27 | cfg = CN(new_allowed=True) 28 | 29 | cfg.OUTPUT_DIR = 'results' 30 | cfg.DEVICE = 'cuda' 31 | cfg.DEBUG = False 32 | cfg.LOGDIR = '' 33 | cfg.VAL_VIS_BATCH_FREQ = 200 34 | cfg.TRAIN_VIS_ITER_FERQ = 1000 35 | cfg.SEED_VALUE = -1 36 | 37 | cfg.TRAIN = CN(new_allowed=True) 38 | 39 | cfg.LOSS = CN(new_allowed=True) 40 | cfg.LOSS.KP_2D_W = 300.0 41 | cfg.LOSS.KP_3D_W = 300.0 42 | cfg.LOSS.SHAPE_W = 0.06 43 | cfg.LOSS.POSE_W = 60.0 44 | cfg.LOSS.VERT_W = 0.0 45 | 46 | # Loss weights for dense correspondences 47 | cfg.LOSS.INDEX_WEIGHTS = 2.0 48 | # Loss weights for surface parts. (24 Parts) 49 | cfg.LOSS.PART_WEIGHTS = 0.3 50 | # Loss weights for UV regression. 51 | cfg.LOSS.POINT_REGRESSION_WEIGHTS = 0.5 52 | 53 | cfg.MODEL = CN(new_allowed=True) 54 | 55 | cfg.MODEL.PyMAF = CN(new_allowed=True) 56 | 57 | ## switch 58 | cfg.TRAIN.BATCH_SIZE = 64 59 | cfg.TRAIN.VAL_LOOP = True 60 | 61 | cfg.TEST = CN(new_allowed=True) 62 | 63 | 64 | def get_cfg_defaults(): 65 | """Get a yacs CfgNode object with default values for my_project.""" 66 | # Return a clone so that the defaults will not be altered 67 | # This is for the "local variable" use pattern 68 | # return cfg.clone() 69 | return cfg 70 | 71 | 72 | def update_cfg(cfg_file): 73 | # cfg = get_cfg_defaults() 74 | cfg.merge_from_file(cfg_file) 75 | # return cfg.clone() 76 | return cfg 77 | 78 | 79 | def parse_args(args): 80 | cfg_file = args.cfg_file 81 | if args.cfg_file is not None: 82 | cfg = update_cfg(args.cfg_file) 83 | else: 84 | cfg = get_cfg_defaults() 85 | 86 | if args.misc is not None: 87 | cfg.merge_from_list(args.misc) 88 | 89 | return cfg 90 | 91 | 92 | def parse_args_extend(args): 93 | if args.resume: 94 | if not os.path.exists(args.log_dir): 95 | raise ValueError('Experiment are set to resume mode, but log directory does not exist.') 96 | 97 | if args.cfg_file is not None: 98 | cfg = update_cfg(args.cfg_file) 99 | else: 100 | cfg = get_cfg_defaults() 101 | # load log's cfg 102 | cfg_file = os.path.join(args.log_dir, 'cfg.yaml') 103 | cfg = update_cfg(cfg_file) 104 | 105 | if args.misc is not None: 106 | cfg.merge_from_list(args.misc) 107 | else: 108 | parse_args(args) 109 | -------------------------------------------------------------------------------- /lib/pymafx/core/path_config.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | 3 | pymafx_data_dir = osp.join(osp.dirname(__file__), "../../../data/HPS/pymafx_data") 4 | 5 | JOINT_REGRESSOR_TRAIN_EXTRA = osp.join(pymafx_data_dir, 'J_regressor_extra.npy') 6 | JOINT_REGRESSOR_H36M = osp.join(pymafx_data_dir, 'J_regressor_h36m.npy') 7 | SMPL_MEAN_PARAMS = osp.join(pymafx_data_dir, 'smpl_mean_params.npz') 8 | SMPL_MODEL_DIR = osp.join(pymafx_data_dir, 'smpl') 9 | CHECKPOINT_FILE = osp.join(pymafx_data_dir, 'PyMAF-X_model_checkpoint.pt') 10 | PARTIAL_MESH_DIR = osp.join(pymafx_data_dir, "partial_mesh") 11 | 12 | MANO_DOWNSAMPLING = osp.join(pymafx_data_dir, 'mano_downsampling.npz') 13 | SMPL_DOWNSAMPLING = osp.join(pymafx_data_dir, 'smpl_downsampling.npz') 14 | -------------------------------------------------------------------------------- /lib/pymafx/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .hmr import hmr 2 | from .pymaf_net import pymaf_net 3 | from .smpl import SMPL 4 | -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuliangXiu/ECON/d8f4e8b7171e30868acd94a1d1f6fcc1238e3e32/lib/pymafx/models/transformers/__init__.py -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/bert/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "1.0.0" 2 | 3 | from .file_utils import PYTORCH_PRETRAINED_BERT_CACHE, cached_path 4 | from .modeling_bert import ( 5 | BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, 6 | BERT_PRETRAINED_MODEL_ARCHIVE_MAP, 7 | BertConfig, 8 | BertModel, 9 | load_tf_weights_in_bert, 10 | ) 11 | from .modeling_graphormer import Graphormer 12 | from .modeling_utils import ( 13 | CONFIG_NAME, 14 | TF_WEIGHTS_NAME, 15 | WEIGHTS_NAME, 16 | Conv1D, 17 | PretrainedConfig, 18 | PreTrainedModel, 19 | prune_layer, 20 | ) 21 | 22 | # from .e2e_body_network import Graphormer_Body_Network 23 | # from .e2e_hand_network import Graphormer_Hand_Network 24 | -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/bert/bert-base-uncased/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "architectures": [ 3 | "BertForMaskedLM" 4 | ], 5 | "attention_probs_dropout_prob": 0.1, 6 | "hidden_act": "gelu", 7 | "hidden_dropout_prob": 0.1, 8 | "hidden_size": 768, 9 | "initializer_range": 0.02, 10 | "intermediate_size": 3072, 11 | "max_position_embeddings": 512, 12 | "num_attention_heads": 12, 13 | "num_hidden_layers": 12, 14 | "type_vocab_size": 2, 15 | "vocab_size": 30522 16 | } 17 | -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/bert/e2e_body_network.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) Microsoft Corporation. 3 | Licensed under the MIT license. 4 | 5 | """ 6 | 7 | import src.modeling.data.config as cfg 8 | import torch 9 | 10 | 11 | class Graphormer_Body_Network(torch.nn.Module): 12 | ''' 13 | End-to-end Graphormer network for human pose and mesh reconstruction from a single image. 14 | ''' 15 | def __init__(self, args, config, backbone, trans_encoder, mesh_sampler): 16 | super(Graphormer_Body_Network, self).__init__() 17 | self.config = config 18 | self.config.device = args.device 19 | self.backbone = backbone 20 | self.trans_encoder = trans_encoder 21 | self.upsampling = torch.nn.Linear(431, 1723) 22 | self.upsampling2 = torch.nn.Linear(1723, 6890) 23 | self.cam_param_fc = torch.nn.Linear(3, 1) 24 | self.cam_param_fc2 = torch.nn.Linear(431, 250) 25 | self.cam_param_fc3 = torch.nn.Linear(250, 3) 26 | self.grid_feat_dim = torch.nn.Linear(1024, 2051) 27 | 28 | def forward(self, images, smpl, mesh_sampler, meta_masks=None, is_train=False): 29 | batch_size = images.size(0) 30 | # Generate T-pose template mesh 31 | template_pose = torch.zeros((1, 72)) 32 | template_pose[:, 0] = 3.1416 # Rectify "upside down" reference mesh in global coord 33 | template_pose = template_pose.cuda(self.config.device) 34 | template_betas = torch.zeros((1, 10)).cuda(self.config.device) 35 | template_vertices = smpl(template_pose, template_betas) 36 | 37 | # template mesh simplification 38 | template_vertices_sub = mesh_sampler.downsample(template_vertices) 39 | template_vertices_sub2 = mesh_sampler.downsample(template_vertices_sub, n1=1, n2=2) 40 | print( 41 | 'template_vertices', template_vertices.shape, template_vertices_sub.shape, 42 | template_vertices_sub2.shape 43 | ) 44 | 45 | # template mesh-to-joint regression 46 | template_3d_joints = smpl.get_h36m_joints(template_vertices) 47 | template_pelvis = template_3d_joints[:, cfg.H36M_J17_NAME.index('Pelvis'), :] 48 | template_3d_joints = template_3d_joints[:, cfg.H36M_J17_TO_J14, :] 49 | num_joints = template_3d_joints.shape[1] 50 | 51 | # normalize 52 | template_3d_joints = template_3d_joints - template_pelvis[:, None, :] 53 | template_vertices_sub2 = template_vertices_sub2 - template_pelvis[:, None, :] 54 | 55 | # concatinate template joints and template vertices, and then duplicate to batch size 56 | ref_vertices = torch.cat([template_3d_joints, template_vertices_sub2], dim=1) 57 | ref_vertices = ref_vertices.expand(batch_size, -1, -1) 58 | print('ref_vertices', ref_vertices.shape) 59 | 60 | # extract grid features and global image features using a CNN backbone 61 | image_feat, grid_feat = self.backbone(images) 62 | print('image_feat, grid_feat', image_feat.shape, grid_feat.shape) 63 | # concatinate image feat and 3d mesh template 64 | image_feat = image_feat.view(batch_size, 1, 2048).expand(-1, ref_vertices.shape[-2], -1) 65 | print('image_feat', image_feat.shape) 66 | # process grid features 67 | grid_feat = torch.flatten(grid_feat, start_dim=2) 68 | grid_feat = grid_feat.transpose(1, 2) 69 | print('grid_feat bf', grid_feat.shape) 70 | grid_feat = self.grid_feat_dim(grid_feat) 71 | print('grid_feat', grid_feat.shape) 72 | # concatinate image feat and template mesh to form the joint/vertex queries 73 | features = torch.cat([ref_vertices, image_feat], dim=2) 74 | print('features', features.shape, ref_vertices.shape, image_feat.shape) 75 | # prepare input tokens including joint/vertex queries and grid features 76 | features = torch.cat([features, grid_feat], dim=1) 77 | print('features', features.shape) 78 | 79 | if is_train == True: 80 | # apply mask vertex/joint modeling 81 | # meta_masks is a tensor of all the masks, randomly generated in dataloader 82 | # we pre-define a [MASK] token, which is a floating-value vector with 0.01s 83 | special_token = torch.ones_like(features[:, :-49, :]).cuda() * 0.01 84 | print('special_token', special_token.shape, meta_masks.shape) 85 | print('meta_masks', torch.unique(meta_masks)) 86 | features[:, :-49, : 87 | ] = features[:, :-49, :] * meta_masks + special_token * (1 - meta_masks) 88 | 89 | # forward pass 90 | if self.config.output_attentions == True: 91 | features, hidden_states, att = self.trans_encoder(features) 92 | else: 93 | features = self.trans_encoder(features) 94 | 95 | pred_3d_joints = features[:, :num_joints, :] 96 | pred_vertices_sub2 = features[:, num_joints:-49, :] 97 | 98 | # learn camera parameters 99 | x = self.cam_param_fc(pred_vertices_sub2) 100 | x = x.transpose(1, 2) 101 | x = self.cam_param_fc2(x) 102 | x = self.cam_param_fc3(x) 103 | cam_param = x.transpose(1, 2) 104 | cam_param = cam_param.squeeze() 105 | 106 | temp_transpose = pred_vertices_sub2.transpose(1, 2) 107 | pred_vertices_sub = self.upsampling(temp_transpose) 108 | pred_vertices_full = self.upsampling2(pred_vertices_sub) 109 | pred_vertices_sub = pred_vertices_sub.transpose(1, 2) 110 | pred_vertices_full = pred_vertices_full.transpose(1, 2) 111 | 112 | if self.config.output_attentions == True: 113 | return cam_param, pred_3d_joints, pred_vertices_sub2, pred_vertices_sub, pred_vertices_full, hidden_states, att 114 | else: 115 | return cam_param, pred_3d_joints, pred_vertices_sub2, pred_vertices_sub, pred_vertices_full 116 | -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/bert/e2e_hand_network.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (c) Microsoft Corporation. 3 | Licensed under the MIT license. 4 | 5 | """ 6 | 7 | import src.modeling.data.config as cfg 8 | import torch 9 | 10 | 11 | class Graphormer_Hand_Network(torch.nn.Module): 12 | ''' 13 | End-to-end Graphormer network for hand pose and mesh reconstruction from a single image. 14 | ''' 15 | def __init__(self, args, config, backbone, trans_encoder): 16 | super(Graphormer_Hand_Network, self).__init__() 17 | self.config = config 18 | self.backbone = backbone 19 | self.trans_encoder = trans_encoder 20 | self.upsampling = torch.nn.Linear(195, 778) 21 | self.cam_param_fc = torch.nn.Linear(3, 1) 22 | self.cam_param_fc2 = torch.nn.Linear(195 + 21, 150) 23 | self.cam_param_fc3 = torch.nn.Linear(150, 3) 24 | self.grid_feat_dim = torch.nn.Linear(1024, 2051) 25 | 26 | def forward(self, images, mesh_model, mesh_sampler, meta_masks=None, is_train=False): 27 | batch_size = images.size(0) 28 | # Generate T-pose template mesh 29 | template_pose = torch.zeros((1, 48)) 30 | template_pose = template_pose.cuda() 31 | template_betas = torch.zeros((1, 10)).cuda() 32 | template_vertices, template_3d_joints = mesh_model.layer(template_pose, template_betas) 33 | template_vertices = template_vertices / 1000.0 34 | template_3d_joints = template_3d_joints / 1000.0 35 | 36 | template_vertices_sub = mesh_sampler.downsample(template_vertices) 37 | 38 | # normalize 39 | template_root = template_3d_joints[:, cfg.J_NAME.index('Wrist'), :] 40 | template_3d_joints = template_3d_joints - template_root[:, None, :] 41 | template_vertices = template_vertices - template_root[:, None, :] 42 | template_vertices_sub = template_vertices_sub - template_root[:, None, :] 43 | num_joints = template_3d_joints.shape[1] 44 | 45 | # concatinate template joints and template vertices, and then duplicate to batch size 46 | ref_vertices = torch.cat([template_3d_joints, template_vertices_sub], dim=1) 47 | ref_vertices = ref_vertices.expand(batch_size, -1, -1) 48 | 49 | # extract grid features and global image features using a CNN backbone 50 | image_feat, grid_feat = self.backbone(images) 51 | # concatinate image feat and mesh template 52 | image_feat = image_feat.view(batch_size, 1, 2048).expand(-1, ref_vertices.shape[-2], -1) 53 | # process grid features 54 | grid_feat = torch.flatten(grid_feat, start_dim=2) 55 | grid_feat = grid_feat.transpose(1, 2) 56 | grid_feat = self.grid_feat_dim(grid_feat) 57 | # concatinate image feat and template mesh to form the joint/vertex queries 58 | features = torch.cat([ref_vertices, image_feat], dim=2) 59 | # prepare input tokens including joint/vertex queries and grid features 60 | features = torch.cat([features, grid_feat], dim=1) 61 | 62 | if is_train == True: 63 | # apply mask vertex/joint modeling 64 | # meta_masks is a tensor of all the masks, randomly generated in dataloader 65 | # we pre-define a [MASK] token, which is a floating-value vector with 0.01s 66 | special_token = torch.ones_like(features[:, :-49, :]).cuda() * 0.01 67 | features[:, :-49, : 68 | ] = features[:, :-49, :] * meta_masks + special_token * (1 - meta_masks) 69 | 70 | # forward pass 71 | if self.config.output_attentions == True: 72 | features, hidden_states, att = self.trans_encoder(features) 73 | else: 74 | features = self.trans_encoder(features) 75 | 76 | pred_3d_joints = features[:, :num_joints, :] 77 | pred_vertices_sub = features[:, num_joints:-49, :] 78 | 79 | # learn camera parameters 80 | x = self.cam_param_fc(features[:, :-49, :]) 81 | x = x.transpose(1, 2) 82 | x = self.cam_param_fc2(x) 83 | x = self.cam_param_fc3(x) 84 | cam_param = x.transpose(1, 2) 85 | cam_param = cam_param.squeeze() 86 | 87 | temp_transpose = pred_vertices_sub.transpose(1, 2) 88 | pred_vertices = self.upsampling(temp_transpose) 89 | pred_vertices = pred_vertices.transpose(1, 2) 90 | 91 | if self.config.output_attentions == True: 92 | return cam_param, pred_3d_joints, pred_vertices_sub, pred_vertices, hidden_states, att 93 | else: 94 | return cam_param, pred_3d_joints, pred_vertices_sub, pred_vertices 95 | -------------------------------------------------------------------------------- /lib/pymafx/models/transformers/tokenlearner.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class SpatialAttention(nn.Module): 7 | def __init__(self) -> None: 8 | super().__init__() 9 | self.conv = nn.Sequential( 10 | nn.Conv2d(2, 1, kernel_size=(1, 1), stride=1), nn.BatchNorm2d(1), nn.ReLU() 11 | ) 12 | 13 | self.sgap = nn.AvgPool2d(2) 14 | 15 | def forward(self, x): 16 | B, H, W, C = x.shape 17 | x = x.reshape(B, C, H, W) 18 | 19 | mx = torch.max(x, 1)[0].unsqueeze(1) 20 | avg = torch.mean(x, 1).unsqueeze(1) 21 | combined = torch.cat([mx, avg], dim=1) 22 | fmap = self.conv(combined) 23 | weight_map = torch.sigmoid(fmap) 24 | out = (x * weight_map).mean(dim=(-2, -1)) 25 | 26 | return out, x * weight_map 27 | 28 | 29 | class TokenLearner(nn.Module): 30 | def __init__(self, S) -> None: 31 | super().__init__() 32 | self.S = S 33 | self.tokenizers = nn.ModuleList([SpatialAttention() for _ in range(S)]) 34 | 35 | def forward(self, x): 36 | B, _, _, C = x.shape 37 | Z = torch.Tensor(B, self.S, C).to(x) 38 | for i in range(self.S): 39 | Ai, _ = self.tokenizers[i](x) # [B, C] 40 | Z[:, i, :] = Ai 41 | return Z 42 | 43 | 44 | class TokenFuser(nn.Module): 45 | def __init__(self, H, W, C, S) -> None: 46 | super().__init__() 47 | self.projection = nn.Linear(S, S, bias=False) 48 | self.Bi = nn.Linear(C, S) 49 | self.spatial_attn = SpatialAttention() 50 | self.S = S 51 | 52 | def forward(self, y, x): 53 | B, S, C = y.shape 54 | B, H, W, C = x.shape 55 | 56 | Y = self.projection(y.reshape(B, C, S)).reshape(B, S, C) 57 | Bw = torch.sigmoid(self.Bi(x)).reshape(B, H * W, S) # [B, HW, S] 58 | BwY = torch.matmul(Bw, Y) 59 | 60 | _, xj = self.spatial_attn(x) 61 | xj = xj.reshape(B, H * W, C) 62 | 63 | out = (BwY + xj).reshape(B, H, W, C) 64 | 65 | return out 66 | -------------------------------------------------------------------------------- /lib/pymafx/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .data_loader import CheckpointDataLoader 2 | from .saver import CheckpointSaver -------------------------------------------------------------------------------- /lib/pymafx/utils/cam_params.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import os 18 | 19 | import joblib 20 | import numpy as np 21 | import torch 22 | from numpy.testing._private.utils import print_assert_equal 23 | 24 | from .geometry import batch_euler2matrix 25 | 26 | 27 | def f_pix2vfov(f_pix, img_h): 28 | 29 | if torch.is_tensor(f_pix): 30 | fov = 2. * torch.arctan(img_h / (2. * f_pix)) 31 | else: 32 | fov = 2. * np.arctan(img_h / (2. * f_pix)) 33 | 34 | return fov 35 | 36 | 37 | def vfov2f_pix(fov, img_h): 38 | 39 | if torch.is_tensor(fov): 40 | f_pix = img_h / 2. / torch.tan(fov / 2.) 41 | else: 42 | f_pix = img_h / 2. / np.tan(fov / 2.) 43 | 44 | return f_pix 45 | 46 | 47 | def read_cam_params(cam_params, orig_shape=None): 48 | # These are predicted camera parameters 49 | # cam_param_folder = CAM_PARAM_FOLDERS[dataset_name][cam_param_type] 50 | 51 | cam_pitch = cam_params['pitch'].item() 52 | cam_roll = cam_params['roll'].item() if 'roll' in cam_params else None 53 | 54 | cam_vfov = cam_params['vfov'].item() if 'vfov' in cam_params else None 55 | 56 | cam_focal_length = cam_params['f_pix'] 57 | 58 | orig_shape = cam_params['orig_resolution'] 59 | 60 | # cam_rotmat = batch_euler2matrix(torch.tensor([[cam_pitch, 0., cam_roll]]).float())[0] 61 | cam_rotmat = batch_euler2matrix(torch.tensor([[cam_pitch, 0., 0.]]).float())[0] 62 | 63 | pred_cam_int = torch.zeros(3, 3) 64 | 65 | cx, cy = orig_shape[1] / 2, orig_shape[0] / 2 66 | 67 | pred_cam_int[0, 0] = cam_focal_length 68 | pred_cam_int[1, 1] = cam_focal_length 69 | 70 | pred_cam_int[:-1, -1] = torch.tensor([cx, cy]) 71 | 72 | cam_int = pred_cam_int.float() 73 | 74 | return cam_rotmat, cam_int, cam_vfov, cam_pitch, cam_roll, cam_focal_length 75 | 76 | 77 | def homo_vector(vector): 78 | """ 79 | vector: B x N x C 80 | h_vector: B x N x (C + 1) 81 | """ 82 | 83 | batch_size, n_pts = vector.shape[:2] 84 | 85 | h_vector = torch.cat([vector, torch.ones((batch_size, n_pts, 1)).to(vector)], dim=-1) 86 | return h_vector 87 | -------------------------------------------------------------------------------- /lib/pymafx/utils/collections.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-present, Facebook, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | ############################################################################## 15 | """A simple attribute dictionary used for representing configuration options.""" 16 | 17 | from __future__ import ( 18 | absolute_import, 19 | division, 20 | print_function, 21 | unicode_literals, 22 | ) 23 | 24 | 25 | class AttrDict(dict): 26 | 27 | IMMUTABLE = '__immutable__' 28 | 29 | def __init__(self, *args, **kwargs): 30 | super().__init__(*args, **kwargs) 31 | self.__dict__[AttrDict.IMMUTABLE] = False 32 | 33 | def __getattr__(self, name): 34 | if name in self.__dict__: 35 | return self.__dict__[name] 36 | elif name in self: 37 | return self[name] 38 | else: 39 | raise AttributeError(name) 40 | 41 | def __setattr__(self, name, value): 42 | if not self.__dict__[AttrDict.IMMUTABLE]: 43 | if name in self.__dict__: 44 | self.__dict__[name] = value 45 | else: 46 | self[name] = value 47 | else: 48 | raise AttributeError( 49 | 'Attempted to set "{}" to "{}", but AttrDict is immutable'.format(name, value) 50 | ) 51 | 52 | def immutable(self, is_immutable): 53 | """Set immutability to is_immutable and recursively apply the setting 54 | to all nested AttrDicts. 55 | """ 56 | self.__dict__[AttrDict.IMMUTABLE] = is_immutable 57 | # Recursively set immutable state 58 | for v in self.__dict__.values(): 59 | if isinstance(v, AttrDict): 60 | v.immutable(is_immutable) 61 | for v in self.values(): 62 | if isinstance(v, AttrDict): 63 | v.immutable(is_immutable) 64 | 65 | def is_immutable(self): 66 | return self.__dict__[AttrDict.IMMUTABLE] 67 | -------------------------------------------------------------------------------- /lib/pymafx/utils/colormap.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-present, Facebook, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | ############################################################################## 15 | """An awesome colormap for really neat visualizations.""" 16 | 17 | from __future__ import ( 18 | absolute_import, 19 | division, 20 | print_function, 21 | unicode_literals, 22 | ) 23 | 24 | import numpy as np 25 | 26 | 27 | def colormap(rgb=False): 28 | color_list = np.array([ 29 | 0.000, 0.447, 0.741, 0.850, 0.325, 0.098, 0.929, 0.694, 0.125, 0.494, 0.184, 0.556, 0.466, 30 | 0.674, 0.188, 0.301, 0.745, 0.933, 0.635, 0.078, 0.184, 0.300, 0.300, 0.300, 0.600, 0.600, 31 | 0.600, 1.000, 0.000, 0.000, 1.000, 0.500, 0.000, 0.749, 0.749, 0.000, 0.000, 1.000, 0.000, 32 | 0.000, 0.000, 1.000, 0.667, 0.000, 1.000, 0.333, 0.333, 0.000, 0.333, 0.667, 0.000, 0.333, 33 | 1.000, 0.000, 0.667, 0.333, 0.000, 0.667, 0.667, 0.000, 0.667, 1.000, 0.000, 1.000, 0.333, 34 | 0.000, 1.000, 0.667, 0.000, 1.000, 1.000, 0.000, 0.000, 0.333, 0.500, 0.000, 0.667, 0.500, 35 | 0.000, 1.000, 0.500, 0.333, 0.000, 0.500, 0.333, 0.333, 0.500, 0.333, 0.667, 0.500, 0.333, 36 | 1.000, 0.500, 0.667, 0.000, 0.500, 0.667, 0.333, 0.500, 0.667, 0.667, 0.500, 0.667, 1.000, 37 | 0.500, 1.000, 0.000, 0.500, 1.000, 0.333, 0.500, 1.000, 0.667, 0.500, 1.000, 1.000, 0.500, 38 | 0.000, 0.333, 1.000, 0.000, 0.667, 1.000, 0.000, 1.000, 1.000, 0.333, 0.000, 1.000, 0.333, 39 | 0.333, 1.000, 0.333, 0.667, 1.000, 0.333, 1.000, 1.000, 0.667, 0.000, 1.000, 0.667, 0.333, 40 | 1.000, 0.667, 0.667, 1.000, 0.667, 1.000, 1.000, 1.000, 0.000, 1.000, 1.000, 0.333, 1.000, 41 | 1.000, 0.667, 1.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 42 | 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, 0.167, 0.000, 0.000, 0.333, 43 | 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 44 | 0.000, 0.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 0.000, 45 | 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, 0.143, 0.143, 0.143, 0.286, 0.286, 46 | 0.286, 0.429, 0.429, 0.429, 0.571, 0.571, 0.571, 0.714, 0.714, 0.714, 0.857, 0.857, 0.857, 47 | 1.000, 1.000, 1.000 48 | ]).astype(np.float32) 49 | color_list = color_list.reshape((-1, 3)) * 255 50 | if not rgb: 51 | color_list = color_list[:, ::-1] 52 | return color_list 53 | -------------------------------------------------------------------------------- /lib/pymafx/utils/data_loader.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | 3 | import torch 4 | from torch.utils.data import DataLoader 5 | from torch.utils.data.sampler import Sampler 6 | 7 | 8 | class RandomSampler(Sampler): 9 | def __init__(self, data_source, checkpoint): 10 | self.data_source = data_source 11 | if checkpoint is not None and checkpoint['dataset_perm'] is not None: 12 | self.dataset_perm = checkpoint['dataset_perm'] 13 | self.perm = self.dataset_perm[checkpoint['batch_size'] * checkpoint['batch_idx']:] 14 | else: 15 | self.dataset_perm = torch.randperm(len(self.data_source)).tolist() 16 | self.perm = torch.randperm(len(self.data_source)).tolist() 17 | 18 | def __iter__(self): 19 | return iter(self.perm) 20 | 21 | def __len__(self): 22 | return len(self.perm) 23 | 24 | 25 | class SequentialSampler(Sampler): 26 | def __init__(self, data_source, checkpoint): 27 | self.data_source = data_source 28 | if checkpoint is not None and checkpoint['dataset_perm'] is not None: 29 | self.dataset_perm = checkpoint['dataset_perm'] 30 | self.perm = self.dataset_perm[checkpoint['batch_size'] * checkpoint['batch_idx']:] 31 | else: 32 | self.dataset_perm = list(range(len(self.data_source))) 33 | self.perm = self.dataset_perm 34 | 35 | def __iter__(self): 36 | return iter(self.perm) 37 | 38 | def __len__(self): 39 | return len(self.perm) 40 | 41 | 42 | class CheckpointDataLoader(DataLoader): 43 | """ 44 | Extends torch.utils.data.DataLoader to handle resuming training from an arbitrary point within an epoch. 45 | """ 46 | def __init__( 47 | self, 48 | dataset, 49 | checkpoint=None, 50 | batch_size=1, 51 | shuffle=False, 52 | num_workers=0, 53 | pin_memory=False, 54 | drop_last=True, 55 | timeout=0, 56 | worker_init_fn=None 57 | ): 58 | 59 | if shuffle: 60 | sampler = RandomSampler(dataset, checkpoint) 61 | else: 62 | sampler = SequentialSampler(dataset, checkpoint) 63 | if checkpoint is not None: 64 | self.checkpoint_batch_idx = checkpoint['batch_idx'] 65 | else: 66 | self.checkpoint_batch_idx = 0 67 | 68 | super(CheckpointDataLoader, self).__init__( 69 | dataset, 70 | sampler=sampler, 71 | shuffle=False, 72 | batch_size=batch_size, 73 | num_workers=num_workers, 74 | drop_last=drop_last, 75 | pin_memory=pin_memory, 76 | timeout=timeout, 77 | worker_init_fn=None 78 | ) 79 | -------------------------------------------------------------------------------- /lib/pymafx/utils/io.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-present, Facebook, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | ############################################################################## 15 | """IO utilities.""" 16 | 17 | from __future__ import ( 18 | absolute_import, 19 | division, 20 | print_function, 21 | unicode_literals, 22 | ) 23 | 24 | import hashlib 25 | import logging 26 | import os 27 | import re 28 | import sys 29 | 30 | from six.moves import cPickle as pickle 31 | 32 | try: 33 | from urllib.request import urlopen 34 | except ImportError: #python2 35 | from urllib2 import urlopen 36 | 37 | logger = logging.getLogger(__name__) 38 | 39 | _DETECTRON_S3_BASE_URL = 'https://s3-us-west-2.amazonaws.com/detectron' 40 | 41 | 42 | def save_object(obj, file_name): 43 | """Save a Python object by pickling it.""" 44 | file_name = os.path.abspath(file_name) 45 | with open(file_name, 'wb') as f: 46 | pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL) 47 | 48 | 49 | def cache_url(url_or_file, cache_dir): 50 | """Download the file specified by the URL to the cache_dir and return the 51 | path to the cached file. If the argument is not a URL, simply return it as 52 | is. 53 | """ 54 | is_url = re.match(r'^(?:http)s?://', url_or_file, re.IGNORECASE) is not None 55 | 56 | if not is_url: 57 | return url_or_file 58 | 59 | url = url_or_file 60 | # assert url.startswith(_DETECTRON_S3_BASE_URL), \ 61 | # ('Detectron only automatically caches URLs in the Detectron S3 ' 62 | # 'bucket: {}').format(_DETECTRON_S3_BASE_URL) 63 | # 64 | # cache_file_path = url.replace(_DETECTRON_S3_BASE_URL, cache_dir) 65 | Len_filename = len(url.split('/')[-1]) 66 | BASE_URL = url[0:-Len_filename - 1] 67 | # 68 | cache_file_path = url.replace(BASE_URL, cache_dir) 69 | if os.path.exists(cache_file_path): 70 | # assert_cache_file_is_ok(url, cache_file_path) 71 | return cache_file_path 72 | 73 | cache_file_dir = os.path.dirname(cache_file_path) 74 | if not os.path.exists(cache_file_dir): 75 | os.makedirs(cache_file_dir) 76 | 77 | logger.info('Downloading remote file {} to {}'.format(url, cache_file_path)) 78 | download_url(url, cache_file_path) 79 | # assert_cache_file_is_ok(url, cache_file_path) 80 | return cache_file_path 81 | 82 | 83 | def assert_cache_file_is_ok(url, file_path): 84 | """Check that cache file has the correct hash.""" 85 | # File is already in the cache, verify that the md5sum matches and 86 | # return local path 87 | cache_file_md5sum = _get_file_md5sum(file_path) 88 | ref_md5sum = _get_reference_md5sum(url) 89 | assert cache_file_md5sum == ref_md5sum, \ 90 | ('Target URL {} appears to be downloaded to the local cache file ' 91 | '{}, but the md5 hash of the local file does not match the ' 92 | 'reference (actual: {} vs. expected: {}). You may wish to delete ' 93 | 'the cached file and try again to trigger automatic ' 94 | 'download.').format(url, file_path, cache_file_md5sum, ref_md5sum) 95 | 96 | 97 | def _progress_bar(count, total): 98 | """Report download progress. 99 | Credit: 100 | https://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console/27871113 101 | """ 102 | bar_len = 60 103 | filled_len = int(round(bar_len * count / float(total))) 104 | 105 | percents = round(100.0 * count / float(total), 1) 106 | bar = '=' * filled_len + '-' * (bar_len - filled_len) 107 | 108 | sys.stdout.write(' [{}] {}% of {:.1f}MB file \r'.format(bar, percents, total / 1024 / 1024)) 109 | sys.stdout.flush() 110 | if count >= total: 111 | sys.stdout.write('\n') 112 | 113 | 114 | def download_url(url, dst_file_path, chunk_size=8192, progress_hook=_progress_bar): 115 | """Download url and write it to dst_file_path. 116 | Credit: 117 | https://stackoverflow.com/questions/2028517/python-urllib2-progress-hook 118 | """ 119 | response = urlopen(url) 120 | total_size = response.info().getheader('Content-Length').strip() 121 | total_size = int(total_size) 122 | bytes_so_far = 0 123 | 124 | with open(dst_file_path, 'wb') as f: 125 | while 1: 126 | chunk = response.read(chunk_size) 127 | bytes_so_far += len(chunk) 128 | if not chunk: 129 | break 130 | if progress_hook: 131 | progress_hook(bytes_so_far, total_size) 132 | f.write(chunk) 133 | 134 | return bytes_so_far 135 | 136 | 137 | def _get_file_md5sum(file_name): 138 | """Compute the md5 hash of a file.""" 139 | hash_obj = hashlib.md5() 140 | with open(file_name, 'r') as f: 141 | hash_obj.update(f.read()) 142 | return hash_obj.hexdigest() 143 | 144 | 145 | def _get_reference_md5sum(url): 146 | """By convention the md5 hash for url is stored in url + '.md5sum'.""" 147 | url_md5sum = url + '.md5sum' 148 | md5sum = urlopen(url_md5sum).read().strip() 149 | return md5sum 150 | -------------------------------------------------------------------------------- /lib/pymafx/utils/part_utils.py: -------------------------------------------------------------------------------- 1 | import neural_renderer as nr 2 | import numpy as np 3 | import torch 4 | from core import path_config 5 | from models import SMPL 6 | 7 | 8 | class PartRenderer(): 9 | """Renderer used to render segmentation masks and part segmentations. 10 | Internally it uses the Neural 3D Mesh Renderer 11 | """ 12 | def __init__(self, focal_length=5000., render_res=224): 13 | # Parameters for rendering 14 | self.focal_length = focal_length 15 | self.render_res = render_res 16 | # We use Neural 3D mesh renderer for rendering masks and part segmentations 17 | self.neural_renderer = nr.Renderer( 18 | dist_coeffs=None, 19 | orig_size=self.render_res, 20 | image_size=render_res, 21 | light_intensity_ambient=1, 22 | light_intensity_directional=0, 23 | anti_aliasing=False 24 | ) 25 | self.faces = torch.from_numpy(SMPL(path_config.SMPL_MODEL_DIR).faces.astype(np.int32) 26 | ).cuda() 27 | textures = np.load(path_config.VERTEX_TEXTURE_FILE) 28 | self.textures = torch.from_numpy(textures).cuda().float() 29 | self.cube_parts = torch.cuda.FloatTensor(np.load(path_config.CUBE_PARTS_FILE)) 30 | 31 | def get_parts(self, parts, mask): 32 | """Process renderer part image to get body part indices.""" 33 | bn, c, h, w = parts.shape 34 | mask = mask.view(-1, 1) 35 | parts_index = torch.floor(100 * parts.permute(0, 2, 3, 1).contiguous().view(-1, 3)).long() 36 | parts = self.cube_parts[parts_index[:, 0], parts_index[:, 1], parts_index[:, 2], None] 37 | parts *= mask 38 | parts = parts.view(bn, h, w).long() 39 | return parts 40 | 41 | def __call__(self, vertices, camera): 42 | """Wrapper function for rendering process.""" 43 | # Estimate camera parameters given a fixed focal length 44 | cam_t = torch.stack([ 45 | camera[:, 1], camera[:, 2], 2 * self.focal_length / 46 | (self.render_res * camera[:, 0] + 1e-9) 47 | ], 48 | dim=-1) 49 | batch_size = vertices.shape[0] 50 | K = torch.eye(3, device=vertices.device) 51 | K[0, 0] = self.focal_length 52 | K[1, 1] = self.focal_length 53 | K[2, 2] = 1 54 | K[0, 2] = self.render_res / 2. 55 | K[1, 2] = self.render_res / 2. 56 | K = K[None, :, :].expand(batch_size, -1, -1) 57 | R = torch.eye(3, device=vertices.device)[None, :, :].expand(batch_size, -1, -1) 58 | faces = self.faces[None, :, :].expand(batch_size, -1, -1) 59 | parts, _, mask = self.neural_renderer( 60 | vertices, 61 | faces, 62 | textures=self.textures.expand(batch_size, -1, -1, -1, -1, -1), 63 | K=K, 64 | R=R, 65 | t=cam_t.unsqueeze(1) 66 | ) 67 | parts = self.get_parts(parts, mask) 68 | return mask, parts 69 | -------------------------------------------------------------------------------- /lib/pymafx/utils/pose_tracker.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | import json 18 | import os 19 | import os.path as osp 20 | import shutil 21 | import subprocess 22 | 23 | import numpy as np 24 | 25 | 26 | def run_openpose( 27 | video_file, 28 | output_folder, 29 | staf_folder, 30 | vis=False, 31 | ): 32 | pwd = os.getcwd() 33 | 34 | os.chdir(staf_folder) 35 | 36 | render = 1 if vis else 0 37 | display = 2 if vis else 0 38 | cmd = [ 39 | 'build/examples/openpose/openpose.bin', '--model_pose', 'BODY_21A', '--tracking', '1', 40 | '--render_pose', 41 | str(render), '--video', video_file, '--write_json', output_folder, '--display', 42 | str(display) 43 | ] 44 | 45 | print('Executing', ' '.join(cmd)) 46 | subprocess.call(cmd) 47 | os.chdir(pwd) 48 | 49 | 50 | def read_posetrack_keypoints(output_folder): 51 | 52 | people = dict() 53 | 54 | for idx, result_file in enumerate(sorted(os.listdir(output_folder))): 55 | json_file = osp.join(output_folder, result_file) 56 | data = json.load(open(json_file)) 57 | # print(idx, data) 58 | for person in data['people']: 59 | person_id = person['person_id'][0] 60 | joints2d = person['pose_keypoints_2d'] 61 | if person_id in people.keys(): 62 | people[person_id]['joints2d'].append(joints2d) 63 | people[person_id]['frames'].append(idx) 64 | else: 65 | people[person_id] = { 66 | 'joints2d': [], 67 | 'frames': [], 68 | } 69 | people[person_id]['joints2d'].append(joints2d) 70 | people[person_id]['frames'].append(idx) 71 | 72 | for k in people.keys(): 73 | people[k]['joints2d'] = np.array(people[k]['joints2d']).reshape( 74 | (len(people[k]['joints2d']), -1, 3) 75 | ) 76 | people[k]['frames'] = np.array(people[k]['frames']) 77 | 78 | return people 79 | 80 | 81 | def run_posetracker(video_file, staf_folder, posetrack_output_folder='/tmp', display=False): 82 | posetrack_output_folder = os.path.join( 83 | posetrack_output_folder, f'{os.path.basename(video_file)}_posetrack' 84 | ) 85 | 86 | # run posetrack on video 87 | run_openpose(video_file, posetrack_output_folder, vis=display, staf_folder=staf_folder) 88 | 89 | people_dict = read_posetrack_keypoints(posetrack_output_folder) 90 | 91 | shutil.rmtree(posetrack_output_folder) 92 | 93 | return people_dict 94 | -------------------------------------------------------------------------------- /lib/pymafx/utils/pose_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Parts of the code are adapted from https://github.com/akanazawa/hmr 3 | """ 4 | from __future__ import absolute_import, division, print_function 5 | 6 | import numpy as np 7 | import torch 8 | 9 | 10 | def compute_similarity_transform(S1, S2): 11 | """ 12 | Computes a similarity transform (sR, t) that takes 13 | a set of 3D points S1 (3 x N) closest to a set of 3D points S2, 14 | where R is an 3x3 rotation matrix, t 3x1 translation, s scale. 15 | i.e. solves the orthogonal Procrutes problem. 16 | """ 17 | transposed = False 18 | if S1.shape[0] != 3 and S1.shape[0] != 2: 19 | S1 = S1.T 20 | S2 = S2.T 21 | transposed = True 22 | assert (S2.shape[1] == S1.shape[1]) 23 | 24 | # 1. Remove mean. 25 | mu1 = S1.mean(axis=1, keepdims=True) 26 | mu2 = S2.mean(axis=1, keepdims=True) 27 | X1 = S1 - mu1 28 | X2 = S2 - mu2 29 | 30 | # 2. Compute variance of X1 used for scale. 31 | var1 = np.sum(X1**2) 32 | 33 | # 3. The outer product of X1 and X2. 34 | K = X1.dot(X2.T) 35 | 36 | # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are 37 | # singular vectors of K. 38 | U, s, Vh = np.linalg.svd(K) 39 | V = Vh.T 40 | # Construct Z that fixes the orientation of R to get det(R)=1. 41 | Z = np.eye(U.shape[0]) 42 | Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T))) 43 | # Construct R. 44 | R = V.dot(Z.dot(U.T)) 45 | 46 | # 5. Recover scale. 47 | scale = np.trace(R.dot(K)) / var1 48 | 49 | # 6. Recover translation. 50 | t = mu2 - scale * (R.dot(mu1)) 51 | 52 | # 7. Error: 53 | S1_hat = scale * R.dot(S1) + t 54 | 55 | if transposed: 56 | S1_hat = S1_hat.T 57 | 58 | return S1_hat 59 | 60 | 61 | def compute_similarity_transform_batch(S1, S2): 62 | """Batched version of compute_similarity_transform.""" 63 | S1_hat = np.zeros_like(S1) 64 | for i in range(S1.shape[0]): 65 | S1_hat[i] = compute_similarity_transform(S1[i], S2[i]) 66 | return S1_hat 67 | 68 | 69 | def reconstruction_error(S1, S2, reduction='mean'): 70 | """Do Procrustes alignment and compute reconstruction error.""" 71 | S1_hat = compute_similarity_transform_batch(S1, S2) 72 | re = np.sqrt(((S1_hat - S2)**2).sum(axis=-1)).mean(axis=-1) 73 | if reduction == 'mean': 74 | re = re.mean() 75 | elif reduction == 'sum': 76 | re = re.sum() 77 | return re, S1_hat 78 | 79 | 80 | # https://math.stackexchange.com/questions/382760/composition-of-two-axis-angle-rotations 81 | def axis_angle_add(theta, roll_axis, alpha): 82 | """Composition of two axis-angle rotations (PyTorch version) 83 | Args: 84 | theta: N x 3 85 | roll_axis: N x 3 86 | alph: N x 1 87 | Returns: 88 | equivalent axis-angle of the composition 89 | """ 90 | alpha = alpha / 2. 91 | 92 | l2norm = torch.norm(theta + 1e-8, p=2, dim=1) 93 | angle = torch.unsqueeze(l2norm, -1) 94 | 95 | normalized = torch.div(theta, angle) 96 | angle = angle * 0.5 97 | b_cos = torch.cos(angle).cpu() 98 | b_sin = torch.sin(angle).cpu() 99 | 100 | a_cos = torch.cos(alpha) 101 | a_sin = torch.sin(alpha) 102 | 103 | dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True) 104 | cross_mm = torch.zeros_like(normalized) 105 | cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1] 106 | cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2] 107 | cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0] 108 | 109 | c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm 110 | c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm 111 | 112 | c_angle = 2 * torch.acos(c_cos) 113 | c_sin = torch.sin(c_angle * 0.5) 114 | c_n = (c_angle / c_sin) * c_sin_n 115 | 116 | return c_n 117 | 118 | 119 | def axis_angle_add_np(theta, roll_axis, alpha): 120 | """Composition of two axis-angle rotations (NumPy version) 121 | Args: 122 | theta: N x 3 123 | roll_axis: N x 3 124 | alph: N x 1 125 | Returns: 126 | equivalent axis-angle of the composition 127 | """ 128 | alpha = alpha / 2. 129 | 130 | angle = np.linalg.norm(theta + 1e-8, ord=2, axis=1, keepdims=True) 131 | normalized = np.divide(theta, angle) 132 | angle = angle * 0.5 133 | 134 | b_cos = np.cos(angle) 135 | b_sin = np.sin(angle) 136 | a_cos = np.cos(alpha) 137 | a_sin = np.sin(alpha) 138 | 139 | dot_mm = np.sum(normalized * roll_axis, axis=1, keepdims=True) 140 | cross_mm = np.zeros_like(normalized) 141 | cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1] 142 | cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2] 143 | cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0] 144 | 145 | c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm 146 | c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm 147 | c_angle = 2 * np.arccos(c_cos) 148 | c_sin = np.sin(c_angle * 0.5) 149 | c_n = (c_angle / c_sin) * c_sin_n 150 | 151 | return c_n 152 | -------------------------------------------------------------------------------- /lib/pymafx/utils/sample_mesh.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import trimesh 5 | 6 | from .utils.libmesh import check_mesh_contains 7 | 8 | 9 | def get_occ_gt( 10 | in_path=None, 11 | vertices=None, 12 | faces=None, 13 | pts_num=1000, 14 | points_sigma=0.01, 15 | with_dp=False, 16 | points=None, 17 | extra_points=None 18 | ): 19 | if in_path is not None: 20 | mesh = trimesh.load(in_path, process=False) 21 | print(type(mesh.vertices), mesh.vertices.shape, mesh.faces.shape) 22 | 23 | mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) 24 | 25 | # print('get_occ_gt', type(mesh.vertices), mesh.vertices.shape, mesh.faces.shape) 26 | 27 | # points_size = 100000 28 | points_padding = 0.1 29 | # points_sigma = 0.01 30 | points_uniform_ratio = 0.5 31 | n_points_uniform = int(pts_num * points_uniform_ratio) 32 | n_points_surface = pts_num - n_points_uniform 33 | 34 | if points is None: 35 | points_scale = 2.0 36 | boxsize = points_scale + points_padding 37 | points_uniform = np.random.rand(n_points_uniform, 3) 38 | points_uniform = boxsize * (points_uniform - 0.5) 39 | points_surface, index_surface = mesh.sample(n_points_surface, return_index=True) 40 | points_surface += points_sigma * np.random.randn(n_points_surface, 3) 41 | points = np.concatenate([points_uniform, points_surface], axis=0) 42 | 43 | if extra_points is not None: 44 | extra_points += points_sigma * np.random.randn(len(extra_points), 3) 45 | points = np.concatenate([points, extra_points], axis=0) 46 | 47 | occupancies = check_mesh_contains(mesh, points) 48 | 49 | index_surface = None 50 | 51 | # points = points.astype(dtype) 52 | 53 | # print('occupancies', occupancies.dtype, np.sum(occupancies), occupancies.shape) 54 | # occupancies = np.packbits(occupancies) 55 | # print('occupancies bit', occupancies.dtype, np.sum(occupancies), occupancies.shape) 56 | 57 | # print('occupancies', points.shape, occupancies.shape, occupancies.dtype, np.sum(occupancies), index_surface.shape) 58 | 59 | return_dict = {} 60 | return_dict['points'] = points 61 | return_dict['points.occ'] = occupancies 62 | return_dict['sf_sidx'] = index_surface 63 | 64 | # export_pointcloud(mesh, modelname, loc, scale, args) 65 | # export_points(mesh, modelname, loc, scale, args) 66 | return return_dict 67 | -------------------------------------------------------------------------------- /lib/pymafx/utils/smooth_bbox.py: -------------------------------------------------------------------------------- 1 | # This script is borrowed from https://github.com/akanazawa/human_dynamics/blob/master/src/util/smooth_bbox.py 2 | # Adhere to their licence to use this script 3 | 4 | import numpy as np 5 | import scipy.signal as signal 6 | from scipy.ndimage.filters import gaussian_filter1d 7 | 8 | 9 | def get_smooth_bbox_params(kps, vis_thresh=2, kernel_size=11, sigma=3): 10 | """ 11 | Computes smooth bounding box parameters from keypoints: 12 | 1. Computes bbox by rescaling the person to be around 150 px. 13 | 2. Linearly interpolates bbox params for missing annotations. 14 | 3. Median filtering 15 | 4. Gaussian filtering. 16 | 17 | Recommended thresholds: 18 | * detect-and-track: 0 19 | * 3DPW: 0.1 20 | 21 | Args: 22 | kps (list): List of kps (Nx3) or None. 23 | vis_thresh (float): Threshold for visibility. 24 | kernel_size (int): Kernel size for median filtering (must be odd). 25 | sigma (float): Sigma for gaussian smoothing. 26 | 27 | Returns: 28 | Smooth bbox params [cx, cy, scale], start index, end index 29 | """ 30 | bbox_params, start, end = get_all_bbox_params(kps, vis_thresh) 31 | smoothed = smooth_bbox_params(bbox_params, kernel_size, sigma) 32 | smoothed = np.vstack((np.zeros((start, 3)), smoothed)) 33 | return smoothed, start, end 34 | 35 | 36 | def kp_to_bbox_param(kp, vis_thresh): 37 | """ 38 | Finds the bounding box parameters from the 2D keypoints. 39 | 40 | Args: 41 | kp (Kx3): 2D Keypoints. 42 | vis_thresh (float): Threshold for visibility. 43 | 44 | Returns: 45 | [center_x, center_y, scale] 46 | """ 47 | if kp is None: 48 | return 49 | vis = kp[:, 2] > vis_thresh 50 | if not np.any(vis): 51 | return 52 | min_pt = np.min(kp[vis, :2], axis=0) 53 | max_pt = np.max(kp[vis, :2], axis=0) 54 | person_height = np.linalg.norm(max_pt - min_pt) 55 | if person_height < 0.5: 56 | return 57 | center = (min_pt + max_pt) / 2. 58 | scale = 150. / person_height 59 | return np.append(center, scale) 60 | 61 | 62 | def get_all_bbox_params(kps, vis_thresh=2): 63 | """ 64 | Finds bounding box parameters for all keypoints. 65 | 66 | Look for sequences in the middle with no predictions and linearly 67 | interpolate the bbox params for those 68 | 69 | Args: 70 | kps (list): List of kps (Kx3) or None. 71 | vis_thresh (float): Threshold for visibility. 72 | 73 | Returns: 74 | bbox_params, start_index (incl), end_index (excl) 75 | """ 76 | # keeps track of how many indices in a row with no prediction 77 | num_to_interpolate = 0 78 | start_index = -1 79 | bbox_params = np.empty(shape=(0, 3), dtype=np.float32) 80 | 81 | for i, kp in enumerate(kps): 82 | bbox_param = kp_to_bbox_param(kp, vis_thresh=vis_thresh) 83 | if bbox_param is None: 84 | num_to_interpolate += 1 85 | continue 86 | 87 | if start_index == -1: 88 | # Found the first index with a prediction! 89 | start_index = i 90 | num_to_interpolate = 0 91 | 92 | if num_to_interpolate > 0: 93 | # Linearly interpolate each param. 94 | previous = bbox_params[-1] 95 | # This will be 3x(n+2) 96 | interpolated = np.array([ 97 | np.linspace(prev, curr, num_to_interpolate + 2) 98 | for prev, curr in zip(previous, bbox_param) 99 | ]) 100 | bbox_params = np.vstack((bbox_params, interpolated.T[1:-1])) 101 | num_to_interpolate = 0 102 | bbox_params = np.vstack((bbox_params, bbox_param)) 103 | 104 | return bbox_params, start_index, i - num_to_interpolate + 1 105 | 106 | 107 | def smooth_bbox_params(bbox_params, kernel_size=11, sigma=8): 108 | """ 109 | Applies median filtering and then gaussian filtering to bounding box 110 | parameters. 111 | 112 | Args: 113 | bbox_params (Nx3): [cx, cy, scale]. 114 | kernel_size (int): Kernel size for median filtering (must be odd). 115 | sigma (float): Sigma for gaussian smoothing. 116 | 117 | Returns: 118 | Smoothed bounding box parameters (Nx3). 119 | """ 120 | smoothed = np.array([signal.medfilt(param, kernel_size) for param in bbox_params.T]).T 121 | return np.array([gaussian_filter1d(traj, sigma) for traj in smoothed.T]).T 122 | -------------------------------------------------------------------------------- /lib/pymafx/utils/transforms.py: -------------------------------------------------------------------------------- 1 | # ------------------------------------------------------------------------------ 2 | # Copyright (c) Microsoft 3 | # Licensed under the MIT License. 4 | # Written by Bin Xiao (Bin.Xiao@microsoft.com) 5 | # ------------------------------------------------------------------------------ 6 | 7 | from __future__ import absolute_import, division, print_function 8 | 9 | import cv2 10 | import numpy as np 11 | 12 | 13 | def flip_back(output_flipped, matched_parts): 14 | ''' 15 | ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width) 16 | ''' 17 | assert output_flipped.ndim == 4,\ 18 | 'output_flipped should be [batch_size, num_joints, height, width]' 19 | 20 | output_flipped = output_flipped[:, :, :, ::-1] 21 | 22 | for pair in matched_parts: 23 | tmp = output_flipped[:, pair[0], :, :].copy() 24 | output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :] 25 | output_flipped[:, pair[1], :, :] = tmp 26 | 27 | return output_flipped 28 | 29 | 30 | def fliplr_joints(joints, joints_vis, width, matched_parts): 31 | """ 32 | flip coords 33 | """ 34 | # Flip horizontal 35 | joints[:, 0] = width - joints[:, 0] - 1 36 | 37 | # Change left-right parts 38 | for pair in matched_parts: 39 | joints[pair[0], :], joints[pair[1], :] = \ 40 | joints[pair[1], :], joints[pair[0], :].copy() 41 | joints_vis[pair[0], :], joints_vis[pair[1], :] = \ 42 | joints_vis[pair[1], :], joints_vis[pair[0], :].copy() 43 | 44 | return joints * joints_vis, joints_vis 45 | 46 | 47 | def transform_preds(coords, center, scale, output_size): 48 | target_coords = np.zeros(coords.shape) 49 | trans = get_affine_transform(center, scale, 0, output_size, inv=1) 50 | for p in range(coords.shape[0]): 51 | target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans) 52 | return target_coords 53 | 54 | 55 | def get_affine_transform( 56 | center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0 57 | ): 58 | if not isinstance(scale, np.ndarray) and not isinstance(scale, list): 59 | # print(scale) 60 | scale = np.array([scale, scale]) 61 | 62 | scale_tmp = scale * 200.0 63 | src_w = scale_tmp[0] 64 | dst_w = output_size[0] 65 | dst_h = output_size[1] 66 | 67 | rot_rad = np.pi * rot / 180 68 | src_dir = get_dir([0, src_w * -0.5], rot_rad) 69 | dst_dir = np.array([0, dst_w * -0.5], np.float32) 70 | 71 | src = np.zeros((3, 2), dtype=np.float32) 72 | dst = np.zeros((3, 2), dtype=np.float32) 73 | src[0, :] = center + scale_tmp * shift 74 | src[1, :] = center + src_dir + scale_tmp * shift 75 | dst[0, :] = [dst_w * 0.5, dst_h * 0.5] 76 | dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir 77 | 78 | src[2:, :] = get_3rd_point(src[0, :], src[1, :]) 79 | dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) 80 | 81 | if inv: 82 | trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) 83 | else: 84 | trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) 85 | 86 | return trans 87 | 88 | 89 | def affine_transform(pt, t): 90 | new_pt = np.array([pt[0], pt[1], 1.]).T 91 | new_pt = np.dot(t, new_pt) 92 | return new_pt[:2] 93 | 94 | 95 | def get_3rd_point(a, b): 96 | direct = a - b 97 | return b + np.array([-direct[1], direct[0]], dtype=np.float32) 98 | 99 | 100 | def get_dir(src_point, rot_rad): 101 | sn, cs = np.sin(rot_rad), np.cos(rot_rad) 102 | 103 | src_result = [0, 0] 104 | src_result[0] = src_point[0] * cs - src_point[1] * sn 105 | src_result[1] = src_point[0] * sn + src_point[1] * cs 106 | 107 | return src_result 108 | 109 | 110 | def crop(img, center, scale, output_size, rot=0): 111 | trans = get_affine_transform(center, scale, rot, output_size) 112 | 113 | dst_img = cv2.warpAffine( 114 | img, trans, (int(output_size[0]), int(output_size[1])), flags=cv2.INTER_LINEAR 115 | ) 116 | 117 | return dst_img 118 | -------------------------------------------------------------------------------- /lib/smplx/.gitignore: -------------------------------------------------------------------------------- 1 | #### joe made this: http://goel.io/joe 2 | 3 | #####=== Python ===##### 4 | 5 | # Byte-compiled / optimized / DLL files 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | 10 | # C extensions 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *.cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # pyenv 80 | .python-version 81 | 82 | # celery beat schedule file 83 | celerybeat-schedule 84 | 85 | # SageMath parsed files 86 | *.sage.py 87 | 88 | # Environments 89 | .env 90 | .venv 91 | env/ 92 | venv/ 93 | ENV/ 94 | env.bak/ 95 | venv.bak/ 96 | 97 | # Spyder project settings 98 | .spyderproject 99 | .spyproject 100 | 101 | # Rope project settings 102 | .ropeproject 103 | 104 | # mkdocs documentation 105 | /site 106 | 107 | # mypy 108 | .mypy_cache/ 109 | models/ 110 | output/ 111 | outputs/ 112 | transfer_data/ 113 | torch-trust-ncg/ 114 | build/ 115 | -------------------------------------------------------------------------------- /lib/smplx/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from .body_models import ( 18 | FLAME, 19 | MANO, 20 | SMPL, 21 | SMPLH, 22 | SMPLX, 23 | FLAMELayer, 24 | MANOLayer, 25 | SMPLHLayer, 26 | SMPLLayer, 27 | SMPLXLayer, 28 | build_layer, 29 | create, 30 | ) 31 | -------------------------------------------------------------------------------- /lib/smplx/joint_names.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | JOINT_NAMES = [ 18 | "pelvis", 19 | "left_hip", 20 | "right_hip", 21 | "spine1", 22 | "left_knee", 23 | "right_knee", 24 | "spine2", 25 | "left_ankle", 26 | "right_ankle", 27 | "spine3", 28 | "left_foot", 29 | "right_foot", 30 | "neck", 31 | "left_collar", 32 | "right_collar", 33 | "head", 34 | "left_shoulder", 35 | "right_shoulder", 36 | "left_elbow", 37 | "right_elbow", 38 | "left_wrist", 39 | "right_wrist", 40 | "jaw", 41 | "left_eye_smplhf", 42 | "right_eye_smplhf", 43 | "left_index1", 44 | "left_index2", 45 | "left_index3", 46 | "left_middle1", 47 | "left_middle2", 48 | "left_middle3", 49 | "left_pinky1", 50 | "left_pinky2", 51 | "left_pinky3", 52 | "left_ring1", 53 | "left_ring2", 54 | "left_ring3", 55 | "left_thumb1", 56 | "left_thumb2", 57 | "left_thumb3", 58 | "right_index1", 59 | "right_index2", 60 | "right_index3", 61 | "right_middle1", 62 | "right_middle2", 63 | "right_middle3", 64 | "right_pinky1", 65 | "right_pinky2", 66 | "right_pinky3", 67 | "right_ring1", 68 | "right_ring2", 69 | "right_ring3", 70 | "right_thumb1", 71 | "right_thumb2", 72 | "right_thumb3", 73 | "nose", 74 | "right_eye", 75 | "left_eye", 76 | "right_ear", 77 | "left_ear", 78 | "left_big_toe", 79 | "left_small_toe", 80 | "left_heel", 81 | "right_big_toe", 82 | "right_small_toe", 83 | "right_heel", 84 | "left_thumb", 85 | "left_index", 86 | "left_middle", 87 | "left_ring", 88 | "left_pinky", 89 | "right_thumb", 90 | "right_index", 91 | "right_middle", 92 | "right_ring", 93 | "right_pinky", 94 | "right_eye_brow1", 95 | "right_eye_brow2", 96 | "right_eye_brow3", 97 | "right_eye_brow4", 98 | "right_eye_brow5", 99 | "left_eye_brow5", 100 | "left_eye_brow4", 101 | "left_eye_brow3", 102 | "left_eye_brow2", 103 | "left_eye_brow1", 104 | "nose1", 105 | "nose2", 106 | "nose3", 107 | "nose4", 108 | "right_nose_2", 109 | "right_nose_1", 110 | "nose_middle", 111 | "left_nose_1", 112 | "left_nose_2", 113 | "right_eye1", 114 | "right_eye2", 115 | "right_eye3", 116 | "right_eye4", 117 | "right_eye5", 118 | "right_eye6", 119 | "left_eye4", 120 | "left_eye3", 121 | "left_eye2", 122 | "left_eye1", 123 | "left_eye6", 124 | "left_eye5", 125 | "right_mouth_1", 126 | "right_mouth_2", 127 | "right_mouth_3", 128 | "mouth_top", 129 | "left_mouth_3", 130 | "left_mouth_2", 131 | "left_mouth_1", 132 | "left_mouth_5", # 59 in OpenPose output 133 | "left_mouth_4", # 58 in OpenPose output 134 | "mouth_bottom", 135 | "right_mouth_4", 136 | "right_mouth_5", 137 | "right_lip_1", 138 | "right_lip_2", 139 | "lip_top", 140 | "left_lip_2", 141 | "left_lip_1", 142 | "left_lip_3", 143 | "lip_bottom", 144 | "right_lip_3", 145 | # Face contour 146 | "right_contour_1", 147 | "right_contour_2", 148 | "right_contour_3", 149 | "right_contour_4", 150 | "right_contour_5", 151 | "right_contour_6", 152 | "right_contour_7", 153 | "right_contour_8", 154 | "contour_middle", 155 | "left_contour_8", 156 | "left_contour_7", 157 | "left_contour_6", 158 | "left_contour_5", 159 | "left_contour_4", 160 | "left_contour_3", 161 | "left_contour_2", 162 | "left_contour_1", 163 | ] 164 | -------------------------------------------------------------------------------- /lib/smplx/utils.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from dataclasses import asdict, dataclass, fields 18 | from typing import NewType, Optional, Union 19 | 20 | import numpy as np 21 | import torch 22 | 23 | Tensor = NewType("Tensor", torch.Tensor) 24 | Array = NewType("Array", np.ndarray) 25 | 26 | 27 | @dataclass 28 | class ModelOutput: 29 | vertices: Optional[Tensor] = None 30 | joints: Optional[Tensor] = None 31 | full_pose: Optional[Tensor] = None 32 | global_orient: Optional[Tensor] = None 33 | transl: Optional[Tensor] = None 34 | 35 | def __getitem__(self, key): 36 | return getattr(self, key) 37 | 38 | def get(self, key, default=None): 39 | return getattr(self, key, default) 40 | 41 | def __iter__(self): 42 | return self.keys() 43 | 44 | def keys(self): 45 | keys = [t.name for t in fields(self)] 46 | return iter(keys) 47 | 48 | def values(self): 49 | values = [getattr(self, t.name) for t in fields(self)] 50 | return iter(values) 51 | 52 | def items(self): 53 | data = [(t.name, getattr(self, t.name)) for t in fields(self)] 54 | return iter(data) 55 | 56 | 57 | @dataclass 58 | class SMPLOutput(ModelOutput): 59 | betas: Optional[Tensor] = None 60 | body_pose: Optional[Tensor] = None 61 | 62 | 63 | @dataclass 64 | class SMPLHOutput(SMPLOutput): 65 | left_hand_pose: Optional[Tensor] = None 66 | right_hand_pose: Optional[Tensor] = None 67 | transl: Optional[Tensor] = None 68 | 69 | 70 | @dataclass 71 | class SMPLXOutput(SMPLHOutput): 72 | expression: Optional[Tensor] = None 73 | jaw_pose: Optional[Tensor] = None 74 | joint_transformation: Optional[Tensor] = None 75 | vertex_transformation: Optional[Tensor] = None 76 | 77 | 78 | @dataclass 79 | class MANOOutput(ModelOutput): 80 | betas: Optional[Tensor] = None 81 | hand_pose: Optional[Tensor] = None 82 | 83 | 84 | @dataclass 85 | class FLAMEOutput(ModelOutput): 86 | betas: Optional[Tensor] = None 87 | expression: Optional[Tensor] = None 88 | jaw_pose: Optional[Tensor] = None 89 | neck_pose: Optional[Tensor] = None 90 | 91 | 92 | def find_joint_kin_chain(joint_id, kinematic_tree): 93 | kin_chain = [] 94 | curr_idx = joint_id 95 | while curr_idx != -1: 96 | kin_chain.append(curr_idx) 97 | curr_idx = kinematic_tree[curr_idx] 98 | return kin_chain 99 | 100 | 101 | def to_tensor(array: Union[Array, Tensor], dtype=torch.float32) -> Tensor: 102 | if torch.is_tensor(array): 103 | return array 104 | else: 105 | return torch.tensor(array, dtype=dtype) 106 | 107 | 108 | class Struct(object): 109 | def __init__(self, **kwargs): 110 | for key, val in kwargs.items(): 111 | setattr(self, key, val) 112 | 113 | 114 | def to_np(array, dtype=np.float32): 115 | if "scipy.sparse" in str(type(array)): 116 | array = array.todense() 117 | return np.array(array, dtype=dtype) 118 | 119 | 120 | def rot_mat_to_euler(rot_mats): 121 | # Calculates rotation matrix to euler angles 122 | # Careful for extreme cases of eular angles like [0.0, pi, 0.0] 123 | 124 | sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] + rot_mats[:, 1, 0] * rot_mats[:, 1, 0]) 125 | return torch.atan2(-rot_mats[:, 2, 0], sy) 126 | -------------------------------------------------------------------------------- /lib/smplx/vertex_ids.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from __future__ import absolute_import, division, print_function 18 | 19 | # Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to 20 | # MSCOCO and OpenPose joints 21 | vertex_ids = { 22 | "smplh": { 23 | "nose": 332, 24 | "reye": 6260, 25 | "leye": 2800, 26 | "rear": 4071, 27 | "lear": 583, 28 | "rthumb": 6191, 29 | "rindex": 5782, 30 | "rmiddle": 5905, 31 | "rring": 6016, 32 | "rpinky": 6133, 33 | "lthumb": 2746, 34 | "lindex": 2319, 35 | "lmiddle": 2445, 36 | "lring": 2556, 37 | "lpinky": 2673, 38 | "LBigToe": 3216, 39 | "LSmallToe": 3226, 40 | "LHeel": 3387, 41 | "RBigToe": 6617, 42 | "RSmallToe": 6624, 43 | "RHeel": 6787, 44 | }, 45 | "smplx": { 46 | "nose": 9120, 47 | "reye": 9929, 48 | "leye": 9448, 49 | "rear": 616, 50 | "lear": 6, 51 | "rthumb": 8079, 52 | "rindex": 7669, 53 | "rmiddle": 7794, 54 | "rring": 7905, 55 | "rpinky": 8022, 56 | "lthumb": 5361, 57 | "lindex": 4933, 58 | "lmiddle": 5058, 59 | "lring": 5169, 60 | "lpinky": 5286, 61 | "LBigToe": 5770, 62 | "LSmallToe": 5780, 63 | "LHeel": 8846, 64 | "RBigToe": 8463, 65 | "RSmallToe": 8474, 66 | "RHeel": 8635, 67 | }, 68 | "mano": { 69 | "thumb": 744, 70 | "index": 320, 71 | "middle": 443, 72 | "ring": 554, 73 | "pinky": 671, 74 | }, 75 | } 76 | -------------------------------------------------------------------------------- /lib/smplx/vertex_joint_selector.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. All rights reserved. 14 | # 15 | # Contact: ps-license@tuebingen.mpg.de 16 | 17 | from __future__ import absolute_import, division, print_function 18 | 19 | import numpy as np 20 | import torch 21 | import torch.nn as nn 22 | 23 | from .utils import to_tensor 24 | 25 | 26 | class VertexJointSelector(nn.Module): 27 | def __init__(self, vertex_ids=None, use_hands=True, use_feet_keypoints=True, **kwargs): 28 | super(VertexJointSelector, self).__init__() 29 | 30 | extra_joints_idxs = [] 31 | 32 | face_keyp_idxs = np.array( 33 | [ 34 | vertex_ids["nose"], 35 | vertex_ids["reye"], 36 | vertex_ids["leye"], 37 | vertex_ids["rear"], 38 | vertex_ids["lear"], 39 | ], 40 | dtype=np.int64, 41 | ) 42 | 43 | extra_joints_idxs = np.concatenate([extra_joints_idxs, face_keyp_idxs]) 44 | 45 | if use_feet_keypoints: 46 | feet_keyp_idxs = np.array( 47 | [ 48 | vertex_ids["LBigToe"], 49 | vertex_ids["LSmallToe"], 50 | vertex_ids["LHeel"], 51 | vertex_ids["RBigToe"], 52 | vertex_ids["RSmallToe"], 53 | vertex_ids["RHeel"], 54 | ], 55 | dtype=np.int32, 56 | ) 57 | 58 | extra_joints_idxs = np.concatenate([extra_joints_idxs, feet_keyp_idxs]) 59 | 60 | if use_hands: 61 | self.tip_names = ["thumb", "index", "middle", "ring", "pinky"] 62 | 63 | tips_idxs = [] 64 | for hand_id in ["l", "r"]: 65 | for tip_name in self.tip_names: 66 | tips_idxs.append(vertex_ids[hand_id + tip_name]) 67 | 68 | extra_joints_idxs = np.concatenate([extra_joints_idxs, tips_idxs]) 69 | 70 | self.register_buffer("extra_joints_idxs", to_tensor(extra_joints_idxs, dtype=torch.long)) 71 | 72 | def forward(self, vertices, joints): 73 | extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs) 74 | joints = torch.cat([joints, extra_joints], dim=1) 75 | 76 | return joints 77 | -------------------------------------------------------------------------------- /lib/torch_utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | # empty 10 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # # 3 | # # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # # and proprietary rights in and to this software, related documentation 5 | # # and any modifications thereto. Any use, reproduction, disclosure or 6 | # # distribution of this software and related documentation without an express 7 | # # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | # # empty 10 | # try: 11 | # from .fused_act import FusedLeakyReLU, fused_leaky_relu 12 | # # from .upfirdn2d import upfirdn2d 13 | 14 | # print('Using custom CUDA kernels') 15 | # except Exception as e: 16 | # print(str(e)) 17 | # print('There was something wrong with the CUDA kernels') 18 | # print('Reverting to native PyTorch implementation') 19 | # from .native_ops import FusedLeakyReLU, fused_leaky_relu 20 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/bias_act.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | #include 10 | #include 11 | #include 12 | #include "bias_act.h" 13 | 14 | //------------------------------------------------------------------------ 15 | 16 | static bool has_same_layout(torch::Tensor x, torch::Tensor y) 17 | { 18 | if (x.dim() != y.dim()) 19 | return false; 20 | for (int64_t i = 0; i < x.dim(); i++) 21 | { 22 | if (x.size(i) != y.size(i)) 23 | return false; 24 | if (x.size(i) >= 2 && x.stride(i) != y.stride(i)) 25 | return false; 26 | } 27 | return true; 28 | } 29 | 30 | //------------------------------------------------------------------------ 31 | 32 | static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::Tensor xref, torch::Tensor yref, torch::Tensor dy, int grad, int dim, int act, float alpha, float gain, float clamp) 33 | { 34 | // Validate arguments. 35 | TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device"); 36 | TORCH_CHECK(b.numel() == 0 || (b.dtype() == x.dtype() && b.device() == x.device()), "b must have the same dtype and device as x"); 37 | TORCH_CHECK(xref.numel() == 0 || (xref.sizes() == x.sizes() && xref.dtype() == x.dtype() && xref.device() == x.device()), "xref must have the same shape, dtype, and device as x"); 38 | TORCH_CHECK(yref.numel() == 0 || (yref.sizes() == x.sizes() && yref.dtype() == x.dtype() && yref.device() == x.device()), "yref must have the same shape, dtype, and device as x"); 39 | TORCH_CHECK(dy.numel() == 0 || (dy.sizes() == x.sizes() && dy.dtype() == x.dtype() && dy.device() == x.device()), "dy must have the same dtype and device as x"); 40 | TORCH_CHECK(x.numel() <= INT_MAX, "x is too large"); 41 | TORCH_CHECK(b.dim() == 1, "b must have rank 1"); 42 | TORCH_CHECK(b.numel() == 0 || (dim >= 0 && dim < x.dim()), "dim is out of bounds"); 43 | TORCH_CHECK(b.numel() == 0 || b.numel() == x.size(dim), "b has wrong number of elements"); 44 | TORCH_CHECK(grad >= 0, "grad must be non-negative"); 45 | 46 | // Validate layout. 47 | TORCH_CHECK(x.is_non_overlapping_and_dense(), "x must be non-overlapping and dense"); 48 | TORCH_CHECK(b.is_contiguous(), "b must be contiguous"); 49 | TORCH_CHECK(xref.numel() == 0 || has_same_layout(xref, x), "xref must have the same layout as x"); 50 | TORCH_CHECK(yref.numel() == 0 || has_same_layout(yref, x), "yref must have the same layout as x"); 51 | TORCH_CHECK(dy.numel() == 0 || has_same_layout(dy, x), "dy must have the same layout as x"); 52 | 53 | // Create output tensor. 54 | const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); 55 | torch::Tensor y = torch::empty_like(x); 56 | TORCH_CHECK(has_same_layout(y, x), "y must have the same layout as x"); 57 | 58 | // Initialize CUDA kernel parameters. 59 | bias_act_kernel_params p; 60 | p.x = x.data_ptr(); 61 | p.b = (b.numel()) ? b.data_ptr() : NULL; 62 | p.xref = (xref.numel()) ? xref.data_ptr() : NULL; 63 | p.yref = (yref.numel()) ? yref.data_ptr() : NULL; 64 | p.dy = (dy.numel()) ? dy.data_ptr() : NULL; 65 | p.y = y.data_ptr(); 66 | p.grad = grad; 67 | p.act = act; 68 | p.alpha = alpha; 69 | p.gain = gain; 70 | p.clamp = clamp; 71 | p.sizeX = (int)x.numel(); 72 | p.sizeB = (int)b.numel(); 73 | p.stepB = (b.numel()) ? (int)x.stride(dim) : 1; 74 | 75 | // Choose CUDA kernel. 76 | void* kernel; 77 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] 78 | { 79 | kernel = choose_bias_act_kernel(p); 80 | }); 81 | TORCH_CHECK(kernel, "no CUDA kernel found for the specified activation func"); 82 | 83 | // Launch CUDA kernel. 84 | p.loopX = 4; 85 | int blockSize = 4 * 32; 86 | int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1; 87 | void* args[] = {&p}; 88 | AT_CUDA_CHECK(cudaLaunchKernel(kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream())); 89 | return y; 90 | } 91 | 92 | //------------------------------------------------------------------------ 93 | 94 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 95 | { 96 | m.def("bias_act", &bias_act); 97 | } 98 | 99 | //------------------------------------------------------------------------ 100 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/bias_act.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | //------------------------------------------------------------------------ 10 | // CUDA kernel parameters. 11 | 12 | struct bias_act_kernel_params 13 | { 14 | const void* x; // [sizeX] 15 | const void* b; // [sizeB] or NULL 16 | const void* xref; // [sizeX] or NULL 17 | const void* yref; // [sizeX] or NULL 18 | const void* dy; // [sizeX] or NULL 19 | void* y; // [sizeX] 20 | 21 | int grad; 22 | int act; 23 | float alpha; 24 | float gain; 25 | float clamp; 26 | 27 | int sizeX; 28 | int sizeB; 29 | int stepB; 30 | int loopX; 31 | }; 32 | 33 | //------------------------------------------------------------------------ 34 | // CUDA kernel selection. 35 | 36 | template void* choose_bias_act_kernel(const bias_act_kernel_params& p); 37 | 38 | //------------------------------------------------------------------------ 39 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/fma.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | """Fused multiply-add, with slightly faster gradients than `torch.addcmul()`.""" 9 | 10 | import torch 11 | 12 | #---------------------------------------------------------------------------- 13 | 14 | 15 | def fma(a, b, c): # => a * b + c 16 | return _FusedMultiplyAdd.apply(a, b, c) 17 | 18 | 19 | #---------------------------------------------------------------------------- 20 | 21 | 22 | class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c 23 | @staticmethod 24 | def forward(ctx, a, b, c): # pylint: disable=arguments-differ 25 | out = torch.addcmul(c, a, b) 26 | ctx.save_for_backward(a, b) 27 | ctx.c_shape = c.shape 28 | return out 29 | 30 | @staticmethod 31 | def backward(ctx, dout): # pylint: disable=arguments-differ 32 | a, b = ctx.saved_tensors 33 | c_shape = ctx.c_shape 34 | da = None 35 | db = None 36 | dc = None 37 | 38 | if ctx.needs_input_grad[0]: 39 | da = _unbroadcast(dout * b, a.shape) 40 | 41 | if ctx.needs_input_grad[1]: 42 | db = _unbroadcast(dout * a, b.shape) 43 | 44 | if ctx.needs_input_grad[2]: 45 | dc = _unbroadcast(dout, c_shape) 46 | 47 | return da, db, dc 48 | 49 | 50 | #---------------------------------------------------------------------------- 51 | 52 | 53 | def _unbroadcast(x, shape): 54 | extra_dims = x.ndim - len(shape) 55 | assert extra_dims >= 0 56 | dim = [ 57 | i 58 | for i in range(x.ndim) if x.shape[i] > 1 and (i < extra_dims or shape[i - extra_dims] == 1) 59 | ] 60 | if len(dim): 61 | x = x.sum(dim=dim, keepdim=True) 62 | if extra_dims: 63 | x = x.reshape(-1, *x.shape[extra_dims + 1:]) 64 | assert x.shape == shape 65 | return x 66 | 67 | 68 | #---------------------------------------------------------------------------- 69 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/fused_act.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.autograd import Function 4 | from torch.nn import functional as F 5 | from torch.utils.cpp_extension import load 6 | 7 | fused = load( 8 | 'fused', 9 | sources=['models/op/fused_bias_act.cpp', 'models/op/fused_bias_act_kernel.cu'], 10 | ) 11 | 12 | 13 | class FusedLeakyReLUFunctionBackward(Function): 14 | @staticmethod 15 | def forward(ctx, grad_output, out, bias, negative_slope, scale): 16 | ctx.save_for_backward(out) 17 | ctx.negative_slope = negative_slope 18 | ctx.scale = scale 19 | 20 | empty = grad_output.new_empty(0) 21 | 22 | grad_input = fused.fused_bias_act(grad_output, empty, out, 3, 1, negative_slope, scale) 23 | 24 | dim = [0] 25 | 26 | if grad_input.ndim > 2: 27 | dim += list(range(2, grad_input.ndim)) 28 | 29 | if bias: 30 | grad_bias = grad_input.sum(dim).detach() 31 | 32 | else: 33 | grad_bias = empty 34 | 35 | return grad_input, grad_bias 36 | 37 | @staticmethod 38 | def backward(ctx, gradgrad_input, gradgrad_bias): 39 | (out, ) = ctx.saved_tensors 40 | gradgrad_out = fused.fused_bias_act( 41 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 42 | ) 43 | 44 | return gradgrad_out, None, None, None, None 45 | 46 | 47 | class FusedLeakyReLUFunction(Function): 48 | @staticmethod 49 | def forward(ctx, input, bias, negative_slope, scale): 50 | # Added for compatibility with apex and autocast 51 | if input.dtype == torch.float16: 52 | bias = bias.half() 53 | 54 | empty = input.new_empty(0) 55 | 56 | ctx.bias = bias is not None 57 | 58 | if bias is None: 59 | bias = empty 60 | 61 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 62 | ctx.save_for_backward(out) 63 | ctx.negative_slope = negative_slope 64 | ctx.scale = scale 65 | 66 | return out 67 | 68 | @staticmethod 69 | def backward(ctx, grad_output): 70 | (out, ) = ctx.saved_tensors 71 | 72 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 73 | grad_output, out, ctx.bias, ctx.negative_slope, ctx.scale 74 | ) 75 | 76 | if not ctx.bias: 77 | grad_bias = None 78 | 79 | return grad_input, grad_bias, None, None 80 | 81 | 82 | class FusedLeakyReLU(nn.Module): 83 | def __init__(self, channel, bias=True, negative_slope=0.2, scale=2**0.5): 84 | super().__init__() 85 | 86 | if bias: 87 | self.bias = nn.Parameter(torch.zeros(channel)) 88 | else: 89 | self.bias = None 90 | 91 | self.negative_slope = negative_slope 92 | self.scale = scale 93 | 94 | def forward(self, input): 95 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 96 | 97 | 98 | def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2**0.5): 99 | if input.device.type == "cpu": 100 | if bias is not None: 101 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 102 | return F.leaky_relu( 103 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 104 | ) * scale 105 | 106 | else: 107 | return F.leaky_relu(input, negative_slope=0.2) * scale 108 | 109 | else: 110 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 111 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | 4 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 5 | int act, int grad, float alpha, float scale); 6 | 7 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 8 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 9 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 10 | 11 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale) { 13 | CHECK_CUDA(input); 14 | CHECK_CUDA(bias); 15 | 16 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 17 | } 18 | 19 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 20 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 21 | } -------------------------------------------------------------------------------- /lib/torch_utils/ops/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | // 3 | // This work is made available under the Nvidia Source Code License-NC. 4 | // To view a copy of this license, visit 5 | // https://nvlabs.github.io/stylegan2/license.html 6 | 7 | #include 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | 18 | template 19 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 20 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 21 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 22 | 23 | scalar_t zero = 0.0; 24 | 25 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 26 | scalar_t x = p_x[xi]; 27 | 28 | if (use_bias) { 29 | x += p_b[(xi / step_b) % size_b]; 30 | } 31 | 32 | scalar_t ref = use_ref ? p_ref[xi] : zero; 33 | 34 | scalar_t y; 35 | 36 | switch (act * 10 + grad) { 37 | default: 38 | case 10: y = x; break; 39 | case 11: y = x; break; 40 | case 12: y = 0.0; break; 41 | 42 | case 30: y = (x > 0.0) ? x : x * alpha; break; 43 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 44 | case 32: y = 0.0; break; 45 | } 46 | 47 | out[xi] = y * scale; 48 | } 49 | } 50 | 51 | 52 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 53 | int act, int grad, float alpha, float scale) { 54 | int curDevice = -1; 55 | cudaGetDevice(&curDevice); 56 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 57 | 58 | auto x = input.contiguous(); 59 | auto b = bias.contiguous(); 60 | auto ref = refer.contiguous(); 61 | 62 | int use_bias = b.numel() ? 1 : 0; 63 | int use_ref = ref.numel() ? 1 : 0; 64 | 65 | int size_x = x.numel(); 66 | int size_b = b.numel(); 67 | int step_b = 1; 68 | 69 | for (int i = 1 + 1; i < x.dim(); i++) { 70 | step_b *= x.size(i); 71 | } 72 | 73 | int loop_x = 4; 74 | int block_size = 4 * 32; 75 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 76 | 77 | auto y = torch::empty_like(x); 78 | 79 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 80 | fused_bias_act_kernel<<>>( 81 | y.data_ptr(), 82 | x.data_ptr(), 83 | b.data_ptr(), 84 | ref.data_ptr(), 85 | act, 86 | grad, 87 | alpha, 88 | scale, 89 | loop_x, 90 | size_x, 91 | step_b, 92 | size_b, 93 | use_bias, 94 | use_ref 95 | ); 96 | }); 97 | 98 | return y; 99 | } -------------------------------------------------------------------------------- /lib/torch_utils/ops/grid_sample_gradfix.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | """Custom replacement for `torch.nn.functional.grid_sample` that 9 | supports arbitrarily high order gradients between the input and output. 10 | Only works on 2D images and assumes 11 | `mode='bilinear'`, `padding_mode='zeros'`, `align_corners=False`.""" 12 | 13 | import warnings 14 | 15 | import torch 16 | 17 | # pylint: disable=redefined-builtin 18 | # pylint: disable=arguments-differ 19 | # pylint: disable=protected-access 20 | 21 | #---------------------------------------------------------------------------- 22 | 23 | enabled = False # Enable the custom op by setting this to true. 24 | 25 | #---------------------------------------------------------------------------- 26 | 27 | 28 | def grid_sample(input, grid): 29 | if _should_use_custom_op(): 30 | return _GridSample2dForward.apply(input, grid) 31 | return torch.nn.functional.grid_sample( 32 | input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False 33 | ) 34 | 35 | 36 | #---------------------------------------------------------------------------- 37 | 38 | 39 | def _should_use_custom_op(): 40 | if not enabled: 41 | return False 42 | if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']): 43 | return True 44 | warnings.warn( 45 | f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().' 46 | ) 47 | return False 48 | 49 | 50 | #---------------------------------------------------------------------------- 51 | 52 | 53 | class _GridSample2dForward(torch.autograd.Function): 54 | @staticmethod 55 | def forward(ctx, input, grid): 56 | assert input.ndim == 4 57 | assert grid.ndim == 4 58 | output = torch.nn.functional.grid_sample( 59 | input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False 60 | ) 61 | ctx.save_for_backward(input, grid) 62 | return output 63 | 64 | @staticmethod 65 | def backward(ctx, grad_output): 66 | input, grid = ctx.saved_tensors 67 | grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid) 68 | return grad_input, grad_grid 69 | 70 | 71 | #---------------------------------------------------------------------------- 72 | 73 | 74 | class _GridSample2dBackward(torch.autograd.Function): 75 | @staticmethod 76 | def forward(ctx, grad_output, input, grid): 77 | op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward') 78 | grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False) 79 | ctx.save_for_backward(grid) 80 | return grad_input, grad_grid 81 | 82 | @staticmethod 83 | def backward(ctx, grad2_grad_input, grad2_grad_grid): 84 | _ = grad2_grad_grid # unused 85 | grid, = ctx.saved_tensors 86 | grad2_grad_output = None 87 | grad2_input = None 88 | grad2_grid = None 89 | 90 | if ctx.needs_input_grad[0]: 91 | grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid) 92 | 93 | assert not ctx.needs_input_grad[2] 94 | return grad2_grad_output, grad2_input, grad2_grid 95 | 96 | 97 | #---------------------------------------------------------------------------- 98 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/native_ops.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | 6 | class FusedLeakyReLU(nn.Module): 7 | def __init__(self, channel, bias=True, negative_slope=0.2, scale=2**0.5): 8 | super().__init__() 9 | 10 | if bias: 11 | self.bias = nn.Parameter(torch.zeros(channel)) 12 | 13 | else: 14 | self.bias = None 15 | 16 | self.negative_slope = negative_slope 17 | self.scale = scale 18 | 19 | def forward(self, input): 20 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 21 | 22 | 23 | def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2**0.5): 24 | if input.dtype == torch.float16: 25 | bias = bias.half() 26 | 27 | if bias is not None: 28 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 29 | return F.leaky_relu( 30 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 31 | ) * scale 32 | 33 | else: 34 | return F.leaky_relu(input, negative_slope=0.2) * scale 35 | 36 | 37 | def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): 38 | up_x, up_y = up, up 39 | down_x, down_y = down, down 40 | pad_x0, pad_x1, pad_y0, pad_y1 = pad[0], pad[1], pad[0], pad[1] 41 | 42 | _, channel, in_h, in_w = input.shape 43 | input = input.reshape(-1, in_h, in_w, 1) 44 | 45 | _, in_h, in_w, minor = input.shape 46 | kernel_h, kernel_w = kernel.shape 47 | 48 | out = input.view(-1, in_h, 1, in_w, 1, minor) 49 | out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) 50 | out = out.view(-1, in_h * up_y, in_w * up_x, minor) 51 | 52 | out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]) 53 | out = out[:, 54 | max(-pad_y0, 0):out.shape[1] - max(-pad_y1, 0), 55 | max(-pad_x0, 0):out.shape[2] - max(-pad_x1, 0), :, ] 56 | 57 | out = out.permute(0, 3, 1, 2) 58 | out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) 59 | w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) 60 | out = F.conv2d(out, w) 61 | out = out.reshape( 62 | -1, 63 | minor, 64 | in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, 65 | in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, 66 | ) 67 | out = out.permute(0, 2, 3, 1) 68 | out = out[:, ::down_y, ::down_x, :] 69 | 70 | out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 71 | out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 72 | 73 | return out.view(-1, channel, out_h, out_w) 74 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | #include 10 | #include 11 | #include 12 | #include "upfirdn2d.h" 13 | 14 | //------------------------------------------------------------------------ 15 | 16 | static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain) 17 | { 18 | // Validate arguments. 19 | TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device"); 20 | TORCH_CHECK(f.device() == x.device(), "f must reside on the same device as x"); 21 | TORCH_CHECK(f.dtype() == torch::kFloat, "f must be float32"); 22 | TORCH_CHECK(x.numel() <= INT_MAX, "x is too large"); 23 | TORCH_CHECK(f.numel() <= INT_MAX, "f is too large"); 24 | TORCH_CHECK(x.dim() == 4, "x must be rank 4"); 25 | TORCH_CHECK(f.dim() == 2, "f must be rank 2"); 26 | TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, "f must be at least 1x1"); 27 | TORCH_CHECK(upx >= 1 && upy >= 1, "upsampling factor must be at least 1"); 28 | TORCH_CHECK(downx >= 1 && downy >= 1, "downsampling factor must be at least 1"); 29 | 30 | // Create output tensor. 31 | const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); 32 | int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx; 33 | int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy; 34 | TORCH_CHECK(outW >= 1 && outH >= 1, "output must be at least 1x1"); 35 | torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format()); 36 | TORCH_CHECK(y.numel() <= INT_MAX, "output is too large"); 37 | 38 | // Initialize CUDA kernel parameters. 39 | upfirdn2d_kernel_params p; 40 | p.x = x.data_ptr(); 41 | p.f = f.data_ptr(); 42 | p.y = y.data_ptr(); 43 | p.up = make_int2(upx, upy); 44 | p.down = make_int2(downx, downy); 45 | p.pad0 = make_int2(padx0, pady0); 46 | p.flip = (flip) ? 1 : 0; 47 | p.gain = gain; 48 | p.inSize = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0)); 49 | p.inStride = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0)); 50 | p.filterSize = make_int2((int)f.size(1), (int)f.size(0)); 51 | p.filterStride = make_int2((int)f.stride(1), (int)f.stride(0)); 52 | p.outSize = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0)); 53 | p.outStride = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0)); 54 | p.sizeMajor = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z; 55 | p.sizeMinor = (p.inStride.z == 1) ? p.inSize.z : 1; 56 | 57 | // Choose CUDA kernel. 58 | upfirdn2d_kernel_spec spec; 59 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] 60 | { 61 | spec = choose_upfirdn2d_kernel(p); 62 | }); 63 | 64 | // Set looping options. 65 | p.loopMajor = (p.sizeMajor - 1) / 16384 + 1; 66 | p.loopMinor = spec.loopMinor; 67 | p.loopX = spec.loopX; 68 | p.launchMinor = (p.sizeMinor - 1) / p.loopMinor + 1; 69 | p.launchMajor = (p.sizeMajor - 1) / p.loopMajor + 1; 70 | 71 | // Compute grid size. 72 | dim3 blockSize, gridSize; 73 | if (spec.tileOutW < 0) // large 74 | { 75 | blockSize = dim3(4, 32, 1); 76 | gridSize = dim3( 77 | ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor, 78 | (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1, 79 | p.launchMajor); 80 | } 81 | else // small 82 | { 83 | blockSize = dim3(256, 1, 1); 84 | gridSize = dim3( 85 | ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor, 86 | (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1, 87 | p.launchMajor); 88 | } 89 | 90 | // Launch CUDA kernel. 91 | void* args[] = {&p}; 92 | AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream())); 93 | return y; 94 | } 95 | 96 | //------------------------------------------------------------------------ 97 | 98 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 99 | { 100 | m.def("upfirdn2d", &upfirdn2d); 101 | } 102 | 103 | //------------------------------------------------------------------------ 104 | -------------------------------------------------------------------------------- /lib/torch_utils/ops/upfirdn2d.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | #include 10 | 11 | //------------------------------------------------------------------------ 12 | // CUDA kernel parameters. 13 | 14 | struct upfirdn2d_kernel_params 15 | { 16 | const void* x; 17 | const float* f; 18 | void* y; 19 | 20 | int2 up; 21 | int2 down; 22 | int2 pad0; 23 | int flip; 24 | float gain; 25 | 26 | int4 inSize; // [width, height, channel, batch] 27 | int4 inStride; 28 | int2 filterSize; // [width, height] 29 | int2 filterStride; 30 | int4 outSize; // [width, height, channel, batch] 31 | int4 outStride; 32 | int sizeMinor; 33 | int sizeMajor; 34 | 35 | int loopMinor; 36 | int loopMajor; 37 | int loopX; 38 | int launchMinor; 39 | int launchMajor; 40 | }; 41 | 42 | //------------------------------------------------------------------------ 43 | // CUDA kernel specialization. 44 | 45 | struct upfirdn2d_kernel_spec 46 | { 47 | void* kernel; 48 | int tileOutW; 49 | int tileOutH; 50 | int loopMinor; 51 | int loopX; 52 | }; 53 | 54 | //------------------------------------------------------------------------ 55 | // CUDA kernel selection. 56 | 57 | template upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p); 58 | 59 | //------------------------------------------------------------------------ 60 | -------------------------------------------------------------------------------- /loose.txt: -------------------------------------------------------------------------------- 1 | renderpeople/rp_yasmin_posed_007 2 | renderpeople/rp_victoria_posed_006 3 | renderpeople/rp_tilda_posed_005 4 | renderpeople/rp_tiffany_posed_015 5 | renderpeople/rp_tanja_posed_018 6 | renderpeople/rp_stephanie_posed_010 7 | renderpeople/rp_stacy_posed_002 8 | renderpeople/rp_serena_posed_027 9 | renderpeople/rp_serena_posed_024 10 | renderpeople/rp_seiko_posed_031 11 | renderpeople/rp_seiko_posed_015 12 | renderpeople/rp_saki_posed_033 13 | renderpeople/rp_rosy_posed_014 14 | renderpeople/rp_rosy_posed_001 15 | renderpeople/rp_roberta_posed_022 16 | renderpeople/rp_rick_posed_016 17 | renderpeople/rp_ray_posed_007 18 | renderpeople/rp_ramon_posed_002 19 | renderpeople/rp_ralph_posed_013 20 | renderpeople/rp_philip_posed_030 21 | renderpeople/rp_petra_posed_008 22 | renderpeople/rp_olivia_posed_014 23 | renderpeople/rp_olivia_posed_007 24 | renderpeople/rp_naomi_posed_034 25 | renderpeople/rp_naomi_posed_030 26 | renderpeople/rp_martha_posed_002 27 | renderpeople/rp_martha_posed_001 28 | renderpeople/rp_marleen_posed_002 29 | renderpeople/rp_lina_posed_004 30 | renderpeople/rp_kylie_posed_017 31 | renderpeople/rp_kylie_posed_006 32 | renderpeople/rp_kylie_posed_003 33 | renderpeople/rp_kent_posed_005 34 | renderpeople/rp_kent_posed_002 35 | renderpeople/rp_julia_posed_022 36 | renderpeople/rp_julia_posed_014 37 | renderpeople/rp_judy_posed_002 38 | renderpeople/rp_jessica_posed_058 39 | renderpeople/rp_jessica_posed_022 40 | renderpeople/rp_jennifer_posed_003 41 | renderpeople/rp_janna_posed_046 42 | renderpeople/rp_janna_posed_043 43 | renderpeople/rp_janna_posed_034 44 | renderpeople/rp_janna_posed_019 45 | renderpeople/rp_janett_posed_016 46 | renderpeople/rp_jamal_posed_012 47 | renderpeople/rp_helen_posed_037 48 | renderpeople/rp_fiona_posed_002 49 | renderpeople/rp_felice_posed_005 50 | renderpeople/rp_felice_posed_004 51 | renderpeople/rp_eve_posed_003 52 | renderpeople/rp_eve_posed_002 53 | renderpeople/rp_eve_posed_001 54 | renderpeople/rp_eric_posed_048 55 | renderpeople/rp_emma_posed_029 56 | renderpeople/rp_ellie_posed_015 57 | renderpeople/rp_ellie_posed_014 58 | renderpeople/rp_debra_posed_016 59 | renderpeople/rp_debra_posed_014 60 | renderpeople/rp_debra_posed_004 61 | renderpeople/rp_corey_posed_020 62 | renderpeople/rp_corey_posed_009 63 | renderpeople/rp_corey_posed_004 64 | renderpeople/rp_cody_posed_016 65 | renderpeople/rp_claudia_posed_034 66 | renderpeople/rp_claudia_posed_033 67 | renderpeople/rp_claudia_posed_024 68 | renderpeople/rp_claudia_posed_025 69 | renderpeople/rp_cindy_posed_020 70 | renderpeople/rp_christine_posed_023 71 | renderpeople/rp_christine_posed_022 72 | renderpeople/rp_christine_posed_020 73 | renderpeople/rp_christine_posed_010 74 | renderpeople/rp_carla_posed_016 75 | renderpeople/rp_caren_posed_009 76 | renderpeople/rp_caren_posed_008 77 | renderpeople/rp_brandon_posed_006 78 | renderpeople/rp_belle_posed_001 79 | renderpeople/rp_beatrice_posed_025 80 | renderpeople/rp_beatrice_posed_024 81 | renderpeople/rp_beatrice_posed_023 82 | renderpeople/rp_beatrice_posed_021 83 | renderpeople/rp_beatrice_posed_019 84 | renderpeople/rp_beatrice_posed_017 85 | renderpeople/rp_anna_posed_008 86 | renderpeople/rp_anna_posed_007 87 | renderpeople/rp_anna_posed_006 88 | renderpeople/rp_anna_posed_003 89 | renderpeople/rp_anna_posed_001 90 | renderpeople/rp_alvin_posed_016 91 | renderpeople/rp_alison_posed_028 92 | renderpeople/rp_alison_posed_024 93 | renderpeople/rp_alison_posed_017 94 | renderpeople/rp_alexandra_posed_022 95 | renderpeople/rp_alexandra_posed_023 96 | renderpeople/rp_alexandra_posed_019 97 | renderpeople/rp_alexandra_posed_018 98 | renderpeople/rp_alexandra_posed_013 99 | renderpeople/rp_alexandra_posed_012 100 | renderpeople/rp_alexandra_posed_011 101 | -------------------------------------------------------------------------------- /pose.txt: -------------------------------------------------------------------------------- 1 | cape/00215-jerseyshort-pose_model-000200 2 | cape/00134-longlong-ballet4_trial2-000250 3 | cape/00134-longlong-badminton_trial1-000230 4 | cape/00134-longlong-frisbee_trial1-000190 5 | cape/03375-shortlong-ballet1_trial1-000210 6 | cape/03375-longlong-babysit_trial2-000110 7 | cape/00134-shortlong-stretch_trial1-000310 8 | cape/03375-shortshort-lean_trial1-000060 9 | cape/03375-shortshort-swim_trial2-000110 10 | cape/03375-longlong-box_trial1-000190 11 | cape/03375-longlong-row_trial2-000150 12 | cape/00134-shortlong-hockey_trial1-000140 13 | cape/00134-shortlong-hockey_trial1-000090 14 | cape/00134-longlong-ski_trial2-000200 15 | cape/00134-longlong-stretch_trial1-000450 16 | cape/00096-shirtshort-soccer-000160 17 | cape/03375-shortshort-hands_up_trial2-000270 18 | cape/03375-shortshort-ballet1_trial1-000110 19 | cape/03375-longlong-babysit_trial2-000150 20 | cape/03375-shortshort-fashion_trial1-000140 21 | cape/00134-shortlong-ballet2_trial1-000110 22 | cape/00134-longlong-ballet2_trial1-000120 23 | cape/00134-shortlong-ballet2_trial1-000120 24 | cape/00134-shortlong-ballet2_trial1-000090 25 | cape/00134-longlong-ballet2_trial2-000110 26 | cape/00134-longlong-volleyball_trial2-000050 27 | cape/00134-longlong-stretch_trial1-000500 28 | cape/00134-longlong-housework_trial1-000380 29 | cape/00134-shortlong-dig_trial1-000150 30 | cape/03375-longlong-catchpick_trial1-000110 31 | cape/03375-shortlong-ballet1_trial1-000250 32 | cape/03375-shortlong-shoulders_trial1-000360 33 | cape/03375-shortlong-slack_trial2-000070 34 | cape/03375-shortlong-shoulders_trial1-000220 35 | cape/03375-shortlong-stretch_trial1-000330 36 | cape/00127-shortlong-ballerina_spin-000080 37 | cape/00127-shortlong-ballerina_spin-000200 38 | cape/00096-shortshort-basketball-000100 39 | cape/00096-shortshort-ballerina_spin-000160 40 | cape/00134-longlong-stretch_trial2-000440 41 | cape/02474-longlong-ATUsquat-000100 42 | cape/03375-longlong-ATUsquat_trial1-000120 43 | cape/02474-longlong-ATUsquat-000110 44 | cape/00134-longlong-ballet1_trial1-000180 45 | cape/00096-shirtlong-ATUsquat-000130 46 | cape/00032-shortshort-pose_model-000030 47 | cape/00134-shortlong-athletics_trial2-000070 48 | cape/00032-longshort-pose_model-000060 49 | cape/00032-shortshort-shoulders_mill-000060 50 | cape/00127-shortlong-pose_model-000430 51 | cape/00122-shortshort-ATUsquat-000120 52 | cape/00032-shortshort-bend_back_and_front-000220 53 | cape/00096-shortshort-squats-000180 54 | cape/00032-shortlong-squats-000090 55 | cape/03375-shortlong-ATUsquat_trial2-000080 56 | cape/03375-shortshort-lean_trial1-000130 57 | cape/03375-blazerlong-music_trial1-000150 58 | cape/03284-longlong-hips-000170 59 | cape/03375-shortlong-shoulders_trial1-000370 60 | cape/03375-shortlong-ballet1_trial1-000290 61 | cape/00215-jerseyshort-shoulders_mill-000320 62 | cape/00215-poloshort-soccer-000110 63 | cape/00122-shortshort-punching-000170 64 | cape/00096-jerseyshort-shoulders_mill-000140 65 | cape/00032-longshort-flying_eagle-000240 66 | cape/00134-shortlong-swim_trial1-000160 67 | cape/03375-shortshort-music_trial1-000120 68 | cape/03375-shortshort-handball_trial1-000120 69 | cape/00215-longshort-punching-000060 70 | cape/00134-shortlong-swim_trial2-000120 71 | cape/03375-shortshort-hands_up_trial1-000140 72 | cape/03375-shortshort-hands_up_trial1-000270 73 | cape/03375-shortshort-volleyball_trial1-000110 74 | cape/03375-shortshort-swim_trial1-000270 75 | cape/03375-longlong-row_trial2-000190 76 | cape/00215-poloshort-flying_eagle-000120 77 | cape/03223-shortshort-flying_eagle-000280 78 | cape/00096-shirtlong-shoulders_mill-000110 79 | cape/00096-shirtshort-pose_model-000190 80 | cape/03375-shortshort-swim_trial1-000190 81 | cape/03375-shortlong-music_trial2-000040 82 | cape/03375-shortlong-babysit_trial2-000070 83 | cape/00215-jerseyshort-flying_eagle-000110 84 | cape/03375-blazerlong-music_trial1-000030 85 | cape/03375-longlong-volleyball_trial2-000230 86 | cape/03375-blazerlong-lean_trial2-000110 87 | cape/03375-longlong-box_trial2-000110 88 | cape/03375-longlong-drinkeat_trial2-000050 89 | cape/00134-shortlong-slack_trial1-000150 90 | cape/03375-shortshort-climb_trial1-000170 91 | cape/00032-longshort-tilt_twist_left-000060 92 | cape/00215-longshort-chicken_wings-000060 93 | cape/00215-poloshort-bend_back_and_front-000130 94 | cape/03223-longshort-flying_eagle-000480 95 | cape/00215-longshort-bend_back_and_front-000100 96 | cape/00215-longshort-tilt_twist_left-000130 97 | cape/00096-longshort-tilt_twist_left-000150 98 | cape/03284-longshort-twist_tilt_left-000080 99 | cape/03223-shortshort-flying_eagle-000270 100 | cape/02474-longshort-improvise-000080 101 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | scikit-image 3 | trimesh 4 | rtree 5 | pytorch_lightning 6 | kornia>0.4.0 7 | chumpy 8 | opencv-python 9 | opencv_contrib_python 10 | scikit-learn 11 | protobuf 12 | dataclasses 13 | mediapipe 14 | einops 15 | boto3 16 | open3d 17 | xatlas 18 | huggingface_hub 19 | fast-simplification 20 | git+https://github.com/YuliangXiu/rembg.git 21 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [yapf] 2 | based_on_style = facebook 3 | column_limit = 100 4 | indent_width = 4 5 | spaces_before_comment = 4 6 | split_all_comma_separated_values = false 7 | split_all_top_level_comma_separated_values = false 8 | dedent_closing_brackets = true 9 | coalesce_brackets = true 10 | split_before_dot = false 11 | each_dict_entry_on_separate_line = false 12 | indent_dictionary_value = false 13 | 14 | [isort] 15 | multi_line_output = 3 16 | line_length = 80 17 | include_trailing_comma = true 18 | skip=./log,./results,./data,./debug,./lib/common/libmesh/setup.py,./lib/common/libvoxelize/setup.py --------------------------------------------------------------------------------