diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..00f45fce85d23a168874aa20d503c5605ce156a1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +data/*/* +data/thuman* +!data/tbfo.ttf +__pycache__ +debug/ +log/ +results/* +.vscode +!.gitignore +force_push.sh +.idea +smplx/ +human_det/ +kaolin/ +neural_voxelization_layer/ +pytorch3d/ +force_push.sh diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a182848e970b578164ba3c949dd8c274709376d0 --- /dev/null +++ b/LICENSE @@ -0,0 +1,53 @@ +License + +Software Copyright License for non-commercial scientific research purposes +Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the ICON model, data and software, (the "Data & Software"), including 3D meshes, images, videos, textures, software, scripts, and animations. By downloading and/or using the Data & Software (including downloading, cloning, installing, and any other use of the corresponding github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Data & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License + +Ownership / Licensees +The Software and the associated materials has been developed at the Max Planck Institute for Intelligent Systems (hereinafter "MPI"). Any copyright or patent right is owned by and proprietary material of the Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”) hereinafter the “Licensor”. + +License Grant +Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right: + + • To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization; + • To use the Model & Software for the sole purpose of performing peaceful non-commercial scientific research, non-commercial education, or non-commercial artistic projects; + • To modify, adapt, translate or create derivative works based upon the Model & Software. + +Any other use, in particular any use for commercial, pornographic, military, or surveillance, purposes is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Data & Software may not be used to create fake, libelous, misleading, or defamatory content of any kind excluding analyses in peer-reviewed scientific research. The Data & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission. + +The Data & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Software to train methods/algorithms/neural networks/etc. for commercial, pornographic, military, surveillance, or defamatory use of any kind. By downloading the Data & Software, you agree not to reverse engineer it. + +No Distribution +The Data & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only. + +Disclaimer of Representations and Warranties +You expressly acknowledge and agree that the Data & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Data & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE DATA & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Data & Software, (ii) that the use of the Data & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Data & Software will not cause any damage of any kind to you or a third party. + +Limitation of Liability +Because this Data & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage. +Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded. +Patent claims generated through the usage of the Data & Software cannot be directed towards the copyright holders. +The Data & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Data & Software and is not responsible for any problems such modifications cause. + +No Maintenance Services +You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Data & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Data & Software at any time. + +Defects of the Data & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication. + +Publications using the Model & Software +You acknowledge that the Data & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Data & Software. + +Citation: + +@inproceedings{xiu2022icon, + title={{ICON}: {I}mplicit {C}lothed humans {O}btained from {N}ormals}, + author={Xiu, Yuliang and Yang, Jinlong and Tzionas, Dimitrios and Black, Michael J.}, + booktitle={IEEE/CVF Conf.~on Computer Vision and Pattern Recognition (CVPR)}, + month = jun, + year={2022} +} + +Commercial licensing opportunities +For commercial uses of the Model & Software, please send email to ps-license@tue.mpg.de + +This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention. diff --git a/README.md b/README.md index 18ce6ddd25b120f0b70ace8230eb44153a161476..e872a98ec96036d9cb45ea2634b69455a4bfcc20 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,12 @@ --- title: ICON -emoji: 🌖 +metaTitle: "Image2Human by Yuliang Xiu" +emoji: 🤼 colorFrom: indigo -colorTo: indigo +colorTo: yellow sdk: gradio sdk_version: 3.1.1 -app_file: app.py -pinned: false -license: other ---- - -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +app_file: ./apps/app.py +pinned: true +python_version: 3.8 +--- \ No newline at end of file diff --git a/apps/ICON.py b/apps/ICON.py new file mode 100644 index 0000000000000000000000000000000000000000..ae3d7e2c93fe971af6afdffd374e78ad4acc6287 --- /dev/null +++ b/apps/ICON.py @@ -0,0 +1,762 @@ +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.common.seg3d_lossless import Seg3dLossless +from lib.dataset.Evaluator import Evaluator +from lib.net import HGPIFuNet +from lib.common.train_util import * +from lib.renderer.gl.init_gl import initialize_GL_context +from lib.common.render import Render +from lib.dataset.mesh_util import SMPLX, update_mesh_shape_prior_losses, get_visibility +import warnings +import logging +import torch +import smplx +import numpy as np +from torch import nn +from skimage.transform import resize +import pytorch_lightning as pl + +torch.backends.cudnn.benchmark = True + +logging.getLogger("lightning").setLevel(logging.ERROR) + +warnings.filterwarnings("ignore") + + +class ICON(pl.LightningModule): + def __init__(self, cfg): + super(ICON, self).__init__() + + self.cfg = cfg + self.batch_size = self.cfg.batch_size + self.lr_G = self.cfg.lr_G + + self.use_sdf = cfg.sdf + self.prior_type = cfg.net.prior_type + self.mcube_res = cfg.mcube_res + self.clean_mesh_flag = cfg.clean_mesh + + self.netG = HGPIFuNet( + self.cfg, + self.cfg.projection_mode, + error_term=nn.SmoothL1Loss() if self.use_sdf else nn.MSELoss(), + ) + + # TODO: replace the renderer from opengl to pytorch3d + self.evaluator = Evaluator( + device=torch.device(f"cuda:{self.cfg.gpus[0]}")) + + self.resolutions = ( + np.logspace( + start=5, + stop=np.log2(self.mcube_res), + base=2, + num=int(np.log2(self.mcube_res) - 4), + endpoint=True, + ) + + 1.0 + ) + self.resolutions = self.resolutions.astype(np.int16).tolist() + + self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"] + self.pamir_keys = ["voxel_verts", + "voxel_faces", "pad_v_num", "pad_f_num"] + + self.reconEngine = Seg3dLossless( + query_func=query_func, + b_min=[[-1.0, 1.0, -1.0]], + b_max=[[1.0, -1.0, 1.0]], + resolutions=self.resolutions, + align_corners=True, + balance_value=0.50, + device=torch.device(f"cuda:{self.cfg.test_gpus[0]}"), + visualize=False, + debug=False, + use_cuda_impl=False, + faster=True, + ) + + self.render = Render( + size=512, device=torch.device(f"cuda:{self.cfg.test_gpus[0]}") + ) + self.smpl_data = SMPLX() + + self.get_smpl_model = lambda smpl_type, gender, age, v_template: smplx.create( + self.smpl_data.model_dir, + kid_template_path=osp.join( + osp.realpath(self.smpl_data.model_dir), + f"{smpl_type}/{smpl_type}_kid_template.npy", + ), + model_type=smpl_type, + gender=gender, + age=age, + v_template=v_template, + use_face_contour=False, + ext="pkl", + ) + + self.in_geo = [item[0] for item in cfg.net.in_geo] + self.in_nml = [item[0] for item in cfg.net.in_nml] + self.in_geo_dim = [item[1] for item in cfg.net.in_geo] + self.in_total = self.in_geo + self.in_nml + self.smpl_dim = cfg.net.smpl_dim + + self.export_dir = None + self.result_eval = {} + + def get_progress_bar_dict(self): + tqdm_dict = super().get_progress_bar_dict() + if "v_num" in tqdm_dict: + del tqdm_dict["v_num"] + return tqdm_dict + + # Training related + def configure_optimizers(self): + + # set optimizer + weight_decay = self.cfg.weight_decay + momentum = self.cfg.momentum + + optim_params_G = [ + {"params": self.netG.if_regressor.parameters(), "lr": self.lr_G} + ] + + if self.cfg.net.use_filter: + optim_params_G.append( + {"params": self.netG.F_filter.parameters(), "lr": self.lr_G} + ) + + if self.cfg.net.prior_type == "pamir": + optim_params_G.append( + {"params": self.netG.ve.parameters(), "lr": self.lr_G} + ) + + if self.cfg.optim == "Adadelta": + + optimizer_G = torch.optim.Adadelta( + optim_params_G, lr=self.lr_G, weight_decay=weight_decay + ) + + elif self.cfg.optim == "Adam": + + optimizer_G = torch.optim.Adam( + optim_params_G, lr=self.lr_G, weight_decay=weight_decay + ) + + elif self.cfg.optim == "RMSprop": + + optimizer_G = torch.optim.RMSprop( + optim_params_G, + lr=self.lr_G, + weight_decay=weight_decay, + momentum=momentum, + ) + + else: + raise NotImplementedError + + # set scheduler + scheduler_G = torch.optim.lr_scheduler.MultiStepLR( + optimizer_G, milestones=self.cfg.schedule, gamma=self.cfg.gamma + ) + + return [optimizer_G], [scheduler_G] + + def training_step(self, batch, batch_idx): + + if not self.cfg.fast_dev: + export_cfg(self.logger, self.cfg) + + self.netG.train() + + in_tensor_dict = { + "sample": batch["samples_geo"].permute(0, 2, 1), + "calib": batch["calib"], + "label": batch["labels_geo"].unsqueeze(1), + } + + for name in self.in_total: + in_tensor_dict.update({name: batch[name]}) + + if self.prior_type == "icon": + for key in self.icon_keys: + in_tensor_dict.update({key: batch[key]}) + elif self.prior_type == "pamir": + for key in self.pamir_keys: + in_tensor_dict.update({key: batch[key]}) + else: + pass + + preds_G, error_G = self.netG(in_tensor_dict) + + acc, iou, prec, recall = self.evaluator.calc_acc( + preds_G.flatten(), + in_tensor_dict["label"].flatten(), + 0.5, + use_sdf=self.cfg.sdf, + ) + + # metrics processing + metrics_log = { + "train_loss": error_G.item(), + "train_acc": acc.item(), + "train_iou": iou.item(), + "train_prec": prec.item(), + "train_recall": recall.item(), + } + + tf_log = tf_log_convert(metrics_log) + bar_log = bar_log_convert(metrics_log) + + if batch_idx % int(self.cfg.freq_show_train) == 0: + + with torch.no_grad(): + self.render_func(in_tensor_dict, dataset="train") + + metrics_return = { + k.replace("train_", ""): torch.tensor(v) for k, v in metrics_log.items() + } + + metrics_return.update( + {"loss": error_G, "log": tf_log, "progress_bar": bar_log}) + + return metrics_return + + def training_epoch_end(self, outputs): + + if [] in outputs: + outputs = outputs[0] + + # metrics processing + metrics_log = { + "train_avgloss": batch_mean(outputs, "loss"), + "train_avgiou": batch_mean(outputs, "iou"), + "train_avgprec": batch_mean(outputs, "prec"), + "train_avgrecall": batch_mean(outputs, "recall"), + "train_avgacc": batch_mean(outputs, "acc"), + } + + tf_log = tf_log_convert(metrics_log) + + return {"log": tf_log} + + def validation_step(self, batch, batch_idx): + + self.netG.eval() + self.netG.training = False + + in_tensor_dict = { + "sample": batch["samples_geo"].permute(0, 2, 1), + "calib": batch["calib"], + "label": batch["labels_geo"].unsqueeze(1), + } + + for name in self.in_total: + in_tensor_dict.update({name: batch[name]}) + + if self.prior_type == "icon": + for key in self.icon_keys: + in_tensor_dict.update({key: batch[key]}) + elif self.prior_type == "pamir": + for key in self.pamir_keys: + in_tensor_dict.update({key: batch[key]}) + else: + pass + + preds_G, error_G = self.netG(in_tensor_dict) + + acc, iou, prec, recall = self.evaluator.calc_acc( + preds_G.flatten(), + in_tensor_dict["label"].flatten(), + 0.5, + use_sdf=self.cfg.sdf, + ) + + if batch_idx % int(self.cfg.freq_show_val) == 0: + with torch.no_grad(): + self.render_func(in_tensor_dict, dataset="val", idx=batch_idx) + + metrics_return = { + "val_loss": error_G, + "val_acc": acc, + "val_iou": iou, + "val_prec": prec, + "val_recall": recall, + } + + return metrics_return + + def validation_epoch_end(self, outputs): + + # metrics processing + metrics_log = { + "val_avgloss": batch_mean(outputs, "val_loss"), + "val_avgacc": batch_mean(outputs, "val_acc"), + "val_avgiou": batch_mean(outputs, "val_iou"), + "val_avgprec": batch_mean(outputs, "val_prec"), + "val_avgrecall": batch_mean(outputs, "val_recall"), + } + + tf_log = tf_log_convert(metrics_log) + + return {"log": tf_log} + + def compute_vis_cmap(self, smpl_type, smpl_verts, smpl_faces): + + (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1) + smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long()) + if smpl_type == "smpl": + smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0])) + else: + smplx_ind = np.arange(smpl_vis.shape[0]) + smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind) + + return { + "smpl_vis": smpl_vis.unsqueeze(0).to(self.device), + "smpl_cmap": smpl_cmap.unsqueeze(0).to(self.device), + "smpl_verts": smpl_verts.unsqueeze(0), + } + + @torch.enable_grad() + def optim_body(self, in_tensor_dict, batch): + + smpl_model = self.get_smpl_model( + batch["type"][0], batch["gender"][0], batch["age"][0], None + ).to(self.device) + in_tensor_dict["smpl_faces"] = ( + torch.tensor(smpl_model.faces.astype(np.int)) + .long() + .unsqueeze(0) + .to(self.device) + ) + + # The optimizer and variables + optimed_pose = torch.tensor( + batch["body_pose"][0], device=self.device, requires_grad=True + ) # [1,23,3,3] + optimed_trans = torch.tensor( + batch["transl"][0], device=self.device, requires_grad=True + ) # [3] + optimed_betas = torch.tensor( + batch["betas"][0], device=self.device, requires_grad=True + ) # [1,10] + optimed_orient = torch.tensor( + batch["global_orient"][0], device=self.device, requires_grad=True + ) # [1,1,3,3] + + optimizer_smpl = torch.optim.SGD( + [optimed_pose, optimed_trans, optimed_betas, optimed_orient], + lr=1e-3, + momentum=0.9, + ) + scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer_smpl, mode="min", factor=0.5, verbose=0, min_lr=1e-5, patience=5 + ) + loop_smpl = range(50) + for i in loop_smpl: + + optimizer_smpl.zero_grad() + + # prior_loss, optimed_pose = dataset.vposer_prior(optimed_pose) + smpl_out = smpl_model( + betas=optimed_betas, + body_pose=optimed_pose, + global_orient=optimed_orient, + transl=optimed_trans, + return_verts=True, + ) + + smpl_verts = smpl_out.vertices[0] * 100.0 + smpl_verts = projection( + smpl_verts, batch["calib"][0], format="tensor") + smpl_verts[:, 1] *= -1 + # render optimized mesh (normal, T_normal, image [-1,1]) + self.render.load_meshes( + smpl_verts, in_tensor_dict["smpl_faces"]) + ( + in_tensor_dict["T_normal_F"], + in_tensor_dict["T_normal_B"], + ) = self.render.get_rgb_image() + + T_mask_F, T_mask_B = self.render.get_silhouette_image() + + with torch.no_grad(): + ( + in_tensor_dict["normal_F"], + in_tensor_dict["normal_B"], + ) = self.netG.normal_filter(in_tensor_dict) + + # mask = torch.abs(in_tensor['T_normal_F']).sum(dim=0, keepdims=True) > 0.0 + diff_F_smpl = torch.abs( + in_tensor_dict["T_normal_F"] - in_tensor_dict["normal_F"] + ) + diff_B_smpl = torch.abs( + in_tensor_dict["T_normal_B"] - in_tensor_dict["normal_B"] + ) + loss = (diff_F_smpl + diff_B_smpl).mean() + + # silhouette loss + smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0] + gt_arr = torch.cat( + [in_tensor_dict["normal_F"][0], in_tensor_dict["normal_B"][0]], dim=2 + ).permute(1, 2, 0) + gt_arr = ((gt_arr + 1.0) * 0.5).to(self.device) + bg_color = ( + torch.Tensor([0.5, 0.5, 0.5]).unsqueeze( + 0).unsqueeze(0).to(self.device) + ) + gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float() + loss += torch.abs(smpl_arr - gt_arr).mean() + + # Image.fromarray(((in_tensor_dict['T_normal_F'][0].permute(1,2,0)+1.0)*0.5*255.0).detach().cpu().numpy().astype(np.uint8)).show() + + # loop_smpl.set_description(f"smpl = {loss:.3f}") + + loss.backward(retain_graph=True) + optimizer_smpl.step() + scheduler_smpl.step(loss) + in_tensor_dict["smpl_verts"] = smpl_verts.unsqueeze(0) + + in_tensor_dict.update( + self.compute_vis_cmap( + batch["type"][0], + in_tensor_dict["smpl_verts"][0], + in_tensor_dict["smpl_faces"][0], + ) + ) + + features, inter = self.netG.filter(in_tensor_dict, return_inter=True) + + return features, inter, in_tensor_dict + + @torch.enable_grad() + def optim_cloth(self, verts_pr, faces_pr, inter): + + # convert from GT to SDF + verts_pr -= (self.resolutions[-1] - 1) / 2.0 + verts_pr /= (self.resolutions[-1] - 1) / 2.0 + + losses = { + "cloth": {"weight": 5.0, "value": 0.0}, + "edge": {"weight": 100.0, "value": 0.0}, + "normal": {"weight": 0.2, "value": 0.0}, + "laplacian": {"weight": 100.0, "value": 0.0}, + "smpl": {"weight": 1.0, "value": 0.0}, + "deform": {"weight": 20.0, "value": 0.0}, + } + + deform_verts = torch.full( + verts_pr.shape, 0.0, device=self.device, requires_grad=True + ) + optimizer_cloth = torch.optim.SGD( + [deform_verts], lr=1e-1, momentum=0.9) + scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer_cloth, mode="min", factor=0.1, verbose=0, min_lr=1e-3, patience=5 + ) + # cloth optimization + loop_cloth = range(100) + + for i in loop_cloth: + + optimizer_cloth.zero_grad() + + self.render.load_meshes( + verts_pr.unsqueeze(0).to(self.device), + faces_pr.unsqueeze(0).to(self.device).long(), + deform_verts, + ) + P_normal_F, P_normal_B = self.render.get_rgb_image() + + update_mesh_shape_prior_losses(self.render.mesh, losses) + diff_F_cloth = torch.abs(P_normal_F[0] - inter[:3]) + diff_B_cloth = torch.abs(P_normal_B[0] - inter[3:]) + losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean() + losses["deform"]["value"] = torch.topk( + torch.abs(deform_verts.flatten()), 30 + )[0].mean() + + # Weighted sum of the losses + cloth_loss = torch.tensor(0.0, device=self.device) + pbar_desc = "" + + for k in losses.keys(): + if k != "smpl": + cloth_loss_per_cls = losses[k]["value"] * \ + losses[k]["weight"] + pbar_desc += f"{k}: {cloth_loss_per_cls:.3f} | " + cloth_loss += cloth_loss_per_cls + + # loop_cloth.set_description(pbar_desc) + cloth_loss.backward(retain_graph=True) + optimizer_cloth.step() + scheduler_cloth.step(cloth_loss) + + # convert from GT to SDF + deform_verts = deform_verts.flatten().detach() + deform_verts[torch.topk(torch.abs(deform_verts), 30)[ + 1]] = deform_verts.mean() + deform_verts = deform_verts.view(-1, 3).cpu() + + verts_pr += deform_verts + verts_pr *= (self.resolutions[-1] - 1) / 2.0 + verts_pr += (self.resolutions[-1] - 1) / 2.0 + + return verts_pr + + def test_step(self, batch, batch_idx): + + # dict_keys(['dataset', 'subject', 'rotation', 'scale', 'calib', + # 'normal_F', 'normal_B', 'image', 'T_normal_F', 'T_normal_B', + # 'z-trans', 'verts', 'faces', 'samples_geo', 'labels_geo', + # 'smpl_verts', 'smpl_faces', 'smpl_vis', 'smpl_cmap', 'pts_signs', + # 'type', 'gender', 'age', 'body_pose', 'global_orient', 'betas', 'transl']) + + if self.evaluator._normal_render is None: + self.evaluator.init_gl() + + self.netG.eval() + self.netG.training = False + in_tensor_dict = {} + + # export paths + mesh_name = batch["subject"][0] + mesh_rot = batch["rotation"][0].item() + ckpt_dir = self.cfg.name + + for kid, key in enumerate(self.cfg.dataset.noise_type): + ckpt_dir += f"_{key}_{self.cfg.dataset.noise_scale[kid]}" + + if self.cfg.optim_cloth: + ckpt_dir += "_optim_cloth" + if self.cfg.optim_body: + ckpt_dir += "_optim_body" + + self.export_dir = osp.join(self.cfg.results_path, ckpt_dir, mesh_name) + os.makedirs(self.export_dir, exist_ok=True) + + for name in self.in_total: + if name in batch.keys(): + in_tensor_dict.update({name: batch[name]}) + + # update the new T_normal_F/B + in_tensor_dict.update( + self.evaluator.render_normal( + batch["smpl_verts"], batch["smpl_faces"]) + ) + + # update the new smpl_vis + (xy, z) = batch["smpl_verts"][0].split([2, 1], dim=1) + smpl_vis = get_visibility( + xy, + z, + torch.as_tensor(self.smpl_data.faces).type_as( + batch["smpl_verts"]).long(), + ) + in_tensor_dict.update({"smpl_vis": smpl_vis.unsqueeze(0)}) + + if self.prior_type == "icon": + for key in self.icon_keys: + in_tensor_dict.update({key: batch[key]}) + elif self.prior_type == "pamir": + for key in self.pamir_keys: + in_tensor_dict.update({key: batch[key]}) + else: + pass + + with torch.no_grad(): + if self.cfg.optim_body: + features, inter, in_tensor_dict = self.optim_body( + in_tensor_dict, batch) + else: + features, inter = self.netG.filter( + in_tensor_dict, return_inter=True) + sdf = self.reconEngine( + opt=self.cfg, netG=self.netG, features=features, proj_matrix=None + ) + + # save inter results + image = ( + in_tensor_dict["image"][0].permute( + 1, 2, 0).detach().cpu().numpy() + 1.0 + ) * 0.5 + smpl_F = ( + in_tensor_dict["T_normal_F"][0].permute( + 1, 2, 0).detach().cpu().numpy() + + 1.0 + ) * 0.5 + smpl_B = ( + in_tensor_dict["T_normal_B"][0].permute( + 1, 2, 0).detach().cpu().numpy() + + 1.0 + ) * 0.5 + image_inter = np.concatenate( + self.tensor2image(512, inter[0]) + [smpl_F, smpl_B, image], axis=1 + ) + Image.fromarray((image_inter * 255.0).astype(np.uint8)).save( + osp.join(self.export_dir, f"{mesh_rot}_inter.png") + ) + + verts_pr, faces_pr = self.reconEngine.export_mesh(sdf) + + if self.clean_mesh_flag: + verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr) + + if self.cfg.optim_cloth: + verts_pr = self.optim_cloth(verts_pr, faces_pr, inter[0].detach()) + + verts_gt = batch["verts"][0] + faces_gt = batch["faces"][0] + + self.result_eval.update( + { + "verts_gt": verts_gt, + "faces_gt": faces_gt, + "verts_pr": verts_pr, + "faces_pr": faces_pr, + "recon_size": (self.resolutions[-1] - 1.0), + "calib": batch["calib"][0], + } + ) + + self.evaluator.set_mesh(self.result_eval, scale_factor=1.0) + self.evaluator.space_transfer() + + chamfer, p2s = self.evaluator.calculate_chamfer_p2s( + sampled_points=1000) + normal_consist = self.evaluator.calculate_normal_consist( + save_demo_img=osp.join(self.export_dir, f"{mesh_rot}_nc.png") + ) + + test_log = {"chamfer": chamfer, "p2s": p2s, "NC": normal_consist} + + return test_log + + def test_epoch_end(self, outputs): + + # make_test_gif("/".join(self.export_dir.split("/")[:-2])) + + accu_outputs = accumulate( + outputs, + rot_num=3, + split={ + "thuman2": (0, 5), + }, + ) + + print(colored(self.cfg.name, "green")) + print(colored(self.cfg.dataset.noise_scale, "green")) + + self.logger.experiment.add_hparams( + hparam_dict={"lr_G": self.lr_G, "bsize": self.batch_size}, + metric_dict=accu_outputs, + ) + + np.save( + osp.join(self.export_dir, "../test_results.npy"), + accu_outputs, + allow_pickle=True, + ) + + return accu_outputs + + def tensor2image(self, height, inter): + + all = [] + for dim in self.in_geo_dim: + img = resize( + np.tile( + ((inter[:dim].cpu().numpy() + 1.0) / + 2.0).transpose(1, 2, 0), + (1, 1, int(3 / dim)), + ), + (height, height), + anti_aliasing=True, + ) + + all.append(img) + inter = inter[dim:] + + return all + + def render_func(self, in_tensor_dict, dataset="title", idx=0): + + for name in in_tensor_dict.keys(): + in_tensor_dict[name] = in_tensor_dict[name][0:1] + + self.netG.eval() + features, inter = self.netG.filter(in_tensor_dict, return_inter=True) + sdf = self.reconEngine( + opt=self.cfg, netG=self.netG, features=features, proj_matrix=None + ) + + if sdf is not None: + render = self.reconEngine.display(sdf) + + image_pred = np.flip(render[:, :, ::-1], axis=0) + height = image_pred.shape[0] + + image_gt = resize( + ((in_tensor_dict["image"].cpu().numpy()[0] + 1.0) / 2.0).transpose( + 1, 2, 0 + ), + (height, height), + anti_aliasing=True, + ) + image_inter = self.tensor2image(height, inter[0]) + image = np.concatenate( + [image_pred, image_gt] + image_inter, axis=1) + + step_id = self.global_step if dataset == "train" else self.global_step + idx + self.logger.experiment.add_image( + tag=f"Occupancy-{dataset}/{step_id}", + img_tensor=image.transpose(2, 0, 1), + global_step=step_id, + ) + + def test_single(self, batch): + + self.netG.eval() + self.netG.training = False + in_tensor_dict = {} + + for name in self.in_total: + if name in batch.keys(): + in_tensor_dict.update({name: batch[name]}) + + if self.prior_type == "icon": + for key in self.icon_keys: + in_tensor_dict.update({key: batch[key]}) + elif self.prior_type == "pamir": + for key in self.pamir_keys: + in_tensor_dict.update({key: batch[key]}) + else: + pass + + features, inter = self.netG.filter(in_tensor_dict, return_inter=True) + sdf = self.reconEngine( + opt=self.cfg, netG=self.netG, features=features, proj_matrix=None + ) + + verts_pr, faces_pr = self.reconEngine.export_mesh(sdf) + + if self.clean_mesh_flag: + verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr) + + verts_pr -= (self.resolutions[-1] - 1) / 2.0 + verts_pr /= (self.resolutions[-1] - 1) / 2.0 + + return verts_pr, faces_pr, inter diff --git a/apps/Normal.py b/apps/Normal.py new file mode 100644 index 0000000000000000000000000000000000000000..9eed9ab6d7d24d9efadc563c0f921a101fe1c37c --- /dev/null +++ b/apps/Normal.py @@ -0,0 +1,213 @@ +from lib.net import NormalNet +from lib.common.train_util import * +import logging +import torch +import numpy as np +from torch import nn +from skimage.transform import resize +import pytorch_lightning as pl + +torch.backends.cudnn.benchmark = True + +logging.getLogger("lightning").setLevel(logging.ERROR) + + +class Normal(pl.LightningModule): + def __init__(self, cfg): + super(Normal, self).__init__() + self.cfg = cfg + self.batch_size = self.cfg.batch_size + self.lr_N = self.cfg.lr_N + + self.schedulers = [] + + self.netG = NormalNet(self.cfg, error_term=nn.SmoothL1Loss()) + + self.in_nml = [item[0] for item in cfg.net.in_nml] + + def get_progress_bar_dict(self): + tqdm_dict = super().get_progress_bar_dict() + if "v_num" in tqdm_dict: + del tqdm_dict["v_num"] + return tqdm_dict + + # Training related + def configure_optimizers(self): + + # set optimizer + weight_decay = self.cfg.weight_decay + momentum = self.cfg.momentum + + optim_params_N_F = [ + {"params": self.netG.netF.parameters(), "lr": self.lr_N}] + optim_params_N_B = [ + {"params": self.netG.netB.parameters(), "lr": self.lr_N}] + + optimizer_N_F = torch.optim.Adam( + optim_params_N_F, lr=self.lr_N, weight_decay=weight_decay + ) + + optimizer_N_B = torch.optim.Adam( + optim_params_N_B, lr=self.lr_N, weight_decay=weight_decay + ) + + scheduler_N_F = torch.optim.lr_scheduler.MultiStepLR( + optimizer_N_F, milestones=self.cfg.schedule, gamma=self.cfg.gamma + ) + + scheduler_N_B = torch.optim.lr_scheduler.MultiStepLR( + optimizer_N_B, milestones=self.cfg.schedule, gamma=self.cfg.gamma + ) + + self.schedulers = [scheduler_N_F, scheduler_N_B] + optims = [optimizer_N_F, optimizer_N_B] + + return optims, self.schedulers + + def render_func(self, render_tensor): + + height = render_tensor["image"].shape[2] + result_list = [] + + for name in render_tensor.keys(): + result_list.append( + resize( + ((render_tensor[name].cpu().numpy()[0] + 1.0) / 2.0).transpose( + 1, 2, 0 + ), + (height, height), + anti_aliasing=True, + ) + ) + result_array = np.concatenate(result_list, axis=1) + + return result_array + + def training_step(self, batch, batch_idx, optimizer_idx): + + export_cfg(self.logger, self.cfg) + + # retrieve the data + in_tensor = {} + for name in self.in_nml: + in_tensor[name] = batch[name] + + FB_tensor = {"normal_F": batch["normal_F"], + "normal_B": batch["normal_B"]} + + self.netG.train() + + preds_F, preds_B = self.netG(in_tensor) + error_NF, error_NB = self.netG.get_norm_error( + preds_F, preds_B, FB_tensor) + + (opt_nf, opt_nb) = self.optimizers() + + opt_nf.zero_grad() + opt_nb.zero_grad() + + self.manual_backward(error_NF, opt_nf) + self.manual_backward(error_NB, opt_nb) + + opt_nf.step() + opt_nb.step() + + if batch_idx > 0 and batch_idx % int(self.cfg.freq_show_train) == 0: + + self.netG.eval() + with torch.no_grad(): + nmlF, nmlB = self.netG(in_tensor) + in_tensor.update({"nmlF": nmlF, "nmlB": nmlB}) + result_array = self.render_func(in_tensor) + + self.logger.experiment.add_image( + tag=f"Normal-train/{self.global_step}", + img_tensor=result_array.transpose(2, 0, 1), + global_step=self.global_step, + ) + + # metrics processing + metrics_log = { + "train_loss-NF": error_NF.item(), + "train_loss-NB": error_NB.item(), + } + + tf_log = tf_log_convert(metrics_log) + bar_log = bar_log_convert(metrics_log) + + return { + "loss": error_NF + error_NB, + "loss-NF": error_NF, + "loss-NB": error_NB, + "log": tf_log, + "progress_bar": bar_log, + } + + def training_epoch_end(self, outputs): + + if [] in outputs: + outputs = outputs[0] + + # metrics processing + metrics_log = { + "train_avgloss": batch_mean(outputs, "loss"), + "train_avgloss-NF": batch_mean(outputs, "loss-NF"), + "train_avgloss-NB": batch_mean(outputs, "loss-NB"), + } + + tf_log = tf_log_convert(metrics_log) + + tf_log["lr-NF"] = self.schedulers[0].get_last_lr()[0] + tf_log["lr-NB"] = self.schedulers[1].get_last_lr()[0] + + return {"log": tf_log} + + def validation_step(self, batch, batch_idx): + + # retrieve the data + in_tensor = {} + for name in self.in_nml: + in_tensor[name] = batch[name] + + FB_tensor = {"normal_F": batch["normal_F"], + "normal_B": batch["normal_B"]} + + self.netG.train() + + preds_F, preds_B = self.netG(in_tensor) + error_NF, error_NB = self.netG.get_norm_error( + preds_F, preds_B, FB_tensor) + + if (batch_idx > 0 and batch_idx % int(self.cfg.freq_show_train) == 0) or ( + batch_idx == 0 + ): + + with torch.no_grad(): + nmlF, nmlB = self.netG(in_tensor) + in_tensor.update({"nmlF": nmlF, "nmlB": nmlB}) + result_array = self.render_func(in_tensor) + + self.logger.experiment.add_image( + tag=f"Normal-val/{self.global_step}", + img_tensor=result_array.transpose(2, 0, 1), + global_step=self.global_step, + ) + + return { + "val_loss": error_NF + error_NB, + "val_loss-NF": error_NF, + "val_loss-NB": error_NB, + } + + def validation_epoch_end(self, outputs): + + # metrics processing + metrics_log = { + "val_avgloss": batch_mean(outputs, "val_loss"), + "val_avgloss-NF": batch_mean(outputs, "val_loss-NF"), + "val_avgloss-NB": batch_mean(outputs, "val_loss-NB"), + } + + tf_log = tf_log_convert(metrics_log) + + return {"log": tf_log} diff --git a/apps/__pycache__/app.cpython-38.pyc b/apps/__pycache__/app.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdbe6b9262dbf510f32e2596fd209c0457b938d0 Binary files /dev/null and b/apps/__pycache__/app.cpython-38.pyc differ diff --git a/apps/app.py b/apps/app.py new file mode 100644 index 0000000000000000000000000000000000000000..15dc741ea1cc0c4723ca202bd99c370a4d4a9af0 --- /dev/null +++ b/apps/app.py @@ -0,0 +1,21 @@ +# install + +import os +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"]="0" +try: + os.system("bash install.sh") +except Exception as e: + print(e) + + +# running + +import gradio as gr + +def image_classifier(inp): + return {'cat': 0.3, 'dog': 0.7} + +demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label") +demo.launch(auth=("icon@tue.mpg.de", "icon_2022"), + auth_message="Register at icon.is.tue.mpg.de/download to get the username and password.") diff --git a/apps/infer.py b/apps/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..ab34478cbfbdbb24cd47ddb77c35facb760d61f1 --- /dev/null +++ b/apps/infer.py @@ -0,0 +1,616 @@ +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import logging +from lib.common.render import query_color, image2vid +from lib.common.config import cfg +from lib.common.cloth_extraction import extract_cloth +from lib.dataset.mesh_util import ( + load_checkpoint, + update_mesh_shape_prior_losses, + get_optim_grid_image, + blend_rgb_norm, + unwrap, + remesh, + tensor2variable, + normal_loss +) + +from lib.dataset.TestDataset import TestDataset +from lib.net.local_affine import LocalAffine +from pytorch3d.structures import Meshes +from apps.ICON import ICON + +import os +from termcolor import colored +import argparse +import numpy as np +from PIL import Image +import trimesh +import pickle +import numpy as np + +import torch +torch.backends.cudnn.benchmark = True + +logging.getLogger("trimesh").setLevel(logging.ERROR) + + +if __name__ == "__main__": + + # loading cfg file + parser = argparse.ArgumentParser() + + parser.add_argument("-gpu", "--gpu_device", type=int, default=0) + parser.add_argument("-colab", action="store_true") + parser.add_argument("-loop_smpl", "--loop_smpl", type=int, default=100) + parser.add_argument("-patience", "--patience", type=int, default=5) + parser.add_argument("-vis_freq", "--vis_freq", type=int, default=10) + parser.add_argument("-loop_cloth", "--loop_cloth", type=int, default=200) + parser.add_argument("-hps_type", "--hps_type", type=str, default="pymaf") + parser.add_argument("-export_video", action="store_true") + parser.add_argument("-in_dir", "--in_dir", type=str, default="./examples") + parser.add_argument("-out_dir", "--out_dir", + type=str, default="./results") + parser.add_argument('-seg_dir', '--seg_dir', type=str, default=None) + parser.add_argument( + "-cfg", "--config", type=str, default="./configs/icon-filter.yaml" + ) + + args = parser.parse_args() + + # cfg read and merge + cfg.merge_from_file(args.config) + cfg.merge_from_file("./lib/pymaf/configs/pymaf_config.yaml") + + cfg_show_list = [ + "test_gpus", + [args.gpu_device], + "mcube_res", + 256, + "clean_mesh", + True, + ] + + cfg.merge_from_list(cfg_show_list) + cfg.freeze() + + os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" + device = torch.device(f"cuda:{args.gpu_device}") + + if args.colab: + print(colored("colab environment...", "red")) + from tqdm.notebook import tqdm + else: + print(colored("normal environment...", "red")) + from tqdm import tqdm + + # load model and dataloader + model = ICON(cfg) + model = load_checkpoint(model, cfg) + + dataset_param = { + 'image_dir': args.in_dir, + 'seg_dir': args.seg_dir, + 'has_det': True, # w/ or w/o detection + 'hps_type': args.hps_type # pymaf/pare/pixie + } + + if args.hps_type == "pixie" and "pamir" in args.config: + print(colored("PIXIE isn't compatible with PaMIR, thus switch to PyMAF", "red")) + dataset_param["hps_type"] = "pymaf" + + dataset = TestDataset(dataset_param, device) + + print(colored(f"Dataset Size: {len(dataset)}", "green")) + + pbar = tqdm(dataset) + + for data in pbar: + + pbar.set_description(f"{data['name']}") + + in_tensor = {"smpl_faces": data["smpl_faces"], "image": data["image"]} + + # The optimizer and variables + optimed_pose = torch.tensor( + data["body_pose"], device=device, requires_grad=True + ) # [1,23,3,3] + optimed_trans = torch.tensor( + data["trans"], device=device, requires_grad=True + ) # [3] + optimed_betas = torch.tensor( + data["betas"], device=device, requires_grad=True + ) # [1,10] + optimed_orient = torch.tensor( + data["global_orient"], device=device, requires_grad=True + ) # [1,1,3,3] + + optimizer_smpl = torch.optim.SGD( + [optimed_pose, optimed_trans, optimed_betas, optimed_orient], + lr=1e-3, + momentum=0.9, + ) + scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer_smpl, + mode="min", + factor=0.5, + verbose=0, + min_lr=1e-5, + patience=args.patience, + ) + + losses = { + "cloth": {"weight": 1e1, "value": 0.0}, # Cloth: Normal_recon - Normal_pred + "stiffness": {"weight": 1e5, "value": 0.0}, # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2) + "rigid": {"weight": 1e5, "value": 0.0}, # Cloth: det(R) = 1 + "edge": {"weight": 0, "value": 0.0}, # Cloth: edge length + "nc": {"weight": 0, "value": 0.0}, # Cloth: normal consistency + "laplacian": {"weight": 1e2, "value": 0.0}, # Cloth: laplacian smoonth + "normal": {"weight": 1e0, "value": 0.0}, # Body: Normal_pred - Normal_smpl + "silhouette": {"weight": 1e1, "value": 0.0}, # Body: Silhouette_pred - Silhouette_smpl + } + + # smpl optimization + + loop_smpl = tqdm( + range(args.loop_smpl if cfg.net.prior_type != "pifu" else 1)) + + per_data_lst = [] + + for i in loop_smpl: + + per_loop_lst = [] + + optimizer_smpl.zero_grad() + + if dataset_param["hps_type"] != "pixie": + smpl_out = dataset.smpl_model( + betas=optimed_betas, + body_pose=optimed_pose, + global_orient=optimed_orient, + pose2rot=False, + ) + + smpl_verts = ((smpl_out.vertices) + + optimed_trans) * data["scale"] + else: + smpl_verts, _, _ = dataset.smpl_model( + shape_params=optimed_betas, + expression_params=tensor2variable(data["exp"], device), + body_pose=optimed_pose, + global_pose=optimed_orient, + jaw_pose=tensor2variable(data["jaw_pose"], device), + left_hand_pose=tensor2variable( + data["left_hand_pose"], device), + right_hand_pose=tensor2variable( + data["right_hand_pose"], device), + ) + + smpl_verts = (smpl_verts + optimed_trans) * data["scale"] + + # render optimized mesh (normal, T_normal, image [-1,1]) + in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal( + smpl_verts * + torch.tensor([1.0, -1.0, -1.0] + ).to(device), in_tensor["smpl_faces"] + ) + T_mask_F, T_mask_B = dataset.render.get_silhouette_image() + + with torch.no_grad(): + in_tensor["normal_F"], in_tensor["normal_B"] = model.netG.normal_filter( + in_tensor + ) + + diff_F_smpl = torch.abs( + in_tensor["T_normal_F"] - in_tensor["normal_F"]) + diff_B_smpl = torch.abs( + in_tensor["T_normal_B"] - in_tensor["normal_B"]) + + loss_F_smpl = normal_loss( + in_tensor["T_normal_F"], in_tensor["normal_F"]) + loss_B_smpl = normal_loss( + in_tensor["T_normal_B"], in_tensor["normal_B"]) + + losses["normal"]["value"] = (loss_F_smpl + loss_B_smpl).mean() + + # silhouette loss + smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0] + gt_arr = torch.cat( + [in_tensor["normal_F"][0], in_tensor["normal_B"][0]], dim=2 + ).permute(1, 2, 0) + gt_arr = ((gt_arr + 1.0) * 0.5).to(device) + bg_color = ( + torch.Tensor([0.5, 0.5, 0.5]).unsqueeze( + 0).unsqueeze(0).to(device) + ) + gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float() + diff_S = torch.abs(smpl_arr - gt_arr) + losses["silhouette"]["value"] = diff_S.mean() + + # Weighted sum of the losses + smpl_loss = 0.0 + pbar_desc = "Body Fitting --- " + for k in ["normal", "silhouette"]: + pbar_desc += f"{k}: {losses[k]['value'] * losses[k]['weight']:.3f} | " + smpl_loss += losses[k]["value"] * losses[k]["weight"] + pbar_desc += f"Total: {smpl_loss:.3f}" + loop_smpl.set_description(pbar_desc) + + if i % args.vis_freq == 0: + + per_loop_lst.extend( + [ + in_tensor["image"], + in_tensor["T_normal_F"], + in_tensor["normal_F"], + diff_F_smpl / 2.0, + diff_S[:, :512].unsqueeze( + 0).unsqueeze(0).repeat(1, 3, 1, 1), + ] + ) + per_loop_lst.extend( + [ + in_tensor["image"], + in_tensor["T_normal_B"], + in_tensor["normal_B"], + diff_B_smpl / 2.0, + diff_S[:, 512:].unsqueeze( + 0).unsqueeze(0).repeat(1, 3, 1, 1), + ] + ) + per_data_lst.append( + get_optim_grid_image( + per_loop_lst, None, nrow=5, type="smpl") + ) + + smpl_loss.backward() + optimizer_smpl.step() + scheduler_smpl.step(smpl_loss) + in_tensor["smpl_verts"] = smpl_verts * \ + torch.tensor([1.0, 1.0, -1.0]).to(device) + + # visualize the optimization process + # 1. SMPL Fitting + # 2. Clothes Refinement + + os.makedirs(os.path.join(args.out_dir, cfg.name, + "refinement"), exist_ok=True) + + # visualize the final results in self-rotation mode + os.makedirs(os.path.join(args.out_dir, cfg.name, "vid"), exist_ok=True) + + # final results rendered as image + # 1. Render the final fitted SMPL (xxx_smpl.png) + # 2. Render the final reconstructed clothed human (xxx_cloth.png) + # 3. Blend the original image with predicted cloth normal (xxx_overlap.png) + + os.makedirs(os.path.join(args.out_dir, cfg.name, "png"), exist_ok=True) + + # final reconstruction meshes + # 1. SMPL mesh (xxx_smpl.obj) + # 2. SMPL params (xxx_smpl.npy) + # 3. clohted mesh (xxx_recon.obj) + # 4. remeshed clothed mesh (xxx_remesh.obj) + # 5. refined clothed mesh (xxx_refine.obj) + + os.makedirs(os.path.join(args.out_dir, cfg.name, "obj"), exist_ok=True) + + if cfg.net.prior_type != "pifu": + + per_data_lst[0].save( + os.path.join( + args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.gif" + ), + save_all=True, + append_images=per_data_lst[1:], + duration=500, + loop=0, + ) + + if args.vis_freq == 1: + image2vid( + per_data_lst, + os.path.join( + args.out_dir, cfg.name, f"refinement/{data['name']}_smpl.avi" + ), + ) + + per_data_lst[-1].save( + os.path.join(args.out_dir, cfg.name, + f"png/{data['name']}_smpl.png") + ) + + norm_pred = ( + ((in_tensor["normal_F"][0].permute(1, 2, 0) + 1.0) * 255.0 / 2.0) + .detach() + .cpu() + .numpy() + .astype(np.uint8) + ) + + norm_orig = unwrap(norm_pred, data) + mask_orig = unwrap( + np.repeat( + data["mask"].permute(1, 2, 0).detach().cpu().numpy(), 3, axis=2 + ).astype(np.uint8), + data, + ) + rgb_norm = blend_rgb_norm(data["ori_image"], norm_orig, mask_orig) + + Image.fromarray( + np.concatenate( + [data["ori_image"].astype(np.uint8), rgb_norm], axis=1) + ).save(os.path.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png")) + + smpl_obj = trimesh.Trimesh( + in_tensor["smpl_verts"].detach().cpu()[0] * + torch.tensor([1.0, -1.0, 1.0]), + in_tensor['smpl_faces'].detach().cpu()[0], + process=False, + maintains_order=True + ) + smpl_obj.export( + f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.obj") + + smpl_info = {'betas': optimed_betas, + 'pose': optimed_pose, + 'orient': optimed_orient, + 'trans': optimed_trans} + + np.save( + f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl.npy", smpl_info, allow_pickle=True) + + # ------------------------------------------------------------------------------------------------------------------ + + # cloth optimization + + per_data_lst = [] + + # cloth recon + in_tensor.update( + dataset.compute_vis_cmap( + in_tensor["smpl_verts"][0], in_tensor["smpl_faces"][0] + ) + ) + + if cfg.net.prior_type == "pamir": + in_tensor.update( + dataset.compute_voxel_verts( + optimed_pose, + optimed_orient, + optimed_betas, + optimed_trans, + data["scale"], + ) + ) + + with torch.no_grad(): + verts_pr, faces_pr, _ = model.test_single(in_tensor) + + recon_obj = trimesh.Trimesh( + verts_pr, faces_pr, process=False, maintains_order=True + ) + recon_obj.export( + os.path.join(args.out_dir, cfg.name, + f"obj/{data['name']}_recon.obj") + ) + + # Isotropic Explicit Remeshing for better geometry topology + verts_refine, faces_refine = remesh(os.path.join(args.out_dir, cfg.name, + f"obj/{data['name']}_recon.obj"), 0.5, device) + + # define local_affine deform verts + mesh_pr = Meshes(verts_refine, faces_refine).to(device) + local_affine_model = LocalAffine( + mesh_pr.verts_padded().shape[1], mesh_pr.verts_padded().shape[0], mesh_pr.edges_packed()).to(device) + optimizer_cloth = torch.optim.Adam( + [{'params': local_affine_model.parameters()}], lr=1e-4, amsgrad=True) + + scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer_cloth, + mode="min", + factor=0.1, + verbose=0, + min_lr=1e-5, + patience=args.patience, + ) + + with torch.no_grad(): + per_loop_lst = [] + rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[ + 0, 1, 2, 3]) + per_loop_lst.extend(rotate_recon_lst) + per_data_lst.append(get_optim_grid_image( + per_loop_lst, None, type="cloth")) + + final = None + + if args.loop_cloth > 0: + + loop_cloth = tqdm(range(args.loop_cloth)) + + for i in loop_cloth: + + per_loop_lst = [] + + optimizer_cloth.zero_grad() + + deformed_verts, stiffness, rigid = local_affine_model( + verts_refine.to(device), return_stiff=True) + mesh_pr = mesh_pr.update_padded(deformed_verts) + + # losses for laplacian, edge, normal consistency + update_mesh_shape_prior_losses(mesh_pr, losses) + + in_tensor["P_normal_F"], in_tensor["P_normal_B"] = dataset.render_normal( + mesh_pr.verts_padded(), mesh_pr.faces_padded()) + + diff_F_cloth = torch.abs( + in_tensor["P_normal_F"] - in_tensor["normal_F"]) + diff_B_cloth = torch.abs( + in_tensor["P_normal_B"] - in_tensor["normal_B"]) + + losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean() + losses["stiffness"]["value"] = torch.mean(stiffness) + losses["rigid"]["value"] = torch.mean(rigid) + + # Weighted sum of the losses + cloth_loss = torch.tensor(0.0, requires_grad=True).to(device) + pbar_desc = "Cloth Refinement --- " + + for k in losses.keys(): + if k not in ["normal", "silhouette"] and losses[k]["weight"] > 0.0: + cloth_loss = cloth_loss + \ + losses[k]["value"] * losses[k]["weight"] + pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.5f} | " + + pbar_desc += f"Total: {cloth_loss:.5f}" + loop_cloth.set_description(pbar_desc) + + # update params + cloth_loss.backward(retain_graph=True) + optimizer_cloth.step() + scheduler_cloth.step(cloth_loss) + + # for vis + with torch.no_grad(): + if i % args.vis_freq == 0: + + rotate_recon_lst = dataset.render.get_rgb_image(cam_ids=[ + 0, 1, 2, 3]) + + per_loop_lst.extend( + [ + in_tensor["image"], + in_tensor["P_normal_F"], + in_tensor["normal_F"], + diff_F_cloth / 2.0, + ] + ) + per_loop_lst.extend( + [ + in_tensor["image"], + in_tensor["P_normal_B"], + in_tensor["normal_B"], + diff_B_cloth / 2.0, + ] + ) + per_loop_lst.extend(rotate_recon_lst) + per_data_lst.append( + get_optim_grid_image( + per_loop_lst, None, type="cloth") + ) + + # gif for optimization + per_data_lst[1].save( + os.path.join( + args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.gif" + ), + save_all=True, + append_images=per_data_lst[2:], + duration=500, + loop=0, + ) + + if args.vis_freq == 1: + image2vid( + per_data_lst, + os.path.join( + args.out_dir, cfg.name, f"refinement/{data['name']}_cloth.avi" + ), + ) + + final = trimesh.Trimesh( + mesh_pr.verts_packed().detach().squeeze(0).cpu(), + mesh_pr.faces_packed().detach().squeeze(0).cpu(), + process=False, maintains_order=True + ) + final_colors = query_color( + mesh_pr.verts_packed().detach().squeeze(0).cpu(), + mesh_pr.faces_packed().detach().squeeze(0).cpu(), + in_tensor["image"], + device=device, + ) + final.visual.vertex_colors = final_colors + final.export( + f"{args.out_dir}/{cfg.name}/obj/{data['name']}_refine.obj") + + # always export visualized png regardless of the cloth refinment + per_data_lst[-1].save( + os.path.join(args.out_dir, cfg.name, + f"png/{data['name']}_cloth.png") + ) + + # always export visualized video regardless of the cloth refinment + if args.export_video: + if final is not None: + verts_lst = [verts_pr, final.vertices] + faces_lst = [faces_pr, final.faces] + else: + verts_lst = [verts_pr] + faces_lst = [faces_pr] + + # self-rotated video + dataset.render.load_meshes( + verts_lst, faces_lst) + dataset.render.get_rendered_video( + [data["ori_image"], rgb_norm], + os.path.join(args.out_dir, cfg.name, + f"vid/{data['name']}_cloth.mp4"), + ) + + # garment extraction from deepfashion images + if not (args.seg_dir is None): + if final is not None: + recon_obj = final.copy() + + os.makedirs(os.path.join( + args.out_dir, cfg.name, "clothes"), exist_ok=True) + os.makedirs(os.path.join(args.out_dir, cfg.name, + "clothes", "info"), exist_ok=True) + for seg in data['segmentations']: + # These matrices work for PyMaf, not sure about the other hps type + K = np.array([[1.0000, 0.0000, 0.0000, 0.0000], + [0.0000, 1.0000, 0.0000, 0.0000], + [0.0000, 0.0000, -0.5000, 0.0000], + [-0.0000, -0.0000, 0.5000, 1.0000]]).T + + R = np.array([[-1., 0., 0.], + [0., 1., 0.], + [0., 0., -1.]]) + + t = np.array([[-0., -0., 100.]]) + clothing_obj = extract_cloth(recon_obj, seg, K, R, t, smpl_obj) + if clothing_obj is not None: + cloth_type = seg['type'].replace(' ', '_') + cloth_info = { + 'betas': optimed_betas, + 'body_pose': optimed_pose, + 'global_orient': optimed_orient, + 'pose2rot': False, + 'clothing_type': cloth_type, + } + + file_id = f"{data['name']}_{cloth_type}" + with open(os.path.join(args.out_dir, cfg.name, "clothes", "info", f"{file_id}_info.pkl"), 'wb') as fp: + pickle.dump(cloth_info, fp) + + clothing_obj.export(os.path.join( + args.out_dir, cfg.name, "clothes", f"{file_id}.obj")) + else: + print( + f"Unable to extract clothing of type {seg['type']} from image {data['name']}") diff --git a/assets/garment_teaser.png b/assets/garment_teaser.png new file mode 100644 index 0000000000000000000000000000000000000000..95200eae30f68d1a16797bf9f95065d0a885d2b5 Binary files /dev/null and b/assets/garment_teaser.png differ diff --git a/assets/intermediate_results.png b/assets/intermediate_results.png new file mode 100644 index 0000000000000000000000000000000000000000..e1ca1303d711b5f5c728fc45aa33fca3994be4bc Binary files /dev/null and b/assets/intermediate_results.png differ diff --git a/assets/teaser.gif b/assets/teaser.gif new file mode 100644 index 0000000000000000000000000000000000000000..30fb76214a17d73fd0760f43430e4fb933c39176 Binary files /dev/null and b/assets/teaser.gif differ diff --git a/configs/icon-filter.yaml b/configs/icon-filter.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a8131461ddd0ebd6cb8ab4e3e2d2301517c88bf1 --- /dev/null +++ b/configs/icon-filter.yaml @@ -0,0 +1,25 @@ +name: icon-filter +ckpt_dir: "./data/ckpt/" +resume_path: "./data/ckpt/icon-filter.ckpt" +normal_path: "./data/ckpt/normal.ckpt" + +test_mode: True +batch_size: 1 + +net: + mlp_dim: [256, 512, 256, 128, 1] + res_layers: [2,3,4] + num_stack: 2 + prior_type: "icon" # icon/pamir/icon + use_filter: True + in_geo: (('normal_F',3), ('normal_B',3)) + in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3)) + smpl_feats: ['sdf', 'norm', 'vis', 'cmap'] + gtype: 'HGPIFuNet' + norm_mlp: 'batch' + hourglass_dim: 6 + smpl_dim: 7 + +# user defined +mcube_res: 512 # occupancy field resolution, higher --> more details +clean_mesh: False # if True, will remove floating pieces \ No newline at end of file diff --git a/configs/icon-nofilter.yaml b/configs/icon-nofilter.yaml new file mode 100644 index 0000000000000000000000000000000000000000..506b9b22a1e97a5bb99ce2ba10662de500dfb147 --- /dev/null +++ b/configs/icon-nofilter.yaml @@ -0,0 +1,25 @@ +name: icon-nofilter +ckpt_dir: "./data/ckpt/" +resume_path: "./data/ckpt/icon-nofilter.ckpt" +normal_path: "./data/ckpt/normal.ckpt" + +test_mode: True +batch_size: 1 + +net: + mlp_dim: [256, 512, 256, 128, 1] + res_layers: [2,3,4] + num_stack: 2 + prior_type: "icon" # icon/pamir/icon + use_filter: False + in_geo: (('normal_F',3), ('normal_B',3)) + in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3)) + smpl_feats: ['sdf', 'norm', 'vis', 'cmap'] + gtype: 'HGPIFuNet' + norm_mlp: 'batch' + hourglass_dim: 6 + smpl_dim: 7 + +# user defined +mcube_res: 512 # occupancy field resolution, higher --> more details +clean_mesh: False # if True, will remove floating pieces diff --git a/configs/pamir.yaml b/configs/pamir.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a06e96779a6b3bd1c0e0f1e53f75cc609c380c8b --- /dev/null +++ b/configs/pamir.yaml @@ -0,0 +1,24 @@ +name: pamir +ckpt_dir: "./data/ckpt/" +resume_path: "./data/ckpt/pamir.ckpt" +normal_path: "./data/ckpt/normal.ckpt" + +test_mode: True +batch_size: 1 + +net: + mlp_dim: [256, 512, 256, 128, 1] + res_layers: [2,3,4] + num_stack: 2 + prior_type: "pamir" # icon/pamir/icon + use_filter: True + in_geo: (('image',3), ('normal_F',3), ('normal_B',3)) + in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3)) + gtype: 'HGPIFuNet' + norm_mlp: 'batch' + hourglass_dim: 6 + voxel_dim: 7 + +# user defined +mcube_res: 512 # occupancy field resolution, higher --> more details +clean_mesh: False # if True, will remove floating pieces \ No newline at end of file diff --git a/configs/pifu.yaml b/configs/pifu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f1d0d23929c7d69dcc121dd87484eeb73d732364 --- /dev/null +++ b/configs/pifu.yaml @@ -0,0 +1,24 @@ +name: pifu +ckpt_dir: "./data/ckpt/" +resume_path: "./data/ckpt/pifu.ckpt" +normal_path: "./data/ckpt/normal.ckpt" + +test_mode: True +batch_size: 1 + +net: + mlp_dim: [256, 512, 256, 128, 1] + res_layers: [2,3,4] + num_stack: 2 + prior_type: "pifu" # icon/pamir/icon + use_filter: True + in_geo: (('image',3), ('normal_F',3), ('normal_B',3)) + in_nml: (('image',3), ('T_normal_F',3), ('T_normal_B',3)) + gtype: 'HGPIFuNet' + norm_mlp: 'batch' + hourglass_dim: 12 + + +# user defined +mcube_res: 512 # occupancy field resolution, higher --> more details +clean_mesh: False # if True, will remove floating pieces \ No newline at end of file diff --git a/environment.yaml b/environment.yaml new file mode 100644 index 0000000000000000000000000000000000000000..23cb949f3c160bef04bccca04456384c3ded0256 --- /dev/null +++ b/environment.yaml @@ -0,0 +1,16 @@ +name: icon +channels: + - pytorch-lts + - nvidia + - conda-forge + - fvcore + - iopath + - bottler + - defaults +dependencies: + - pytorch + - torchvision + - fvcore + - iopath + - nvidiacub + - pyembree \ No newline at end of file diff --git a/examples/22097467bffc92d4a5c4246f7d4edb75.png b/examples/22097467bffc92d4a5c4246f7d4edb75.png new file mode 100644 index 0000000000000000000000000000000000000000..7139dd8327b4cc0f2e152fb6c481baefd9b3feb0 Binary files /dev/null and b/examples/22097467bffc92d4a5c4246f7d4edb75.png differ diff --git a/examples/44c0f84c957b6b9bdf77662af5bb7078.png b/examples/44c0f84c957b6b9bdf77662af5bb7078.png new file mode 100644 index 0000000000000000000000000000000000000000..c0b7006e8d022dd0b7e831123e00386ae7eabb6a Binary files /dev/null and b/examples/44c0f84c957b6b9bdf77662af5bb7078.png differ diff --git a/examples/5a6a25963db2f667441d5076972c207c.png b/examples/5a6a25963db2f667441d5076972c207c.png new file mode 100644 index 0000000000000000000000000000000000000000..035a0f74ba5f8377c3f44ed1d5f3bc16719e17d9 Binary files /dev/null and b/examples/5a6a25963db2f667441d5076972c207c.png differ diff --git a/examples/8da7ceb94669c2f65cbd28022e1f9876.png b/examples/8da7ceb94669c2f65cbd28022e1f9876.png new file mode 100644 index 0000000000000000000000000000000000000000..ef8f0d8a7221334f87c28feb1af2cbdb27fb3c9a Binary files /dev/null and b/examples/8da7ceb94669c2f65cbd28022e1f9876.png differ diff --git a/examples/923d65f767c85a42212cae13fba3750b.png b/examples/923d65f767c85a42212cae13fba3750b.png new file mode 100644 index 0000000000000000000000000000000000000000..60d12587cd7ee1a227dc6e8e4769e81efa21af26 Binary files /dev/null and b/examples/923d65f767c85a42212cae13fba3750b.png differ diff --git a/examples/959c4c726a69901ce71b93a9242ed900.png b/examples/959c4c726a69901ce71b93a9242ed900.png new file mode 100644 index 0000000000000000000000000000000000000000..adfb527b5b1cf4a73f4cc75a64646ae7a3791705 Binary files /dev/null and b/examples/959c4c726a69901ce71b93a9242ed900.png differ diff --git a/examples/c9856a2bc31846d684cbb965457fad59.png b/examples/c9856a2bc31846d684cbb965457fad59.png new file mode 100644 index 0000000000000000000000000000000000000000..ff8bc4d132348fb4d8961091300165673998e680 Binary files /dev/null and b/examples/c9856a2bc31846d684cbb965457fad59.png differ diff --git a/examples/e1e7622af7074a022f5d96dc16672517.png b/examples/e1e7622af7074a022f5d96dc16672517.png new file mode 100644 index 0000000000000000000000000000000000000000..a38826046400c3c3cbe013e834ba2a37af1b4a3e Binary files /dev/null and b/examples/e1e7622af7074a022f5d96dc16672517.png differ diff --git a/examples/fb9d20fdb93750584390599478ecf86e.png b/examples/fb9d20fdb93750584390599478ecf86e.png new file mode 100644 index 0000000000000000000000000000000000000000..1150008a65b386c876410d7a0c7b2a655bebbe24 Binary files /dev/null and b/examples/fb9d20fdb93750584390599478ecf86e.png differ diff --git a/examples/segmentation/003883.jpg b/examples/segmentation/003883.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9e9d08a207946883f2ba92f608422d236b3cc51f Binary files /dev/null and b/examples/segmentation/003883.jpg differ diff --git a/examples/segmentation/003883.json b/examples/segmentation/003883.json new file mode 100644 index 0000000000000000000000000000000000000000..438947e9d91bb578bc8344769aea65195d0eeb2d --- /dev/null +++ b/examples/segmentation/003883.json @@ -0,0 +1,136 @@ +{ + "item2": { + "segmentation": [ + [ + 232.29572649572654, 34.447388414055126, 237.0364672364673, + 40.57084520417861, 244.9377018043686, 47.089363722697165, + 252.04881291547974, 49.65726495726508, 262.5179487179489, + 51.43504273504287, 269.233998100665, 50.447388414055204, + 277.5446343779678, 49.12725546058881, 285.64339981006657, + 46.16429249762584, 294.9273504273506, 41.22602089268754, + 299.9377967711301, 36.514245014245084, 304.67853751187084, + 30.588319088319132, 306.0612535612536, 25.65004748338083, + 307.64150047483383, 23.477207977207982, 311.19705603038943, + 24.859924026590704, 317.12298195631536, 28.020417853751216, + 323.04890788224134, 29.008072174738874, 331.34520417853764, + 30.193257359924065, 339.4439696106365, 34.7364672364673, + 346.75261158594515, 39.279677113010536, 350.11063627730323, + 44.61301044634389, 355.00541310541314, 61.422317188983875, + 358.9560303893638, 77.6198480531815, 362.1165242165243, + 90.26182336182353, 364.88195631528976, 103.29886039886063, + 367.6473884140552, 118.11367521367552, 369.42516619183294, + 129.37293447293484, 369.2324786324788, 132.60550807217476, + 365.6769230769232, 134.77834757834762, 359.15840455840464, + 138.3339031339032, 353.43000949667623, 140.70427350427357, + 351.4547008547009, 141.4943969610637, 351.25716999050337, + 138.5314339981007, 351.05963912630585, 136.75365622032294, + 345.7263057929725, 137.34624881291552, 337.8250712250712, + 139.51908831908838, 331.5040835707502, 141.09933523266864, + 324.7880341880341, 143.66723646723653, 322.2201329534662, + 146.43266856600198, 322.2201329534662, 151.5684710351378, + 323.0102564102563, 160.6548907882243, 324.95185185185176, + 173.44615384615395, 325.34691358024685, 190.23627730294416, + 325.93950617283946, 205.64368471035164, 325.93950617283946, + 215.71775878442577, 325.93950617283946, 220.06343779677147, + 322.7790123456789, 223.22393162393197, 315.0753086419752, + 228.55726495726532, 309.34691358024673, 230.53257359924066, + 290.1866096866098, 230.87929724596398, 263.91500474833805, + 229.6941120607788, 236.45821462488112, 229.29905033238373, + 218.48290598290572, 226.73114909781583, 202.65650522317188, + 224.82811016144353, 197.71823361823357, 221.07502374169044, + 195.15033238366567, 214.55650522317188, 195.74292497625825, + 200.53181386514711, 197.125641025641, 180.5811965811964, + 197.33285849952523, 164.68736942070285, 198.51804368471042, + 154.21823361823365, 198.51804368471042, 138.61329534662863, + 193.5797720797721, 136.4404558404558, 185.08594491927823, + 133.08243114909774, 177.77730294396957, 128.73675213675205, + 174.41927825261152, 128.53922127255453, 173.82668566001894, + 133.2799620132953, 174.02421652421646, 136.24292497625825, + 172.83903133903127, 137.03304843304838, 167.11063627730283, + 134.86020892687554, 159.9995251661917, 130.51452991452985, + 159.01187084520404, 129.1318138651471, 159.60446343779662, + 123.60094966761622, 162.6012345679013, 111.57578347578357, + 165.95925925925934, 98.53874643874646, 170.30493827160504, + 82.7362773029439, 173.92307692307693, 70.05584045584048, + 177.08357075023744, 54.84596391263053, 180.58129154795822, + 41.73190883190885, 183.14919278252614, 34.423266856600165, + 188.51623931623936, 30.279962013295354, 195.6273504273505, + 25.539221272554588, 201.75080721747398, 22.971320037986676, + 211.23228869895553, 22.37872744539408, 221.10883190883212, + 20.996011396011355, 224.8619183285852, 20.996011396011355, + 226.04710351377042, 23.56391263057927, 229.01006647673339, + 30.279962013295354 + ] + ], + "category_id": 1, + "category_name": "short sleeve top" + }, + "item1": { + "segmentation": [ + [ + 201.51804815682925, 224.7401022799914, 218.41555508203712, + 227.23317707223518, 236.42109524824218, 228.89522693373104, + 256.91971020669104, 229.44924355422967, 280.188408267633, + 230.2802684849776, 296.53189857234224, 230.2802684849776, + 313.7064138077994, 229.72625186447897, 315.32667803111013, + 236.8076070743661, 317.8197528233539, 240.96273172810572, + 318.65077775410185, 246.2258896228426, 321.4208608565949, + 253.15109737907534, 322.8059024078415, 265.0624547197956, + 324.74496057958663, 273.6497123375242, 325.9612827615598, + 284.4076070743661, 325.40726614106114, 299.9200724483274, + 324.29923290006394, 316.8175793735353, 322.0831664180694, + 325.9588536117625, 320.16803750266354, 336.5366716386107, + 316.0129128489239, 344.01589601534204, 315.18188791817596, + 357.86631152780745, 312.4118048156829, 368.1156190070319, + 308.5336884721926, 378.64193479650567, 306.31762199019806, + 385.29013424248905, 305.76360536969946, 398.3095248242066, + 305.48659705945016, 409.6668655444283, 304.94393777967184, + 419.3418708715109, 302.7278712976774, 427.0981035584915, + 301.3428297464308, 433.74630300447495, 301.3428297464308, + 445.3806520349459, 300.5118048156829, 461.72414233965515, + 299.89735776688684, 467.352311953974, 297.9582995951417, + 477.60161943319844, 295.1882164926486, 491.7290432559132, + 293.52616663115276, 497.2692094608994, 291.8641167696569, + 503.36339228638417, 291.3101001491583, 510.8426166631155, + 289.37104197741314, 513.8897080758579, 287.4433411463882, + 519.2043682079693, 283.0112081823993, 519.7583848284679, + 275.5319838056679, 519.4813765182186, 270.26882591093107, + 518.096334966972, 265.8366929469421, 513.6642020029831, + 263.62062646494763, 509.78608565949276, 264.7286597059449, + 498.9827615597697, 265.2826763264435, 478.76115491157015, + 266.1137012571914, 467.1268058810992, 266.1137012571914, + 454.6614319198803, 264.17464308544623, 441.64204133816276, + 263.06660984444903, 424.19051779245626, 261.5834221180482, + 407.2581504368212, 259.92137225655233, 396.45482633709815, + 257.1512891540592, 380.1113360323889, 257.42829746430857, + 359.05870445344146, 256.8742808438099, 338.56008949499255, + 256.8742808438099, 321.3855742595354, 254.10419774131685, + 320.5545493287875, 251.05710632857443, 326.6487321542723, + 249.39505646707858, 339.1141061154912, 249.11804815682927, + 356.28862135094835, 248.28702322608135, 372.3551033454083, + 245.23993181333896, 387.59056040912026, 243.5766673769444, + 409.1404219049649, 241.91461751544855, 424.92989558917554, + 240.52957596420202, 440.4423609631369, 238.86752610270617, + 455.40080971659955, 238.86752610270617, 470.91327509056083, + 238.31350948220754, 486.42574046452216, 238.81966759002768, + 501.19639889196685, 239.6506925207756, 511.168698060942, + 236.0495844875346, 515.6008310249309, 229.40138504155118, + 519.4789473684212, 221.6451523545705, 520.3099722991692, + 216.65900277008296, 517.2628808864267, 213.33490304709125, + 509.50664819944615, 208.3487534626037, 491.50110803324105, + 205.8556786703599, 475.1576177285318, 203.63961218836545, + 460.75318559556774, 203.63961218836545, 443.3016620498613, + 203.63961218836545, 421.9720221606645, 200.59252077562303, + 415.60083102493036, 197.5052844662264, 406.9847858512679, + 195.28921798423193, 392.0263370978052, 193.35015981248677, + 370.97370551885774, 190.857085020243, 343.82689111442545, + 187.8099936075006, 322.77425953547794, 187.0028979330919, + 309.89237161730256, 186.17187300234397, 291.33281483059886, + 188.11093117408916, 266.67907521841033, 191.15802258683155, + 250.3355849137011, 196.69818879181773, 234.82311953973982 + ] + ], + "category_id": 8, + "category_name": "trousers" + } +} diff --git a/examples/segmentation/028009.jpg b/examples/segmentation/028009.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d0d6453e012c2b1cb9dacb0e712b7723d8b5971b Binary files /dev/null and b/examples/segmentation/028009.jpg differ diff --git a/examples/segmentation/028009.json b/examples/segmentation/028009.json new file mode 100644 index 0000000000000000000000000000000000000000..705031ff1b0d6d043bed95605a94dc51bc4cdf96 --- /dev/null +++ b/examples/segmentation/028009.json @@ -0,0 +1,191 @@ +{ + "item2": { + "segmentation": [ + [ + 314.7474747474744, 204.84848484848482, 328.9696969696967, + 209.7373737373737, 342.74747474747454, 211.95959595959593, + 360.0808080808079, 211.07070707070704, 375.19191919191906, + 210.18181818181816, 384.5252525252524, 207.07070707070704, + 390.30303030303025, 204.84848484848482, 396.080808080808, + 201.29292929292924, 402.3030303030303, 204.40404040404036, + 412.969696969697, 203.9595959595959, 425.8585858585859, + 206.18181818181813, 434.3030303030304, 211.95959595959593, + 439.63636363636374, 223.0707070707071, 444.0808080808082, + 234.18181818181824, 448.52525252525265, 250.62626262626276, + 449.41414141414157, 260.848484848485, 452.08080808080825, + 279.0707070707073, 456.08080808080825, 300.84848484848516, + 457.858585858586, 308.40404040404076, 460.5252525252526, + 315.7575757575756, 460.96969696969705, 329.97979797979787, + 460.5252525252526, 345.9797979797979, 456.969696969697, + 363.75757575757575, 453.41414141414145, 373.5353535353536, + 450.3030303030303, 385.97979797979804, 447.1919191919192, + 393.9797979797981, 443.6363636363636, 401.9797979797981, + 438.3030303030303, 403.7575757575759, 433.85858585858585, + 401.09090909090924, 430.7474747474747, 393.0909090909092, + 426.7474747474747, 383.3131313131314, 424.9696969696969, + 374.8686868686869, 424.9696969696969, 369.0909090909091, + 423.63636363636357, 363.3131313131313, 423.63636363636357, + 359.3131313131313, 423.63636363636357, 352.6464646464646, + 420.9696969696969, 350.86868686868684, 422.74747474747466, + 345.53535353535347, 422.74747474747466, 340.64646464646455, + 422.74747474747466, 332.2020202020201, 421.8585858585858, + 321.53535353535335, 418.74747474747466, 313.0909090909089, + 416.5252525252524, 306.4242424242422, 412.9696969696969, + 314.8686868686867, 410.3030303030302, 320.20202020202004, + 411.6363636363635, 327.3131313131312, 414.74747474747466, + 336.2020202020201, 418.74747474747466, 351.7575757575757, + 420.9696969696969, 365.0909090909091, 423.1919191919191, + 377.0909090909091, 423.1919191919191, 385.0909090909092, + 424.5252525252525, 398.42424242424255, 396.0808080808079, + 398.42424242424255, 374.7474747474745, 400.6464646464648, + 354.7474747474744, 400.6464646464648, 331.6363636363632, + 400.6464646464648, 313.41414141414094, 400.6464646464648, + 305.4141414141409, 399.3131313131314, 297.4141414141409, + 396.6464646464648, 284.525252525252, 396.2020202020203, + 282.8686868686866, 391.59595959595964, 282.42424242424215, + 373.81818181818176, 282.42424242424215, 358.26262626262616, + 281.09090909090884, 334.70707070707056, 281.5353535353533, + 313.37373737373713, 283.31313131313107, 297.3737373737371, + 282.8686868686866, 283.1515151515148, 280.6464646464644, + 266.7070707070703, 271.313131313131, 253.3737373737369, + 264.6464646464643, 246.70707070707022, 257.5353535353532, + 239.59595959595907, 249.9797979797976, 228.9292929292924, + 242.42424242424204, 220.92929292929236, 233.17171717171723, + 209.01010101010093, 225.1717171717172, 194.78787878787867, + 222.06060606060606, 185.4545454545453, 224.2828282828283, + 179.6767676767675, 230.0606060606061, 171.67676767676747, + 232.72727272727278, 169.89898989898967, 243.83838383838392, + 167.67676767676744, 256.2828282828284, 165.4545454545452, + 274.06060606060623, 165.4545454545452, 291.8383838383841, + 167.67676767676744, 302.5050505050508, 168.1212121212119, + 310.94949494949526, 177.0101010101008, 314.0606060606064, + 181.45454545454527, 314.94949494949526, 187.2323232323231, + 312.7272727272731, 193.01010101010087, 307.8383838383842, + 191.2323232323231, 302.94949494949526, 193.45454545454533, + 292.727272727273, 193.45454545454533, 290.50505050505075, + 195.67676767676755, 287.39393939393966, 197.45454545454533, + 285.61616161616183, 197.45454545454533, 283.3939393939396, + 193.89898989898978, 278.94949494949515, 197.45454545454533, + 274.94949494949515, 199.67676767676755, 279.83838383838406, + 201.45454545454535, 286.50505050505075, 201.45454545454535, + 291.8383838383841, 201.8989898989898, 296.2828282828286, + 202.7878787878787, 303.3939393939397, 202.34343434343424 + ] + ], + "category_id": 2, + "category_name": "long sleeve top" + }, + "item1": { + "segmentation": [ + [ + 346.9494949494949, 660.6868686868687, 397.6161616161618, + 661.5757575757576, 398.06060606060623, 674.0202020202021, + 398.94949494949515, 691.3535353535356, 397.6161616161618, + 710.0202020202022, 395.838383838384, 726.0202020202023, + 393.1717171717173, 742.0202020202023, 346.9494949494949, + 738.9090909090912, 346.50505050505046, 724.2424242424245, + 347.3939393939394, 713.5757575757578, 348.72727272727275, + 706.0202020202022, 349.17171717171715, 686.0202020202022, + 348.72727272727275, 675.7979797979799, 347.3939393939394, + 667.7979797979799 + ], + [ + 283.71717171717165, 396.68686868686876, 289.9393939393939, + 396.68686868686876, 303.27272727272725, 397.1313131313132, + 312.16161616161617, 399.7979797979799, 334.3838383838385, + 400.68686868686876, 351.7171717171719, 400.68686868686876, + 361.93939393939417, 401.5757575757577, 376.60606060606085, + 401.5757575757577, 390.82828282828314, 398.46464646464653, + 410.3838383838388, 397.5757575757577, 425.0505050505055, + 394.46464646464653, 431.71717171717216, 422.9090909090911, + 434.38383838383885, 447.79797979798, 430.38383838383885, + 478.0202020202024, 423.2727272727277, 507.79797979798025, + 418.3838383838388, 530.0202020202025, 411.8787878787878, + 557.3333333333333, 403.43434343434336, 590.6666666666666, + 400.7676767676767, 611.5555555555557, 399.8787878787878, + 619.1111111111112, 399.8787878787878, 630.6666666666669, + 398.10101010101, 635.1111111111113, 399.43434343434336, + 641.7777777777779, 399.43434343434336, 656.4444444444447, + 398.10101010101, 662.666666666667, 347.4343434343432, 662.666666666667, + 346.1010101010098, 637.7777777777779, 347.4343434343432, + 610.6666666666667, 349.21212121212096, 576.4444444444445, + 350.98989898989873, 556.4444444444443, 349.6565656565654, + 541.3333333333331, 348.32323232323205, 535.9999999999998, + 348.32323232323205, 523.5555555555553, 349.21212121212096, + 505.33333333333303, 342.5454545454543, 511.5555555555553, + 338.9898989898987, 516.8888888888887, 334.5454545454542, + 523.5555555555553, 325.6565656565653, 543.111111111111, + 319.87878787878753, 556.4444444444443, 314.1010101010097, + 568.4444444444443, 307.8787878787875, 583.1111111111111, + 300.3232323232319, 608.0000000000001, 298.10101010100965, + 617.7777777777778, 298.5454545454541, 624.0000000000001, + 295.43434343434296, 628.0000000000001, 293.2121212121208, + 628.0000000000001, 293.6565656565652, 632.4444444444446, + 291.43434343434296, 638.6666666666669, 290.54545454545405, + 644.4444444444447, 292.3232323232319, 648.8888888888891, + 303.8787878787875, 667.1111111111114, 313.65656565656525, + 684.0000000000003, 319.87878787878753, 700.8888888888893, + 322.54545454545416, 712.8888888888894, 324.323232323232, + 720.0000000000005, 327.87878787878753, 731.5555555555561, + 330.9898989898987, 738.6666666666672, 331.87878787878753, + 743.1111111111117, 334.5454545454542, 745.7777777777783, + 336.3232323232325, 749.1313131313133, 338.54545454545473, + 754.0202020202022, 338.54545454545473, 757.5757575757577, + 341.6565656565658, 760.6868686868688, 344.76767676767696, + 767.3535353535356, 345.2121212121214, 770.9090909090911, + 346.9898989898992, 754.0202020202022, 347.43434343434365, + 738.909090909091, 393.2121212121216, 740.6868686868687, + 389.65656565656604, 764.6868686868688, 386.5454545454549, + 784.2424242424245, 384.3232323232327, 806.9090909090912, + 382.54545454545485, 812.686868686869, 381.13131313131316, + 818.7070707070708, 378.020202020202, 828.4848484848485, + 375.35353535353534, 839.5959595959597, 374.9090909090909, + 854.2626262626264, 373.1313131313131, 856.9292929292931, + 376.24242424242425, 864.9292929292931, 372.24242424242425, + 874.2626262626264, 366.4646464646464, 880.9292929292932, + 357.13131313131305, 872.9292929292932, 345.13131313131305, + 868.0404040404043, 337.131313131313, 867.1515151515154, + 337.131313131313, 856.0404040404042, 338.4646464646463, + 850.7070707070709, 336.2424242424241, 846.2626262626264, + 335.3535353535352, 841.3737373737375, 338.4646464646463, + 827.5959595959597, 342.0202020202019, 815.5959595959596, + 344.6868686868686, 809.3737373737374, 344.6868686868686, + 796.4848484848484, 344.6868686868686, 786.7070707070707, + 346.0202020202019, 779.151515151515, 344.24242424242414, + 776.0404040404039, 343.3535353535352, 786.2626262626262, + 342.0202020202019, 796.0404040404039, 338.90909090909076, + 801.8181818181818, 333.57575757575745, 809.3737373737374, + 326.02020202020185, 813.8181818181819, 320.242424242424, + 812.4848484848485, 318.02020202020185, 810.7070707070707, + 317.13131313131294, 807.1515151515151, 315.79797979797956, + 803.5959595959596, 313.57575757575734, 799.5959595959596, + 311.3535353535351, 793.8181818181818, 306.90909090909065, + 791.1515151515151, 305.57575757575734, 787.5959595959595, + 304.242424242424, 782.7070707070706, 302.02020202020174, + 776.4848484848484, 298.90909090909065, 773.8181818181816, + 294.90909090909065, 771.151515151515, 290.34343434343435, + 758.909090909091, 284.5656565656566, 742.020202020202, + 278.78787878787875, 729.5757575757575, 270.3434343434343, + 713.131313131313, 257.8989898989898, 689.1313131313129, + 247.2323232323231, 669.1313131313128, 239.23232323232307, + 657.5757575757573, 233.89898989898973, 642.9090909090905, + 233.0101010101008, 634.0202020202016, 233.45454545454527, + 630.0202020202016, 235.23232323232304, 611.7979797979793, + 241.93939393939402, 583.0707070707073, 245.93939393939405, + 567.5151515151516, 251.2727272727274, 540.4040404040404, + 256.1616161616163, 518.6262626262626, 260.60606060606074, + 501.2929292929292, 263.7171717171719, 493.7373737373736, + 268.16161616161634, 481.73737373737356, 270.38383838383857, + 469.73737373737356, 272.6060606060608, 462.18181818181796, + 276.1616161616164, 457.7373737373735, 276.1616161616164, + 454.1818181818179, 277.05050505050525, 450.1818181818179, + 278.828282828283, 433.292929292929, 278.3838383838386, + 419.0707070707067, 278.828282828283, 417.29292929292893, + 281.0505050505053, 414.1818181818178, 281.93939393939417, + 404.8484848484844, 283.71717171717194, 401.2929292929289 + ] + ], + "category_id": 8, + "category_name": "trousers" + } +} diff --git a/examples/slack_trial2-000150.png b/examples/slack_trial2-000150.png new file mode 100644 index 0000000000000000000000000000000000000000..e2f13ebb881e0bef4d57971cf0744ccfb831fda0 Binary files /dev/null and b/examples/slack_trial2-000150.png differ diff --git a/fetch_data.sh b/fetch_data.sh new file mode 100644 index 0000000000000000000000000000000000000000..ca7ea81c569a34f287fa589653ff40312ce617eb --- /dev/null +++ b/fetch_data.sh @@ -0,0 +1,60 @@ +#!/bin/bash +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; } + +mkdir -p data/smpl_related/models + +# username and password input +echo -e "\nYou need to register at https://icon.is.tue.mpg.de/, according to Installation Instruction." +read -p "Username (ICON):" username +read -p "Password (ICON):" password +username=$(urle $username) +password=$(urle $password) + +# SMPL (Male, Female) +echo -e "\nDownloading SMPL..." +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 +unzip data/smpl_related/models/SMPL_python_v.1.0.0.zip -d data/smpl_related/models +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 +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 +cd data/smpl_related/models +rm -rf *.zip __MACOSX smpl/models smpl/smpl_webuser +cd ../../.. + +# SMPL (Neutral, from SMPLIFY) +echo -e "\nDownloading SMPLify..." +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 +unzip data/smpl_related/models/mpips_smplify_public_v2.zip -d data/smpl_related/models +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 +cd data/smpl_related/models +rm -rf *.zip smplify_public +cd ../../.. + +# ICON +echo -e "\nDownloading ICON..." +wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=icon&sfile=icon_data.zip&resume=1' -O './data/icon_data.zip' --no-check-certificate --continue +cd data && unzip icon_data.zip +mv smpl_data smpl_related/ +rm -f icon_data.zip +cd .. + +function download_for_training () { + + # SMPL-X (optional) + echo -e "\nDownloading SMPL-X..." + wget --post-data "username=$1&password=$2" '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 + unzip data/smpl_related/models/models_smplx_v1_1.zip -d data/smpl_related + rm -f data/smpl_related/models/models_smplx_v1_1.zip + + # SMIL (optional) + echo -e "\nDownloading SMIL..." + wget --post-data "username=$1&password=$2" 'https://download.is.tue.mpg.de/download.php?domain=agora&sfile=smpl_kid_template.npy&resume=1' -O './data/smpl_related/models/smpl/smpl_kid_template.npy' --no-check-certificate --continue + wget --post-data "username=$1&password=$2" 'https://download.is.tue.mpg.de/download.php?domain=agora&sfile=smplx_kid_template.npy&resume=1' -O './data/smpl_related/models/smplx/smplx_kid_template.npy' --no-check-certificate --continue +} + + +read -p "(optional) Download models used for training (y/n)?" choice +case "$choice" in + y|Y ) download_for_training $username $password;; + n|N ) echo "Great job! Try the demo for now!";; + * ) echo "Invalid input! Please use y|Y or n|N";; +esac \ No newline at end of file diff --git a/install.sh b/install.sh new file mode 100644 index 0000000000000000000000000000000000000000..64b3a3ef17bf2a09007e672f34f92c2691289477 --- /dev/null +++ b/install.sh @@ -0,0 +1,16 @@ +# # conda installation +# wget https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.3-Linux-x86_64.sh +# chmod +x Miniconda3-py38_4.10.3-Linux-x86_64.sh +# bash Miniconda3-py38_4.10.3-Linux-x86_64.sh -b -f -p /home/user/.local +# rm Miniconda3-py38_4.10.3-Linux-x86_64.sh +# conda config --env --set always_yes true +# conda update -n base -c defaults conda -y + +# # conda environment setup +# conda env create -f environment.yaml +# conda init bash +# source /home/user/.bashrc +# source activate icon +nvidia-smi +pip install torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu111 +pip install -r requirement.txt \ No newline at end of file diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/common/__init__.py b/lib/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/common/cloth_extraction.py b/lib/common/cloth_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..8ecf8f1288457577a531547d67f00bf53d863134 --- /dev/null +++ b/lib/common/cloth_extraction.py @@ -0,0 +1,170 @@ +import numpy as np +import json +import os +import itertools +import trimesh +from matplotlib.path import Path +from collections import Counter +from sklearn.neighbors import KNeighborsClassifier + + +def load_segmentation(path, shape): + """ + Get a segmentation mask for a given image + Arguments: + path: path to the segmentation json file + shape: shape of the output mask + Returns: + Returns a segmentation mask + """ + with open(path) as json_file: + dict = json.load(json_file) + segmentations = [] + for key, val in dict.items(): + if not key.startswith('item'): + continue + + # Each item can have multiple polygons. Combine them to one + # segmentation_coord = list(itertools.chain.from_iterable(val['segmentation'])) + # segmentation_coord = np.round(np.array(segmentation_coord)).astype(int) + + coordinates = [] + for segmentation_coord in val['segmentation']: + # The format before is [x1,y1, x2, y2, ....] + x = segmentation_coord[::2] + y = segmentation_coord[1::2] + xy = np.vstack((x, y)).T + coordinates.append(xy) + + segmentations.append( + {'type': val['category_name'], 'type_id': val['category_id'], 'coordinates': coordinates}) + + return segmentations + + +def smpl_to_recon_labels(recon, smpl, k=1): + """ + Get the bodypart labels for the recon object by using the labels from the corresponding smpl object + Arguments: + recon: trimesh object (fully clothed model) + shape: trimesh object (smpl model) + k: number of nearest neighbours to use + Returns: + Returns a dictionary containing the bodypart and the corresponding indices + """ + smpl_vert_segmentation = json.load( + open(os.path.join(os.path.dirname(__file__), 'smpl_vert_segmentation.json'))) + n = smpl.vertices.shape[0] + y = np.array([None] * n) + for key, val in smpl_vert_segmentation.items(): + y[val] = key + + classifier = KNeighborsClassifier(n_neighbors=1) + classifier.fit(smpl.vertices, y) + + y_pred = classifier.predict(recon.vertices) + + recon_labels = {} + for key in smpl_vert_segmentation.keys(): + recon_labels[key] = list(np.argwhere( + y_pred == key).flatten().astype(int)) + + return recon_labels + + +def extract_cloth(recon, segmentation, K, R, t, smpl=None): + """ + Extract a portion of a mesh using 2d segmentation coordinates + Arguments: + recon: fully clothed mesh + seg_coord: segmentation coordinates in 2D (NDC) + K: intrinsic matrix of the projection + R: rotation matrix of the projection + t: translation vector of the projection + Returns: + Returns a submesh using the segmentation coordinates + """ + seg_coord = segmentation['coord_normalized'] + mesh = trimesh.Trimesh(recon.vertices, recon.faces) + extrinsic = np.zeros((3, 4)) + extrinsic[:3, :3] = R + extrinsic[:, 3] = t + P = K[:3, :3] @ extrinsic + + P_inv = np.linalg.pinv(P) + + # Each segmentation can contain multiple polygons + # We need to check them separately + points_so_far = [] + faces = recon.faces + for polygon in seg_coord: + n = len(polygon) + coords_h = np.hstack((polygon, np.ones((n, 1)))) + # Apply the inverse projection on homogeneus 2D coordinates to get the corresponding 3d Coordinates + XYZ = P_inv @ coords_h[:, :, None] + XYZ = XYZ.reshape((XYZ.shape[0], XYZ.shape[1])) + XYZ = XYZ[:, :3] / XYZ[:, 3, None] + + p = Path(XYZ[:, :2]) + + grid = p.contains_points(recon.vertices[:, :2]) + indeces = np.argwhere(grid == True) + points_so_far += list(indeces.flatten()) + + if smpl is not None: + num_verts = recon.vertices.shape[0] + recon_labels = smpl_to_recon_labels(recon, smpl) + body_parts_to_remove = ['rightHand', 'leftToeBase', 'leftFoot', 'rightFoot', 'head', + 'leftHandIndex1', 'rightHandIndex1', 'rightToeBase', 'leftHand', 'rightHand'] + type = segmentation['type_id'] + + # Remove additional bodyparts that are most likely not part of the segmentation but might intersect (e.g. hand in front of torso) + # https://github.com/switchablenorms/DeepFashion2 + # Short sleeve clothes + if type == 1 or type == 3 or type == 10: + body_parts_to_remove += ['leftForeArm', 'rightForeArm'] + # No sleeves at all or lower body clothes + elif type == 5 or type == 6 or type == 12 or type == 13 or type == 8 or type == 9: + body_parts_to_remove += ['leftForeArm', + 'rightForeArm', 'leftArm', 'rightArm'] + # Shorts + elif type == 7: + body_parts_to_remove += ['leftLeg', 'rightLeg', + 'leftForeArm', 'rightForeArm', 'leftArm', 'rightArm'] + + verts_to_remove = list(itertools.chain.from_iterable( + [recon_labels[part] for part in body_parts_to_remove])) + + label_mask = np.zeros(num_verts, dtype=bool) + label_mask[verts_to_remove] = True + + seg_mask = np.zeros(num_verts, dtype=bool) + seg_mask[points_so_far] = True + + # Remove points that belong to other bodyparts + # If a vertice in pointsSoFar is included in the bodyparts to remove, then these points should be removed + extra_verts_to_remove = np.array(list(seg_mask) and list(label_mask)) + + combine_mask = np.zeros(num_verts, dtype=bool) + combine_mask[points_so_far] = True + combine_mask[extra_verts_to_remove] = False + + all_indices = np.argwhere(combine_mask == True).flatten() + + i_x = np.where(np.in1d(faces[:, 0], all_indices))[0] + i_y = np.where(np.in1d(faces[:, 1], all_indices))[0] + i_z = np.where(np.in1d(faces[:, 2], all_indices))[0] + + faces_to_keep = np.array(list(set(i_x).union(i_y).union(i_z))) + mask = np.zeros(len(recon.faces), dtype=bool) + if len(faces_to_keep) > 0: + mask[faces_to_keep] = True + + mesh.update_faces(mask) + mesh.remove_unreferenced_vertices() + + # mesh.rezero() + + return mesh + + return None diff --git a/lib/common/config.py b/lib/common/config.py new file mode 100644 index 0000000000000000000000000000000000000000..2a69e8ca0ab7e42ae38055bb0b31e07ba42a4aed --- /dev/null +++ b/lib/common/config.py @@ -0,0 +1,218 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from yacs.config import CfgNode as CN +import os + +_C = CN(new_allowed=True) + +# needed by trainer +_C.name = 'default' +_C.gpus = [0] +_C.test_gpus = [1] +_C.root = "./data/" +_C.ckpt_dir = './data/ckpt/' +_C.resume_path = '' +_C.normal_path = '' +_C.corr_path = '' +_C.results_path = './data/results/' +_C.projection_mode = 'orthogonal' +_C.num_views = 1 +_C.sdf = False +_C.sdf_clip = 5.0 + +_C.lr_G = 1e-3 +_C.lr_C = 1e-3 +_C.lr_N = 2e-4 +_C.weight_decay = 0.0 +_C.momentum = 0.0 +_C.optim = 'RMSprop' +_C.schedule = [5, 10, 15] +_C.gamma = 0.1 + +_C.overfit = False +_C.resume = False +_C.test_mode = False +_C.test_uv = False +_C.draw_geo_thres = 0.60 +_C.num_sanity_val_steps = 2 +_C.fast_dev = 0 +_C.get_fit = False +_C.agora = False +_C.optim_cloth = False +_C.optim_body = False +_C.mcube_res = 256 +_C.clean_mesh = True +_C.remesh = False + +_C.batch_size = 4 +_C.num_threads = 8 + +_C.num_epoch = 10 +_C.freq_plot = 0.01 +_C.freq_show_train = 0.1 +_C.freq_show_val = 0.2 +_C.freq_eval = 0.5 +_C.accu_grad_batch = 4 + +_C.test_items = ['sv', 'mv', 'mv-fusion', 'hybrid', 'dc-pred', 'gt'] + +_C.net = CN() +_C.net.gtype = 'HGPIFuNet' +_C.net.ctype = 'resnet18' +_C.net.classifierIMF = 'MultiSegClassifier' +_C.net.netIMF = 'resnet18' +_C.net.norm = 'group' +_C.net.norm_mlp = 'group' +_C.net.norm_color = 'group' +_C.net.hg_down = 'ave_pool' +_C.net.num_views = 1 + +# kernel_size, stride, dilation, padding + +_C.net.conv1 = [7, 2, 1, 3] +_C.net.conv3x3 = [3, 1, 1, 1] + +_C.net.num_stack = 4 +_C.net.num_hourglass = 2 +_C.net.hourglass_dim = 256 +_C.net.voxel_dim = 32 +_C.net.resnet_dim = 120 +_C.net.mlp_dim = [320, 1024, 512, 256, 128, 1] +_C.net.mlp_dim_knn = [320, 1024, 512, 256, 128, 3] +_C.net.mlp_dim_color = [513, 1024, 512, 256, 128, 3] +_C.net.mlp_dim_multiseg = [1088, 2048, 1024, 500] +_C.net.res_layers = [2, 3, 4] +_C.net.filter_dim = 256 +_C.net.smpl_dim = 3 + +_C.net.cly_dim = 3 +_C.net.soft_dim = 64 +_C.net.z_size = 200.0 +_C.net.N_freqs = 10 +_C.net.geo_w = 0.1 +_C.net.norm_w = 0.1 +_C.net.dc_w = 0.1 +_C.net.C_cat_to_G = False + +_C.net.skip_hourglass = True +_C.net.use_tanh = True +_C.net.soft_onehot = True +_C.net.no_residual = True +_C.net.use_attention = False + +_C.net.prior_type = "sdf" +_C.net.smpl_feats = ['sdf', 'cmap', 'norm', 'vis'] +_C.net.use_filter = True +_C.net.use_cc = False +_C.net.use_PE = False +_C.net.use_IGR = False +_C.net.in_geo = () +_C.net.in_nml = () + +_C.dataset = CN() +_C.dataset.root = '' +_C.dataset.set_splits = [0.95, 0.04] +_C.dataset.types = [ + "3dpeople", "axyz", "renderpeople", "renderpeople_p27", "humanalloy" +] +_C.dataset.scales = [1.0, 100.0, 1.0, 1.0, 100.0 / 39.37] +_C.dataset.rp_type = "pifu900" +_C.dataset.th_type = 'train' +_C.dataset.input_size = 512 +_C.dataset.rotation_num = 3 +_C.dataset.num_precomp = 10 # Number of segmentation classifiers +_C.dataset.num_multiseg = 500 # Number of categories per classifier +_C.dataset.num_knn = 10 # for loss/error +_C.dataset.num_knn_dis = 20 # for accuracy +_C.dataset.num_verts_max = 20000 +_C.dataset.zray_type = False +_C.dataset.online_smpl = False +_C.dataset.noise_type = ['z-trans', 'pose', 'beta'] +_C.dataset.noise_scale = [0.0, 0.0, 0.0] +_C.dataset.num_sample_geo = 10000 +_C.dataset.num_sample_color = 0 +_C.dataset.num_sample_seg = 0 +_C.dataset.num_sample_knn = 10000 + +_C.dataset.sigma_geo = 5.0 +_C.dataset.sigma_color = 0.10 +_C.dataset.sigma_seg = 0.10 +_C.dataset.thickness_threshold = 20.0 +_C.dataset.ray_sample_num = 2 +_C.dataset.semantic_p = False +_C.dataset.remove_outlier = False + +_C.dataset.train_bsize = 1.0 +_C.dataset.val_bsize = 1.0 +_C.dataset.test_bsize = 1.0 + + +def get_cfg_defaults(): + """Get a yacs CfgNode object with default values for my_project.""" + # Return a clone so that the defaults will not be altered + # This is for the "local variable" use pattern + return _C.clone() + + +# Alternatively, provide a way to import the defaults as +# a global singleton: +cfg = _C # users can `from config import cfg` + +# cfg = get_cfg_defaults() +# cfg.merge_from_file('./configs/example.yaml') + +# # Now override from a list (opts could come from the command line) +# opts = ['dataset.root', './data/XXXX', 'learning_rate', '1e-2'] +# cfg.merge_from_list(opts) + + +def update_cfg(cfg_file): + # cfg = get_cfg_defaults() + _C.merge_from_file(cfg_file) + # return cfg.clone() + return _C + + +def parse_args(args): + cfg_file = args.cfg_file + if args.cfg_file is not None: + cfg = update_cfg(args.cfg_file) + else: + cfg = get_cfg_defaults() + + # if args.misc is not None: + # cfg.merge_from_list(args.misc) + + return cfg + + +def parse_args_extend(args): + if args.resume: + if not os.path.exists(args.log_dir): + raise ValueError( + 'Experiment are set to resume mode, but log directory does not exist.' + ) + + # load log's cfg + cfg_file = os.path.join(args.log_dir, 'cfg.yaml') + cfg = update_cfg(cfg_file) + + if args.misc is not None: + cfg.merge_from_list(args.misc) + else: + parse_args(args) diff --git a/lib/common/render.py b/lib/common/render.py new file mode 100644 index 0000000000000000000000000000000000000000..fce2d1658a24a8af3b393a92aa7354d06ad79952 --- /dev/null +++ b/lib/common/render.py @@ -0,0 +1,387 @@ +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from pytorch3d.renderer import ( + BlendParams, + blending, + look_at_view_transform, + FoVOrthographicCameras, + PointLights, + RasterizationSettings, + PointsRasterizationSettings, + PointsRenderer, + AlphaCompositor, + PointsRasterizer, + MeshRenderer, + MeshRasterizer, + SoftPhongShader, + SoftSilhouetteShader, + TexturesVertex, +) +from pytorch3d.renderer.mesh import TexturesVertex +from pytorch3d.structures import Meshes +from lib.dataset.mesh_util import SMPLX, get_visibility + +import lib.common.render_utils as util +import torch +import numpy as np +from PIL import Image +from tqdm import tqdm +import os +import cv2 +import math +from termcolor import colored + + +def image2vid(images, vid_path): + + w, h = images[0].size + videodims = (w, h) + fourcc = cv2.VideoWriter_fourcc(*'XVID') + video = cv2.VideoWriter(vid_path, fourcc, 30, videodims) + for image in images: + video.write(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)) + video.release() + + +def query_color(verts, faces, image, device): + """query colors from points and image + + Args: + verts ([B, 3]): [query verts] + faces ([M, 3]): [query faces] + image ([B, 3, H, W]): [full image] + + Returns: + [np.float]: [return colors] + """ + + verts = verts.float().to(device) + faces = faces.long().to(device) + + (xy, z) = verts.split([2, 1], dim=1) + visibility = get_visibility(xy, z, faces[:, [0, 2, 1]]).flatten() + uv = xy.unsqueeze(0).unsqueeze(2) # [B, N, 2] + uv = uv * torch.tensor([1.0, -1.0]).type_as(uv) + colors = (torch.nn.functional.grid_sample(image, uv, align_corners=True)[ + 0, :, :, 0].permute(1, 0) + 1.0) * 0.5 * 255.0 + colors[visibility == 0.0] = ((Meshes(verts.unsqueeze(0), faces.unsqueeze( + 0)).verts_normals_padded().squeeze(0) + 1.0) * 0.5 * 255.0)[visibility == 0.0] + + return colors.detach().cpu() + + +class cleanShader(torch.nn.Module): + def __init__(self, device="cpu", cameras=None, blend_params=None): + super().__init__() + self.cameras = cameras + self.blend_params = blend_params if blend_params is not None else BlendParams() + + def forward(self, fragments, meshes, **kwargs): + cameras = kwargs.get("cameras", self.cameras) + if cameras is None: + msg = "Cameras must be specified either at initialization \ + or in the forward pass of TexturedSoftPhongShader" + + raise ValueError(msg) + + # get renderer output + blend_params = kwargs.get("blend_params", self.blend_params) + texels = meshes.sample_textures(fragments) + images = blending.softmax_rgb_blend( + texels, fragments, blend_params, znear=-256, zfar=256 + ) + + return images + + +class Render: + def __init__(self, size=512, device=torch.device("cuda:0")): + self.device = device + self.mesh_y_center = 100.0 + self.dis = 100.0 + self.scale = 1.0 + self.size = size + self.cam_pos = [(0, 100, 100)] + + self.mesh = None + self.deform_mesh = None + self.pcd = None + self.renderer = None + self.meshRas = None + self.type = None + self.knn = None + self.knn_inverse = None + + self.smpl_seg = None + self.smpl_cmap = None + + self.smplx = SMPLX() + + self.uv_rasterizer = util.Pytorch3dRasterizer(self.size) + + def get_camera(self, cam_id): + + R, T = look_at_view_transform( + eye=[self.cam_pos[cam_id]], + at=((0, self.mesh_y_center, 0),), + up=((0, 1, 0),), + ) + + camera = FoVOrthographicCameras( + device=self.device, + R=R, + T=T, + znear=100.0, + zfar=-100.0, + max_y=100.0, + min_y=-100.0, + max_x=100.0, + min_x=-100.0, + scale_xyz=(self.scale * np.ones(3),), + ) + + return camera + + def init_renderer(self, camera, type="clean_mesh", bg="gray"): + + if "mesh" in type: + + # rasterizer + self.raster_settings_mesh = RasterizationSettings( + image_size=self.size, + blur_radius=np.log(1.0 / 1e-4) * 1e-7, + faces_per_pixel=30, + ) + self.meshRas = MeshRasterizer( + cameras=camera, raster_settings=self.raster_settings_mesh + ) + + if bg == "black": + blendparam = BlendParams(1e-4, 1e-4, (0.0, 0.0, 0.0)) + elif bg == "white": + blendparam = BlendParams(1e-4, 1e-8, (1.0, 1.0, 1.0)) + elif bg == "gray": + blendparam = BlendParams(1e-4, 1e-8, (0.5, 0.5, 0.5)) + + if type == "ori_mesh": + + lights = PointLights( + device=self.device, + ambient_color=((0.8, 0.8, 0.8),), + diffuse_color=((0.2, 0.2, 0.2),), + specular_color=((0.0, 0.0, 0.0),), + location=[[0.0, 200.0, 0.0]], + ) + + self.renderer = MeshRenderer( + rasterizer=self.meshRas, + shader=SoftPhongShader( + device=self.device, + cameras=camera, + lights=lights, + blend_params=blendparam, + ), + ) + + if type == "silhouette": + self.raster_settings_silhouette = RasterizationSettings( + image_size=self.size, + blur_radius=np.log(1.0 / 1e-4 - 1.0) * 5e-5, + faces_per_pixel=50, + cull_backfaces=True, + ) + + self.silhouetteRas = MeshRasterizer( + cameras=camera, raster_settings=self.raster_settings_silhouette + ) + self.renderer = MeshRenderer( + rasterizer=self.silhouetteRas, shader=SoftSilhouetteShader() + ) + + if type == "pointcloud": + self.raster_settings_pcd = PointsRasterizationSettings( + image_size=self.size, radius=0.006, points_per_pixel=10 + ) + + self.pcdRas = PointsRasterizer( + cameras=camera, raster_settings=self.raster_settings_pcd + ) + self.renderer = PointsRenderer( + rasterizer=self.pcdRas, + compositor=AlphaCompositor(background_color=(0, 0, 0)), + ) + + if type == "clean_mesh": + + self.renderer = MeshRenderer( + rasterizer=self.meshRas, + shader=cleanShader( + device=self.device, cameras=camera, blend_params=blendparam + ), + ) + + def VF2Mesh(self, verts, faces): + + if not torch.is_tensor(verts): + verts = torch.tensor(verts) + if not torch.is_tensor(faces): + faces = torch.tensor(faces) + + if verts.ndimension() == 2: + verts = verts.unsqueeze(0).float() + if faces.ndimension() == 2: + faces = faces.unsqueeze(0).long() + + verts = verts.to(self.device) + faces = faces.to(self.device) + + mesh = Meshes(verts, faces).to(self.device) + + mesh.textures = TexturesVertex( + verts_features=(mesh.verts_normals_padded() + 1.0) * 0.5 + ) + + return mesh + + def load_meshes(self, verts, faces): + """load mesh into the pytorch3d renderer + + Args: + verts ([N,3]): verts + faces ([N,3]): faces + offset ([N,3]): offset + """ + + # camera setting + self.scale = 100.0 + self.mesh_y_center = 0.0 + + self.cam_pos = [ + (0, self.mesh_y_center, 100.0), + (100.0, self.mesh_y_center, 0), + (0, self.mesh_y_center, -100.0), + (-100.0, self.mesh_y_center, 0), + ] + + self.type = "color" + + if isinstance(verts, list): + self.meshes = [] + for V, F in zip(verts, faces): + self.meshes.append(self.VF2Mesh(V, F)) + else: + self.meshes = [self.VF2Mesh(verts, faces)] + + def get_depth_map(self, cam_ids=[0, 2]): + + depth_maps = [] + for cam_id in cam_ids: + self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray") + fragments = self.meshRas(self.meshes[0]) + depth_map = fragments.zbuf[..., 0].squeeze(0) + if cam_id == 2: + depth_map = torch.fliplr(depth_map) + depth_maps.append(depth_map) + + return depth_maps + + def get_rgb_image(self, cam_ids=[0, 2]): + + images = [] + for cam_id in range(len(self.cam_pos)): + if cam_id in cam_ids: + self.init_renderer(self.get_camera( + cam_id), "clean_mesh", "gray") + if len(cam_ids) == 4: + rendered_img = ( + self.renderer(self.meshes[0])[ + 0:1, :, :, :3].permute(0, 3, 1, 2) + - 0.5 + ) * 2.0 + else: + rendered_img = ( + self.renderer(self.meshes[0])[ + 0:1, :, :, :3].permute(0, 3, 1, 2) + - 0.5 + ) * 2.0 + if cam_id == 2 and len(cam_ids) == 2: + rendered_img = torch.flip(rendered_img, dims=[3]) + images.append(rendered_img) + + return images + + def get_rendered_video(self, images, save_path): + + self.cam_pos = [] + for angle in range(360): + self.cam_pos.append( + ( + 100.0 * math.cos(np.pi / 180 * angle), + self.mesh_y_center, + 100.0 * math.sin(np.pi / 180 * angle), + ) + ) + + old_shape = np.array(images[0].shape[:2]) + new_shape = np.around( + (self.size / old_shape[0]) * old_shape).astype(np.int) + + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + video = cv2.VideoWriter( + save_path, fourcc, 30, (self.size * len(self.meshes) + + new_shape[1] * len(images), self.size) + ) + + pbar = tqdm(range(len(self.cam_pos))) + pbar.set_description(colored(f"exporting video {os.path.basename(save_path)}...", "blue")) + for cam_id in pbar: + self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray") + + img_lst = [ + np.array(Image.fromarray(img).resize(new_shape[::-1])).astype(np.uint8)[ + :, :, [2, 1, 0] + ] + for img in images + ] + + for mesh in self.meshes: + rendered_img = ( + (self.renderer(mesh)[0, :, :, :3] * 255.0) + .detach() + .cpu() + .numpy() + .astype(np.uint8) + ) + + img_lst.append(rendered_img) + final_img = np.concatenate(img_lst, axis=1) + video.write(final_img) + + video.release() + + def get_silhouette_image(self, cam_ids=[0, 2]): + + images = [] + for cam_id in range(len(self.cam_pos)): + if cam_id in cam_ids: + self.init_renderer(self.get_camera(cam_id), "silhouette") + rendered_img = self.renderer(self.meshes[0])[0:1, :, :, 3] + if cam_id == 2 and len(cam_ids) == 2: + rendered_img = torch.flip(rendered_img, dims=[2]) + images.append(rendered_img) + + return images diff --git a/lib/common/render_utils.py b/lib/common/render_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..09b38cadc8a5b66d765f9f62596709fa7325c773 --- /dev/null +++ b/lib/common/render_utils.py @@ -0,0 +1,221 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import torch +from torch import nn +import trimesh +import math +from typing import NewType +from pytorch3d.structures import Meshes +from pytorch3d.renderer.mesh import rasterize_meshes + +Tensor = NewType('Tensor', torch.Tensor) + + +def solid_angles(points: Tensor, + triangles: Tensor, + thresh: float = 1e-8) -> Tensor: + ''' Compute solid angle between the input points and triangles + Follows the method described in: + The Solid Angle of a Plane Triangle + A. VAN OOSTEROM AND J. STRACKEE + IEEE TRANSACTIONS ON BIOMEDICAL ENGINEERING, + VOL. BME-30, NO. 2, FEBRUARY 1983 + Parameters + ----------- + points: BxQx3 + Tensor of input query points + triangles: BxFx3x3 + Target triangles + thresh: float + float threshold + Returns + ------- + solid_angles: BxQxF + A tensor containing the solid angle between all query points + and input triangles + ''' + # Center the triangles on the query points. Size should be BxQxFx3x3 + centered_tris = triangles[:, None] - points[:, :, None, None] + + # BxQxFx3 + norms = torch.norm(centered_tris, dim=-1) + + # Should be BxQxFx3 + cross_prod = torch.cross(centered_tris[:, :, :, 1], + centered_tris[:, :, :, 2], + dim=-1) + # Should be BxQxF + numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1) + del cross_prod + + dot01 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 1]).sum(dim=-1) + dot12 = (centered_tris[:, :, :, 1] * centered_tris[:, :, :, 2]).sum(dim=-1) + dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1) + del centered_tris + + denominator = (norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] + + dot02 * norms[:, :, :, 1] + dot12 * norms[:, :, :, 0]) + del dot01, dot12, dot02, norms + + # Should be BxQ + solid_angle = torch.atan2(numerator, denominator) + del numerator, denominator + + torch.cuda.empty_cache() + + return 2 * solid_angle + + +def winding_numbers(points: Tensor, + triangles: Tensor, + thresh: float = 1e-8) -> Tensor: + ''' Uses winding_numbers to compute inside/outside + Robust inside-outside segmentation using generalized winding numbers + Alec Jacobson, + Ladislav Kavan, + Olga Sorkine-Hornung + Fast Winding Numbers for Soups and Clouds SIGGRAPH 2018 + Gavin Barill + NEIL G. Dickson + Ryan Schmidt + David I.W. Levin + and Alec Jacobson + Parameters + ----------- + points: BxQx3 + Tensor of input query points + triangles: BxFx3x3 + Target triangles + thresh: float + float threshold + Returns + ------- + winding_numbers: BxQ + A tensor containing the Generalized winding numbers + ''' + # The generalized winding number is the sum of solid angles of the point + # with respect to all triangles. + return 1 / (4 * math.pi) * solid_angles(points, triangles, + thresh=thresh).sum(dim=-1) + + +def batch_contains(verts, faces, points): + + B = verts.shape[0] + N = points.shape[1] + + verts = verts.detach().cpu() + faces = faces.detach().cpu() + points = points.detach().cpu() + contains = torch.zeros(B, N) + + for i in range(B): + contains[i] = torch.as_tensor( + trimesh.Trimesh(verts[i], faces[i]).contains(points[i])) + + return 2.0 * (contains - 0.5) + + +def dict2obj(d): + # if isinstance(d, list): + # d = [dict2obj(x) for x in d] + if not isinstance(d, dict): + return d + + class C(object): + pass + + o = C() + for k in d: + o.__dict__[k] = dict2obj(d[k]) + return o + + +def face_vertices(vertices, faces): + """ + :param vertices: [batch size, number of vertices, 3] + :param faces: [batch size, number of faces, 3] + :return: [batch size, number of faces, 3, 3] + """ + + bs, nv = vertices.shape[:2] + bs, nf = faces.shape[:2] + device = vertices.device + faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * + nv)[:, None, None] + vertices = vertices.reshape((bs * nv, vertices.shape[-1])) + + return vertices[faces.long()] + + +class Pytorch3dRasterizer(nn.Module): + """ Borrowed from https://github.com/facebookresearch/pytorch3d + Notice: + x,y,z are in image space, normalized + can only render squared image now + """ + + def __init__(self, image_size=224): + """ + use fixed raster_settings for rendering faces + """ + super().__init__() + raster_settings = { + 'image_size': image_size, + 'blur_radius': 0.0, + 'faces_per_pixel': 1, + 'bin_size': None, + 'max_faces_per_bin': None, + 'perspective_correct': True, + 'cull_backfaces': True, + } + raster_settings = dict2obj(raster_settings) + self.raster_settings = raster_settings + + def forward(self, vertices, faces, attributes=None): + fixed_vertices = vertices.clone() + fixed_vertices[..., :2] = -fixed_vertices[..., :2] + meshes_screen = Meshes(verts=fixed_vertices.float(), + faces=faces.long()) + raster_settings = self.raster_settings + pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( + meshes_screen, + image_size=raster_settings.image_size, + blur_radius=raster_settings.blur_radius, + faces_per_pixel=raster_settings.faces_per_pixel, + bin_size=raster_settings.bin_size, + max_faces_per_bin=raster_settings.max_faces_per_bin, + perspective_correct=raster_settings.perspective_correct, + ) + vismask = (pix_to_face > -1).float() + D = attributes.shape[-1] + attributes = attributes.clone() + attributes = attributes.view(attributes.shape[0] * attributes.shape[1], + 3, attributes.shape[-1]) + N, H, W, K, _ = bary_coords.shape + mask = pix_to_face == -1 + pix_to_face = pix_to_face.clone() + pix_to_face[mask] = 0 + idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D) + pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D) + pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2) + pixel_vals[mask] = 0 # Replace masked values in output. + pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) + pixel_vals = torch.cat( + [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) + return pixel_vals diff --git a/lib/common/seg3d_lossless.py b/lib/common/seg3d_lossless.py new file mode 100644 index 0000000000000000000000000000000000000000..bd8133934939251db7b8b47198fdacb11384cd78 --- /dev/null +++ b/lib/common/seg3d_lossless.py @@ -0,0 +1,604 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + + +from .seg3d_utils import ( + create_grid3D, + plot_mask3D, + SmoothConv3D, +) + +import torch +import torch.nn as nn +import numpy as np +import torch.nn.functional as F +import mcubes +from kaolin.ops.conversions import voxelgrids_to_trianglemeshes +import logging + +logging.getLogger("lightning").setLevel(logging.ERROR) + + +class Seg3dLossless(nn.Module): + def __init__(self, + query_func, + b_min, + b_max, + resolutions, + channels=1, + balance_value=0.5, + align_corners=False, + visualize=False, + debug=False, + use_cuda_impl=False, + faster=False, + use_shadow=False, + **kwargs): + """ + align_corners: same with how you process gt. (grid_sample / interpolate) + """ + super().__init__() + self.query_func = query_func + self.register_buffer( + 'b_min', + torch.tensor(b_min).float().unsqueeze(1)) # [bz, 1, 3] + self.register_buffer( + 'b_max', + torch.tensor(b_max).float().unsqueeze(1)) # [bz, 1, 3] + + # ti.init(arch=ti.cuda) + # self.mciso_taichi = MCISO(dim=3, N=resolutions[-1]-1) + + if type(resolutions[0]) is int: + resolutions = torch.tensor([(res, res, res) + for res in resolutions]) + else: + resolutions = torch.tensor(resolutions) + self.register_buffer('resolutions', resolutions) + self.batchsize = self.b_min.size(0) + assert self.batchsize == 1 + self.balance_value = balance_value + self.channels = channels + assert self.channels == 1 + self.align_corners = align_corners + self.visualize = visualize + self.debug = debug + self.use_cuda_impl = use_cuda_impl + self.faster = faster + self.use_shadow = use_shadow + + for resolution in resolutions: + assert resolution[0] % 2 == 1 and resolution[1] % 2 == 1, \ + f"resolution {resolution} need to be odd becuase of align_corner." + + # init first resolution + init_coords = create_grid3D(0, + resolutions[-1] - 1, + steps=resolutions[0]) # [N, 3] + init_coords = init_coords.unsqueeze(0).repeat(self.batchsize, 1, + 1) # [bz, N, 3] + self.register_buffer('init_coords', init_coords) + + # some useful tensors + calculated = torch.zeros( + (self.resolutions[-1][2], self.resolutions[-1][1], + self.resolutions[-1][0]), + dtype=torch.bool) + self.register_buffer('calculated', calculated) + + gird8_offsets = torch.stack( + torch.meshgrid([ + torch.tensor([-1, 0, 1]), + torch.tensor([-1, 0, 1]), + torch.tensor([-1, 0, 1]) + ])).int().view(3, -1).t() # [27, 3] + self.register_buffer('gird8_offsets', gird8_offsets) + + # smooth convs + self.smooth_conv3x3 = SmoothConv3D(in_channels=1, + out_channels=1, + kernel_size=3) + self.smooth_conv5x5 = SmoothConv3D(in_channels=1, + out_channels=1, + kernel_size=5) + self.smooth_conv7x7 = SmoothConv3D(in_channels=1, + out_channels=1, + kernel_size=7) + self.smooth_conv9x9 = SmoothConv3D(in_channels=1, + out_channels=1, + kernel_size=9) + + def batch_eval(self, coords, **kwargs): + """ + coords: in the coordinates of last resolution + **kwargs: for query_func + """ + coords = coords.detach() + # normalize coords to fit in [b_min, b_max] + if self.align_corners: + coords2D = coords.float() / (self.resolutions[-1] - 1) + else: + step = 1.0 / self.resolutions[-1].float() + coords2D = coords.float() / self.resolutions[-1] + step / 2 + coords2D = coords2D * (self.b_max - self.b_min) + self.b_min + # query function + occupancys = self.query_func(**kwargs, points=coords2D) + if type(occupancys) is list: + occupancys = torch.stack(occupancys) # [bz, C, N] + assert len(occupancys.size()) == 3, \ + "query_func should return a occupancy with shape of [bz, C, N]" + return occupancys + + def forward(self, **kwargs): + if self.faster: + return self._forward_faster(**kwargs) + else: + return self._forward(**kwargs) + + def _forward_faster(self, **kwargs): + """ + In faster mode, we make following changes to exchange accuracy for speed: + 1. no conflict checking: 4.88 fps -> 6.56 fps + 2. smooth_conv9x9 ~ smooth_conv3x3 for different resolution + 3. last step no examine + """ + final_W = self.resolutions[-1][0] + final_H = self.resolutions[-1][1] + final_D = self.resolutions[-1][2] + + for resolution in self.resolutions: + W, H, D = resolution + stride = (self.resolutions[-1] - 1) / (resolution - 1) + + # first step + if torch.equal(resolution, self.resolutions[0]): + coords = self.init_coords.clone() # torch.long + occupancys = self.batch_eval(coords, **kwargs) + occupancys = occupancys.view(self.batchsize, self.channels, D, + H, W) + if (occupancys > 0.5).sum() == 0: + # return F.interpolate( + # occupancys, size=(final_D, final_H, final_W), + # mode="linear", align_corners=True) + return None + + if self.visualize: + self.plot(occupancys, coords, final_D, final_H, final_W) + + with torch.no_grad(): + coords_accum = coords / stride + + # last step + elif torch.equal(resolution, self.resolutions[-1]): + + with torch.no_grad(): + # here true is correct! + valid = F.interpolate( + (occupancys > self.balance_value).float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + # here true is correct! + occupancys = F.interpolate(occupancys.float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + # is_boundary = (valid > 0.0) & (valid < 1.0) + is_boundary = valid == 0.5 + + # next steps + else: + coords_accum *= 2 + + with torch.no_grad(): + # here true is correct! + valid = F.interpolate( + (occupancys > self.balance_value).float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + # here true is correct! + occupancys = F.interpolate(occupancys.float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + is_boundary = (valid > 0.0) & (valid < 1.0) + + with torch.no_grad(): + if torch.equal(resolution, self.resolutions[1]): + is_boundary = (self.smooth_conv9x9(is_boundary.float()) + > 0)[0, 0] + elif torch.equal(resolution, self.resolutions[2]): + is_boundary = (self.smooth_conv7x7(is_boundary.float()) + > 0)[0, 0] + else: + is_boundary = (self.smooth_conv3x3(is_boundary.float()) + > 0)[0, 0] + + coords_accum = coords_accum.long() + is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1], + coords_accum[0, :, 0]] = False + point_coords = is_boundary.permute( + 2, 1, 0).nonzero(as_tuple=False).unsqueeze(0) + point_indices = (point_coords[:, :, 2] * H * W + + point_coords[:, :, 1] * W + + point_coords[:, :, 0]) + + R, C, D, H, W = occupancys.shape + + # inferred value + coords = point_coords * stride + + if coords.size(1) == 0: + continue + occupancys_topk = self.batch_eval(coords, **kwargs) + + # put mask point predictions to the right places on the upsampled grid. + R, C, D, H, W = occupancys.shape + point_indices = point_indices.unsqueeze(1).expand(-1, C, -1) + occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( + 2, point_indices, occupancys_topk).view(R, C, D, H, W)) + + with torch.no_grad(): + voxels = coords / stride + coords_accum = torch.cat([voxels, coords_accum], + dim=1).unique(dim=1) + + return occupancys[0, 0] + + def _forward(self, **kwargs): + """ + output occupancy field would be: + (bz, C, res, res) + """ + final_W = self.resolutions[-1][0] + final_H = self.resolutions[-1][1] + final_D = self.resolutions[-1][2] + + calculated = self.calculated.clone() + + for resolution in self.resolutions: + W, H, D = resolution + stride = (self.resolutions[-1] - 1) / (resolution - 1) + + if self.visualize: + this_stage_coords = [] + + # first step + if torch.equal(resolution, self.resolutions[0]): + coords = self.init_coords.clone() # torch.long + occupancys = self.batch_eval(coords, **kwargs) + occupancys = occupancys.view(self.batchsize, self.channels, D, + H, W) + + if self.visualize: + self.plot(occupancys, coords, final_D, final_H, final_W) + + with torch.no_grad(): + coords_accum = coords / stride + calculated[coords[0, :, 2], coords[0, :, 1], + coords[0, :, 0]] = True + + # next steps + else: + coords_accum *= 2 + + with torch.no_grad(): + # here true is correct! + valid = F.interpolate( + (occupancys > self.balance_value).float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + # here true is correct! + occupancys = F.interpolate(occupancys.float(), + size=(D, H, W), + mode="trilinear", + align_corners=True) + + is_boundary = (valid > 0.0) & (valid < 1.0) + + with torch.no_grad(): + # TODO + if self.use_shadow and torch.equal(resolution, + self.resolutions[-1]): + # larger z means smaller depth here + depth_res = resolution[2].item() + depth_index = torch.linspace(0, + depth_res - 1, + steps=depth_res).type_as( + occupancys.device) + depth_index_max = torch.max( + (occupancys > self.balance_value) * + (depth_index + 1), + dim=-1, + keepdim=True)[0] - 1 + shadow = depth_index < depth_index_max + is_boundary[shadow] = False + is_boundary = is_boundary[0, 0] + else: + is_boundary = (self.smooth_conv3x3(is_boundary.float()) + > 0)[0, 0] + # is_boundary = is_boundary[0, 0] + + is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1], + coords_accum[0, :, 0]] = False + point_coords = is_boundary.permute( + 2, 1, 0).nonzero(as_tuple=False).unsqueeze(0) + point_indices = (point_coords[:, :, 2] * H * W + + point_coords[:, :, 1] * W + + point_coords[:, :, 0]) + + R, C, D, H, W = occupancys.shape + # interpolated value + occupancys_interp = torch.gather( + occupancys.reshape(R, C, D * H * W), 2, + point_indices.unsqueeze(1)) + + # inferred value + coords = point_coords * stride + + if coords.size(1) == 0: + continue + occupancys_topk = self.batch_eval(coords, **kwargs) + if self.visualize: + this_stage_coords.append(coords) + + # put mask point predictions to the right places on the upsampled grid. + R, C, D, H, W = occupancys.shape + point_indices = point_indices.unsqueeze(1).expand(-1, C, -1) + occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( + 2, point_indices, occupancys_topk).view(R, C, D, H, W)) + + with torch.no_grad(): + # conflicts + conflicts = ((occupancys_interp - self.balance_value) * + (occupancys_topk - self.balance_value) < 0)[0, + 0] + + if self.visualize: + self.plot(occupancys, coords, final_D, final_H, + final_W) + + voxels = coords / stride + coords_accum = torch.cat([voxels, coords_accum], + dim=1).unique(dim=1) + calculated[coords[0, :, 2], coords[0, :, 1], + coords[0, :, 0]] = True + + while conflicts.sum() > 0: + if self.use_shadow and torch.equal(resolution, + self.resolutions[-1]): + break + + with torch.no_grad(): + conflicts_coords = coords[0, conflicts, :] + + if self.debug: + self.plot(occupancys, + conflicts_coords.unsqueeze(0), + final_D, + final_H, + final_W, + title='conflicts') + + conflicts_boundary = (conflicts_coords.int() + + self.gird8_offsets.unsqueeze(1) * + stride.int()).reshape( + -1, 3).long().unique(dim=0) + conflicts_boundary[:, 0] = ( + conflicts_boundary[:, 0].clamp( + 0, + calculated.size(2) - 1)) + conflicts_boundary[:, 1] = ( + conflicts_boundary[:, 1].clamp( + 0, + calculated.size(1) - 1)) + conflicts_boundary[:, 2] = ( + conflicts_boundary[:, 2].clamp( + 0, + calculated.size(0) - 1)) + + coords = conflicts_boundary[calculated[ + conflicts_boundary[:, 2], conflicts_boundary[:, 1], + conflicts_boundary[:, 0]] == False] + + if self.debug: + self.plot(occupancys, + coords.unsqueeze(0), + final_D, + final_H, + final_W, + title='coords') + + coords = coords.unsqueeze(0) + point_coords = coords / stride + point_indices = (point_coords[:, :, 2] * H * W + + point_coords[:, :, 1] * W + + point_coords[:, :, 0]) + + R, C, D, H, W = occupancys.shape + # interpolated value + occupancys_interp = torch.gather( + occupancys.reshape(R, C, D * H * W), 2, + point_indices.unsqueeze(1)) + + # inferred value + coords = point_coords * stride + + if coords.size(1) == 0: + break + occupancys_topk = self.batch_eval(coords, **kwargs) + if self.visualize: + this_stage_coords.append(coords) + + with torch.no_grad(): + # conflicts + conflicts = ((occupancys_interp - self.balance_value) * + (occupancys_topk - self.balance_value) < + 0)[0, 0] + + # put mask point predictions to the right places on the upsampled grid. + point_indices = point_indices.unsqueeze(1).expand( + -1, C, -1) + occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( + 2, point_indices, occupancys_topk).view(R, C, D, H, W)) + + with torch.no_grad(): + voxels = coords / stride + coords_accum = torch.cat([voxels, coords_accum], + dim=1).unique(dim=1) + calculated[coords[0, :, 2], coords[0, :, 1], + coords[0, :, 0]] = True + + if self.visualize: + this_stage_coords = torch.cat(this_stage_coords, dim=1) + self.plot(occupancys, this_stage_coords, final_D, final_H, + final_W) + + return occupancys[0, 0] + + def plot(self, + occupancys, + coords, + final_D, + final_H, + final_W, + title='', + **kwargs): + final = F.interpolate(occupancys.float(), + size=(final_D, final_H, final_W), + mode="trilinear", + align_corners=True) # here true is correct! + x = coords[0, :, 0].to("cpu") + y = coords[0, :, 1].to("cpu") + z = coords[0, :, 2].to("cpu") + + plot_mask3D(final[0, 0].to("cpu"), title, (x, y, z), **kwargs) + + def find_vertices(self, sdf, direction="front"): + ''' + - direction: "front" | "back" | "left" | "right" + ''' + resolution = sdf.size(2) + if direction == "front": + pass + elif direction == "left": + sdf = sdf.permute(2, 1, 0) + elif direction == "back": + inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() + sdf = sdf[inv_idx, :, :] + elif direction == "right": + inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() + sdf = sdf[:, :, inv_idx] + sdf = sdf.permute(2, 1, 0) + + inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() + sdf = sdf[inv_idx, :, :] + sdf_all = sdf.permute(2, 1, 0) + + # shadow + grad_v = (sdf_all > 0.5) * torch.linspace( + resolution, 1, steps=resolution).to(sdf.device) + grad_c = torch.ones_like(sdf_all) * torch.linspace( + 0, resolution - 1, steps=resolution).to(sdf.device) + max_v, max_c = grad_v.max(dim=2) + shadow = grad_c > max_c.view(resolution, resolution, 1) + keep = (sdf_all > 0.5) & (~shadow) + + p1 = keep.nonzero(as_tuple=False).t() # [3, N] + p2 = p1.clone() # z + p2[2, :] = (p2[2, :] - 2).clamp(0, resolution) + p3 = p1.clone() # y + p3[1, :] = (p3[1, :] - 2).clamp(0, resolution) + p4 = p1.clone() # x + p4[0, :] = (p4[0, :] - 2).clamp(0, resolution) + + v1 = sdf_all[p1[0, :], p1[1, :], p1[2, :]] + v2 = sdf_all[p2[0, :], p2[1, :], p2[2, :]] + v3 = sdf_all[p3[0, :], p3[1, :], p3[2, :]] + v4 = sdf_all[p4[0, :], p4[1, :], p4[2, :]] + + X = p1[0, :].long() # [N,] + Y = p1[1, :].long() # [N,] + Z = p2[2, :].float() * (0.5 - v1) / (v2 - v1) + \ + p1[2, :].float() * (v2 - 0.5) / (v2 - v1) # [N,] + Z = Z.clamp(0, resolution) + + # normal + norm_z = v2 - v1 + norm_y = v3 - v1 + norm_x = v4 - v1 + # print (v2.min(dim=0)[0], v2.max(dim=0)[0], v3.min(dim=0)[0], v3.max(dim=0)[0]) + + norm = torch.stack([norm_x, norm_y, norm_z], dim=1) + norm = norm / torch.norm(norm, p=2, dim=1, keepdim=True) + + return X, Y, Z, norm + + def render_normal(self, resolution, X, Y, Z, norm): + image = torch.ones((1, 3, resolution, resolution), + dtype=torch.float32).to(norm.device) + color = (norm + 1) / 2.0 + color = color.clamp(0, 1) + image[0, :, Y, X] = color.t() + return image + + def display(self, sdf): + + # render + X, Y, Z, norm = self.find_vertices(sdf, direction="front") + image1 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) + X, Y, Z, norm = self.find_vertices(sdf, direction="left") + image2 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) + X, Y, Z, norm = self.find_vertices(sdf, direction="right") + image3 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) + X, Y, Z, norm = self.find_vertices(sdf, direction="back") + image4 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) + + image = torch.cat([image1, image2, image3, image4], axis=3) + image = image.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255.0 + + return np.uint8(image) + + def export_mesh(self, occupancys): + + final = occupancys[1:, 1:, 1:].contiguous() + + if final.shape[0] > 256: + # for voxelgrid larger than 256^3, the required GPU memory will be > 9GB + # thus we use CPU marching_cube to avoid "CUDA out of memory" + occu_arr = final.detach().cpu().numpy() # non-smooth surface + # occu_arr = mcubes.smooth(final.detach().cpu().numpy()) # smooth surface + vertices, triangles = mcubes.marching_cubes( + occu_arr, self.balance_value) + verts = torch.as_tensor(vertices[:, [2, 1, 0]]) + faces = torch.as_tensor(triangles.astype( + np.long), dtype=torch.long)[:, [0, 2, 1]] + else: + torch.cuda.empty_cache() + vertices, triangles = voxelgrids_to_trianglemeshes( + final.unsqueeze(0)) + verts = vertices[0][:, [2, 1, 0]].cpu() + faces = triangles[0][:, [0, 2, 1]].cpu() + + return verts, faces diff --git a/lib/common/seg3d_utils.py b/lib/common/seg3d_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..74d6cf044e684eb2fb95f54629b7f81be929848e --- /dev/null +++ b/lib/common/seg3d_utils.py @@ -0,0 +1,392 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import torch +import torch.nn as nn +import torch.nn.functional as F +import matplotlib.pyplot as plt + + +def plot_mask2D(mask, + title="", + point_coords=None, + figsize=10, + point_marker_size=5): + ''' + Simple plotting tool to show intermediate mask predictions and points + where PointRend is applied. + + Args: + mask (Tensor): mask prediction of shape HxW + title (str): title for the plot + point_coords ((Tensor, Tensor)): x and y point coordinates + figsize (int): size of the figure to plot + point_marker_size (int): marker size for points + ''' + + H, W = mask.shape + plt.figure(figsize=(figsize, figsize)) + if title: + title += ", " + plt.title("{}resolution {}x{}".format(title, H, W), fontsize=30) + plt.ylabel(H, fontsize=30) + plt.xlabel(W, fontsize=30) + plt.xticks([], []) + plt.yticks([], []) + plt.imshow(mask.detach(), + interpolation="nearest", + cmap=plt.get_cmap('gray')) + if point_coords is not None: + plt.scatter(x=point_coords[0], + y=point_coords[1], + color="red", + s=point_marker_size, + clip_on=True) + plt.xlim(-0.5, W - 0.5) + plt.ylim(H - 0.5, -0.5) + plt.show() + + +def plot_mask3D(mask=None, + title="", + point_coords=None, + figsize=1500, + point_marker_size=8, + interactive=True): + ''' + Simple plotting tool to show intermediate mask predictions and points + where PointRend is applied. + + Args: + mask (Tensor): mask prediction of shape DxHxW + title (str): title for the plot + point_coords ((Tensor, Tensor, Tensor)): x and y and z point coordinates + figsize (int): size of the figure to plot + point_marker_size (int): marker size for points + ''' + import trimesh + import vtkplotter + from skimage import measure + + vp = vtkplotter.Plotter(title=title, size=(figsize, figsize)) + vis_list = [] + + if mask is not None: + mask = mask.detach().to("cpu").numpy() + mask = mask.transpose(2, 1, 0) + + # marching cube to find surface + verts, faces, normals, values = measure.marching_cubes_lewiner( + mask, 0.5, gradient_direction='ascent') + + # create a mesh + mesh = trimesh.Trimesh(verts, faces) + mesh.visual.face_colors = [200, 200, 250, 100] + vis_list.append(mesh) + + if point_coords is not None: + point_coords = torch.stack(point_coords, 1).to("cpu").numpy() + + # import numpy as np + # select_x = np.logical_and(point_coords[:, 0] >= 16, point_coords[:, 0] <= 112) + # select_y = np.logical_and(point_coords[:, 1] >= 48, point_coords[:, 1] <= 272) + # select_z = np.logical_and(point_coords[:, 2] >= 16, point_coords[:, 2] <= 112) + # select = np.logical_and(np.logical_and(select_x, select_y), select_z) + # point_coords = point_coords[select, :] + + pc = vtkplotter.Points(point_coords, r=point_marker_size, c='red') + vis_list.append(pc) + + vp.show(*vis_list, + bg="white", + axes=1, + interactive=interactive, + azimuth=30, + elevation=30) + + +def create_grid3D(min, max, steps): + if type(min) is int: + min = (min, min, min) # (x, y, z) + if type(max) is int: + max = (max, max, max) # (x, y) + if type(steps) is int: + steps = (steps, steps, steps) # (x, y, z) + arrangeX = torch.linspace(min[0], max[0], steps[0]).long() + arrangeY = torch.linspace(min[1], max[1], steps[1]).long() + arrangeZ = torch.linspace(min[2], max[2], steps[2]).long() + gridD, girdH, gridW = torch.meshgrid([arrangeZ, arrangeY, arrangeX]) + coords = torch.stack([gridW, girdH, + gridD]) # [2, steps[0], steps[1], steps[2]] + coords = coords.view(3, -1).t() # [N, 3] + return coords + + +def create_grid2D(min, max, steps): + if type(min) is int: + min = (min, min) # (x, y) + if type(max) is int: + max = (max, max) # (x, y) + if type(steps) is int: + steps = (steps, steps) # (x, y) + arrangeX = torch.linspace(min[0], max[0], steps[0]).long() + arrangeY = torch.linspace(min[1], max[1], steps[1]).long() + girdH, gridW = torch.meshgrid([arrangeY, arrangeX]) + coords = torch.stack([gridW, girdH]) # [2, steps[0], steps[1]] + coords = coords.view(2, -1).t() # [N, 2] + return coords + + +class SmoothConv2D(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3): + super().__init__() + assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}" + self.padding = (kernel_size - 1) // 2 + + weight = torch.ones( + (in_channels, out_channels, kernel_size, kernel_size), + dtype=torch.float32) / (kernel_size**2) + self.register_buffer('weight', weight) + + def forward(self, input): + return F.conv2d(input, self.weight, padding=self.padding) + + +class SmoothConv3D(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3): + super().__init__() + assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}" + self.padding = (kernel_size - 1) // 2 + + weight = torch.ones( + (in_channels, out_channels, kernel_size, kernel_size, kernel_size), + dtype=torch.float32) / (kernel_size**3) + self.register_buffer('weight', weight) + + def forward(self, input): + return F.conv3d(input, self.weight, padding=self.padding) + + +def build_smooth_conv3D(in_channels=1, + out_channels=1, + kernel_size=3, + padding=1): + smooth_conv = torch.nn.Conv3d(in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + padding=padding) + smooth_conv.weight.data = torch.ones( + (in_channels, out_channels, kernel_size, kernel_size, kernel_size), + dtype=torch.float32) / (kernel_size**3) + smooth_conv.bias.data = torch.zeros(out_channels) + return smooth_conv + + +def build_smooth_conv2D(in_channels=1, + out_channels=1, + kernel_size=3, + padding=1): + smooth_conv = torch.nn.Conv2d(in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + padding=padding) + smooth_conv.weight.data = torch.ones( + (in_channels, out_channels, kernel_size, kernel_size), + dtype=torch.float32) / (kernel_size**2) + smooth_conv.bias.data = torch.zeros(out_channels) + return smooth_conv + + +def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points, + **kwargs): + """ + Find `num_points` most uncertain points from `uncertainty_map` grid. + Args: + uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty + values for a set of points on a regular H x W x D grid. + num_points (int): The number of points P to select. + Returns: + point_indices (Tensor): A tensor of shape (N, P) that contains indices from + [0, H x W x D) of the most uncertain points. + point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized + coordinates of the most uncertain points from the H x W x D grid. + """ + R, _, D, H, W = uncertainty_map.shape + # h_step = 1.0 / float(H) + # w_step = 1.0 / float(W) + # d_step = 1.0 / float(D) + + num_points = min(D * H * W, num_points) + point_scores, point_indices = torch.topk(uncertainty_map.view( + R, D * H * W), + k=num_points, + dim=1) + point_coords = torch.zeros(R, + num_points, + 3, + dtype=torch.float, + device=uncertainty_map.device) + # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step + # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step + # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step + point_coords[:, :, 0] = (point_indices % W).to(torch.float) # x + point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float) # y + point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float) # z + print(f"resolution {D} x {H} x {W}", point_scores.min(), + point_scores.max()) + return point_indices, point_coords + + +def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points, + clip_min): + """ + Find `num_points` most uncertain points from `uncertainty_map` grid. + Args: + uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty + values for a set of points on a regular H x W x D grid. + num_points (int): The number of points P to select. + Returns: + point_indices (Tensor): A tensor of shape (N, P) that contains indices from + [0, H x W x D) of the most uncertain points. + point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized + coordinates of the most uncertain points from the H x W x D grid. + """ + R, _, D, H, W = uncertainty_map.shape + # h_step = 1.0 / float(H) + # w_step = 1.0 / float(W) + # d_step = 1.0 / float(D) + + assert R == 1, "batchsize > 1 is not implemented!" + uncertainty_map = uncertainty_map.view(D * H * W) + indices = (uncertainty_map >= clip_min).nonzero().squeeze(1) + num_points = min(num_points, indices.size(0)) + point_scores, point_indices = torch.topk(uncertainty_map[indices], + k=num_points, + dim=0) + point_indices = indices[point_indices].unsqueeze(0) + + point_coords = torch.zeros(R, + num_points, + 3, + dtype=torch.float, + device=uncertainty_map.device) + # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step + # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step + # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step + point_coords[:, :, 0] = (point_indices % W).to(torch.float) # x + point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float) # y + point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float) # z + # print (f"resolution {D} x {H} x {W}", point_scores.min(), point_scores.max()) + return point_indices, point_coords + + +def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points, + **kwargs): + """ + Find `num_points` most uncertain points from `uncertainty_map` grid. + Args: + uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty + values for a set of points on a regular H x W grid. + num_points (int): The number of points P to select. + Returns: + point_indices (Tensor): A tensor of shape (N, P) that contains indices from + [0, H x W) of the most uncertain points. + point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized + coordinates of the most uncertain points from the H x W grid. + """ + R, _, H, W = uncertainty_map.shape + # h_step = 1.0 / float(H) + # w_step = 1.0 / float(W) + + num_points = min(H * W, num_points) + point_scores, point_indices = torch.topk(uncertainty_map.view(R, H * W), + k=num_points, + dim=1) + point_coords = torch.zeros(R, + num_points, + 2, + dtype=torch.long, + device=uncertainty_map.device) + # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step + # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step + point_coords[:, :, 0] = (point_indices % W).to(torch.long) + point_coords[:, :, 1] = (point_indices // W).to(torch.long) + # print (point_scores.min(), point_scores.max()) + return point_indices, point_coords + + +def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points, + clip_min): + """ + Find `num_points` most uncertain points from `uncertainty_map` grid. + Args: + uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty + values for a set of points on a regular H x W grid. + num_points (int): The number of points P to select. + Returns: + point_indices (Tensor): A tensor of shape (N, P) that contains indices from + [0, H x W) of the most uncertain points. + point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized + coordinates of the most uncertain points from the H x W grid. + """ + R, _, H, W = uncertainty_map.shape + # h_step = 1.0 / float(H) + # w_step = 1.0 / float(W) + + assert R == 1, "batchsize > 1 is not implemented!" + uncertainty_map = uncertainty_map.view(H * W) + indices = (uncertainty_map >= clip_min).nonzero().squeeze(1) + num_points = min(num_points, indices.size(0)) + point_scores, point_indices = torch.topk(uncertainty_map[indices], + k=num_points, + dim=0) + point_indices = indices[point_indices].unsqueeze(0) + + point_coords = torch.zeros(R, + num_points, + 2, + dtype=torch.long, + device=uncertainty_map.device) + # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step + # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step + point_coords[:, :, 0] = (point_indices % W).to(torch.long) + point_coords[:, :, 1] = (point_indices // W).to(torch.long) + # print (point_scores.min(), point_scores.max()) + return point_indices, point_coords + + +def calculate_uncertainty(logits, classes=None, balance_value=0.5): + """ + We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the + foreground class in `classes`. + Args: + logits (Tensor): A tensor of shape (R, C, ...) or (R, 1, ...) for class-specific or + class-agnostic, where R is the total number of predicted masks in all images and C is + the number of foreground classes. The values are logits. + classes (list): A list of length R that contains either predicted of ground truth class + for eash predicted mask. + Returns: + scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with + the most uncertain locations having the highest uncertainty score. + """ + if logits.shape[1] == 1: + gt_class_logits = logits + else: + gt_class_logits = logits[ + torch.arange(logits.shape[0], device=logits.device), + classes].unsqueeze(1) + return -torch.abs(gt_class_logits - balance_value) diff --git a/lib/common/smpl_vert_segmentation.json b/lib/common/smpl_vert_segmentation.json new file mode 100644 index 0000000000000000000000000000000000000000..b3244cce450e13f1095a1c3af676f4c8fdea5633 --- /dev/null +++ b/lib/common/smpl_vert_segmentation.json @@ -0,0 +1,7440 @@ +{ + "rightHand": [ + 5442, + 5443, + 5444, + 5445, + 5446, + 5447, + 5448, + 5449, + 5450, + 5451, + 5452, + 5453, + 5454, + 5455, + 5456, + 5457, + 5458, + 5459, + 5460, + 5461, + 5462, + 5463, + 5464, + 5465, + 5466, + 5467, + 5468, + 5469, + 5470, + 5471, + 5472, + 5473, + 5474, + 5475, + 5476, + 5477, + 5478, + 5479, + 5480, + 5481, + 5482, + 5483, + 5484, + 5485, + 5486, + 5487, + 5492, + 5493, + 5494, + 5495, + 5496, + 5497, + 5502, + 5503, + 5504, + 5505, + 5506, + 5507, + 5508, + 5509, + 5510, + 5511, + 5512, + 5513, + 5514, + 5515, + 5516, + 5517, + 5518, + 5519, + 5520, + 5521, + 5522, + 5523, + 5524, + 5525, + 5526, + 5527, + 5530, + 5531, + 5532, + 5533, + 5534, + 5535, + 5536, + 5537, + 5538, + 5539, + 5540, + 5541, + 5542, + 5543, + 5544, + 5545, + 5546, + 5547, + 5548, + 5549, + 5550, + 5551, + 5552, + 5553, + 5554, + 5555, + 5556, + 5557, + 5558, + 5559, + 5560, + 5561, + 5562, + 5569, + 5571, + 5574, + 5575, + 5576, + 5577, + 5578, + 5579, + 5580, + 5581, + 5582, + 5583, + 5588, + 5589, + 5592, + 5593, + 5594, + 5595, + 5596, + 5597, + 5598, + 5599, + 5600, + 5601, + 5602, + 5603, + 5604, + 5605, + 5610, + 5611, + 5612, + 5613, + 5614, + 5621, + 5622, + 5625, + 5631, + 5632, + 5633, + 5634, + 5635, + 5636, + 5637, + 5638, + 5639, + 5640, + 5641, + 5643, + 5644, + 5645, + 5646, + 5649, + 5650, + 5652, + 5653, + 5654, + 5655, + 5656, + 5657, + 5658, + 5659, + 5660, + 5661, + 5662, + 5663, + 5664, + 5667, + 5670, + 5671, + 5672, + 5673, + 5674, + 5675, + 5682, + 5683, + 5684, + 5685, + 5686, + 5687, + 5688, + 5689, + 5690, + 5692, + 5695, + 5697, + 5698, + 5699, + 5700, + 5701, + 5707, + 5708, + 5709, + 5710, + 5711, + 5712, + 5713, + 5714, + 5715, + 5716, + 5717, + 5718, + 5719, + 5720, + 5721, + 5723, + 5724, + 5725, + 5726, + 5727, + 5728, + 5729, + 5730, + 5731, + 5732, + 5735, + 5736, + 5737, + 5738, + 5739, + 5740, + 5745, + 5746, + 5748, + 5749, + 5750, + 5751, + 5752, + 6056, + 6057, + 6066, + 6067, + 6158, + 6159, + 6160, + 6161, + 6162, + 6163, + 6164, + 6165, + 6166, + 6167, + 6168, + 6169, + 6170, + 6171, + 6172, + 6173, + 6174, + 6175, + 6176, + 6177, + 6178, + 6179, + 6180, + 6181, + 6182, + 6183, + 6184, + 6185, + 6186, + 6187, + 6188, + 6189, + 6190, + 6191, + 6192, + 6193, + 6194, + 6195, + 6196, + 6197, + 6198, + 6199, + 6200, + 6201, + 6202, + 6203, + 6204, + 6205, + 6206, + 6207, + 6208, + 6209, + 6210, + 6211, + 6212, + 6213, + 6214, + 6215, + 6216, + 6217, + 6218, + 6219, + 6220, + 6221, + 6222, + 6223, + 6224, + 6225, + 6226, + 6227, + 6228, + 6229, + 6230, + 6231, + 6232, + 6233, + 6234, + 6235, + 6236, + 6237, + 6238, + 6239 + ], + "rightUpLeg": [ + 4320, + 4321, + 4323, + 4324, + 4333, + 4334, + 4335, + 4336, + 4337, + 4338, + 4339, + 4340, + 4356, + 4357, + 4358, + 4359, + 4360, + 4361, + 4362, + 4363, + 4364, + 4365, + 4366, + 4367, + 4383, + 4384, + 4385, + 4386, + 4387, + 4388, + 4389, + 4390, + 4391, + 4392, + 4393, + 4394, + 4395, + 4396, + 4397, + 4398, + 4399, + 4400, + 4401, + 4419, + 4420, + 4421, + 4422, + 4430, + 4431, + 4432, + 4433, + 4434, + 4435, + 4436, + 4437, + 4438, + 4439, + 4440, + 4441, + 4442, + 4443, + 4444, + 4445, + 4446, + 4447, + 4448, + 4449, + 4450, + 4451, + 4452, + 4453, + 4454, + 4455, + 4456, + 4457, + 4458, + 4459, + 4460, + 4461, + 4462, + 4463, + 4464, + 4465, + 4466, + 4467, + 4468, + 4469, + 4470, + 4471, + 4472, + 4473, + 4474, + 4475, + 4476, + 4477, + 4478, + 4479, + 4480, + 4481, + 4482, + 4483, + 4484, + 4485, + 4486, + 4487, + 4488, + 4489, + 4490, + 4491, + 4492, + 4493, + 4494, + 4495, + 4496, + 4497, + 4498, + 4499, + 4500, + 4501, + 4502, + 4503, + 4504, + 4505, + 4506, + 4507, + 4508, + 4509, + 4510, + 4511, + 4512, + 4513, + 4514, + 4515, + 4516, + 4517, + 4518, + 4519, + 4520, + 4521, + 4522, + 4523, + 4524, + 4525, + 4526, + 4527, + 4528, + 4529, + 4530, + 4531, + 4532, + 4623, + 4624, + 4625, + 4626, + 4627, + 4628, + 4629, + 4630, + 4631, + 4632, + 4633, + 4634, + 4645, + 4646, + 4647, + 4648, + 4649, + 4650, + 4651, + 4652, + 4653, + 4654, + 4655, + 4656, + 4657, + 4658, + 4659, + 4660, + 4670, + 4671, + 4672, + 4673, + 4704, + 4705, + 4706, + 4707, + 4708, + 4709, + 4710, + 4711, + 4712, + 4713, + 4745, + 4746, + 4757, + 4758, + 4759, + 4760, + 4801, + 4802, + 4829, + 4834, + 4835, + 4836, + 4837, + 4838, + 4839, + 4840, + 4841, + 4924, + 4925, + 4926, + 4928, + 4929, + 4930, + 4931, + 4932, + 4933, + 4934, + 4935, + 4936, + 4948, + 4949, + 4950, + 4951, + 4952, + 4970, + 4971, + 4972, + 4973, + 4983, + 4984, + 4985, + 4986, + 4987, + 4988, + 4989, + 4990, + 4991, + 4992, + 4993, + 5004, + 5005, + 6546, + 6547, + 6548, + 6549, + 6552, + 6553, + 6554, + 6555, + 6556, + 6873, + 6877 + ], + "leftArm": [ + 626, + 627, + 628, + 629, + 634, + 635, + 680, + 681, + 716, + 717, + 718, + 719, + 769, + 770, + 771, + 772, + 773, + 774, + 775, + 776, + 777, + 778, + 779, + 780, + 784, + 785, + 786, + 787, + 788, + 789, + 790, + 791, + 792, + 793, + 1231, + 1232, + 1233, + 1234, + 1258, + 1259, + 1260, + 1261, + 1271, + 1281, + 1282, + 1310, + 1311, + 1314, + 1315, + 1340, + 1341, + 1342, + 1343, + 1355, + 1356, + 1357, + 1358, + 1376, + 1377, + 1378, + 1379, + 1380, + 1381, + 1382, + 1383, + 1384, + 1385, + 1386, + 1387, + 1388, + 1389, + 1390, + 1391, + 1392, + 1393, + 1394, + 1395, + 1396, + 1397, + 1398, + 1399, + 1400, + 1402, + 1403, + 1405, + 1406, + 1407, + 1408, + 1409, + 1410, + 1411, + 1412, + 1413, + 1414, + 1415, + 1416, + 1428, + 1429, + 1430, + 1431, + 1432, + 1433, + 1438, + 1439, + 1440, + 1441, + 1442, + 1443, + 1444, + 1445, + 1502, + 1505, + 1506, + 1507, + 1508, + 1509, + 1510, + 1538, + 1541, + 1542, + 1543, + 1545, + 1619, + 1620, + 1621, + 1622, + 1631, + 1632, + 1633, + 1634, + 1635, + 1636, + 1637, + 1638, + 1639, + 1640, + 1641, + 1642, + 1645, + 1646, + 1647, + 1648, + 1649, + 1650, + 1651, + 1652, + 1653, + 1654, + 1655, + 1656, + 1658, + 1659, + 1661, + 1662, + 1664, + 1666, + 1667, + 1668, + 1669, + 1670, + 1671, + 1672, + 1673, + 1674, + 1675, + 1676, + 1677, + 1678, + 1679, + 1680, + 1681, + 1682, + 1683, + 1684, + 1696, + 1697, + 1698, + 1703, + 1704, + 1705, + 1706, + 1707, + 1708, + 1709, + 1710, + 1711, + 1712, + 1713, + 1714, + 1715, + 1716, + 1717, + 1718, + 1719, + 1720, + 1725, + 1731, + 1732, + 1733, + 1734, + 1735, + 1737, + 1739, + 1740, + 1745, + 1746, + 1747, + 1748, + 1749, + 1751, + 1761, + 1830, + 1831, + 1844, + 1845, + 1846, + 1850, + 1851, + 1854, + 1855, + 1858, + 1860, + 1865, + 1866, + 1867, + 1869, + 1870, + 1871, + 1874, + 1875, + 1876, + 1877, + 1878, + 1882, + 1883, + 1888, + 1889, + 1892, + 1900, + 1901, + 1902, + 1903, + 1904, + 1909, + 2819, + 2820, + 2821, + 2822, + 2895, + 2896, + 2897, + 2898, + 2899, + 2900, + 2901, + 2902, + 2903, + 2945, + 2946, + 2974, + 2975, + 2976, + 2977, + 2978, + 2979, + 2980, + 2981, + 2982, + 2983, + 2984, + 2985, + 2986, + 2987, + 2988, + 2989, + 2990, + 2991, + 2992, + 2993, + 2994, + 2995, + 2996, + 3002, + 3013 + ], + "leftLeg": [ + 995, + 998, + 999, + 1002, + 1004, + 1005, + 1008, + 1010, + 1012, + 1015, + 1016, + 1018, + 1019, + 1043, + 1044, + 1047, + 1048, + 1049, + 1050, + 1051, + 1052, + 1053, + 1054, + 1055, + 1056, + 1057, + 1058, + 1059, + 1060, + 1061, + 1062, + 1063, + 1064, + 1065, + 1066, + 1067, + 1068, + 1069, + 1070, + 1071, + 1072, + 1073, + 1074, + 1075, + 1076, + 1077, + 1078, + 1079, + 1080, + 1081, + 1082, + 1083, + 1084, + 1085, + 1086, + 1087, + 1088, + 1089, + 1090, + 1091, + 1092, + 1093, + 1094, + 1095, + 1096, + 1097, + 1098, + 1099, + 1100, + 1101, + 1102, + 1103, + 1104, + 1105, + 1106, + 1107, + 1108, + 1109, + 1110, + 1111, + 1112, + 1113, + 1114, + 1115, + 1116, + 1117, + 1118, + 1119, + 1120, + 1121, + 1122, + 1123, + 1124, + 1125, + 1126, + 1127, + 1128, + 1129, + 1130, + 1131, + 1132, + 1133, + 1134, + 1135, + 1136, + 1148, + 1149, + 1150, + 1151, + 1152, + 1153, + 1154, + 1155, + 1156, + 1157, + 1158, + 1175, + 1176, + 1177, + 1178, + 1179, + 1180, + 1181, + 1182, + 1183, + 1369, + 1370, + 1371, + 1372, + 1373, + 1374, + 1375, + 1464, + 1465, + 1466, + 1467, + 1468, + 1469, + 1470, + 1471, + 1472, + 1473, + 1474, + 1522, + 1523, + 1524, + 1525, + 1526, + 1527, + 1528, + 1529, + 1530, + 1531, + 1532, + 3174, + 3175, + 3176, + 3177, + 3178, + 3179, + 3180, + 3181, + 3182, + 3183, + 3184, + 3185, + 3186, + 3187, + 3188, + 3189, + 3190, + 3191, + 3192, + 3193, + 3194, + 3195, + 3196, + 3197, + 3198, + 3199, + 3200, + 3201, + 3202, + 3203, + 3204, + 3205, + 3206, + 3207, + 3208, + 3209, + 3210, + 3319, + 3320, + 3321, + 3322, + 3323, + 3324, + 3325, + 3326, + 3327, + 3328, + 3329, + 3330, + 3331, + 3332, + 3333, + 3334, + 3335, + 3432, + 3433, + 3434, + 3435, + 3436, + 3469, + 3472, + 3473, + 3474 + ], + "leftToeBase": [ + 3211, + 3212, + 3213, + 3214, + 3215, + 3216, + 3217, + 3218, + 3219, + 3220, + 3221, + 3222, + 3223, + 3224, + 3225, + 3226, + 3227, + 3228, + 3229, + 3230, + 3231, + 3232, + 3233, + 3234, + 3235, + 3236, + 3237, + 3238, + 3239, + 3240, + 3241, + 3242, + 3243, + 3244, + 3245, + 3246, + 3247, + 3248, + 3249, + 3250, + 3251, + 3252, + 3253, + 3254, + 3255, + 3256, + 3257, + 3258, + 3259, + 3260, + 3261, + 3262, + 3263, + 3264, + 3265, + 3266, + 3267, + 3268, + 3269, + 3270, + 3271, + 3272, + 3273, + 3274, + 3275, + 3276, + 3277, + 3278, + 3279, + 3280, + 3281, + 3282, + 3283, + 3284, + 3285, + 3286, + 3287, + 3288, + 3289, + 3290, + 3291, + 3292, + 3293, + 3294, + 3295, + 3296, + 3297, + 3298, + 3299, + 3300, + 3301, + 3302, + 3303, + 3304, + 3305, + 3306, + 3307, + 3308, + 3309, + 3310, + 3311, + 3312, + 3313, + 3314, + 3315, + 3316, + 3317, + 3318, + 3336, + 3337, + 3340, + 3342, + 3344, + 3346, + 3348, + 3350, + 3352, + 3354, + 3357, + 3358, + 3360, + 3362 + ], + "leftFoot": [ + 3327, + 3328, + 3329, + 3330, + 3331, + 3332, + 3333, + 3334, + 3335, + 3336, + 3337, + 3338, + 3339, + 3340, + 3341, + 3342, + 3343, + 3344, + 3345, + 3346, + 3347, + 3348, + 3349, + 3350, + 3351, + 3352, + 3353, + 3354, + 3355, + 3356, + 3357, + 3358, + 3359, + 3360, + 3361, + 3362, + 3363, + 3364, + 3365, + 3366, + 3367, + 3368, + 3369, + 3370, + 3371, + 3372, + 3373, + 3374, + 3375, + 3376, + 3377, + 3378, + 3379, + 3380, + 3381, + 3382, + 3383, + 3384, + 3385, + 3386, + 3387, + 3388, + 3389, + 3390, + 3391, + 3392, + 3393, + 3394, + 3395, + 3396, + 3397, + 3398, + 3399, + 3400, + 3401, + 3402, + 3403, + 3404, + 3405, + 3406, + 3407, + 3408, + 3409, + 3410, + 3411, + 3412, + 3413, + 3414, + 3415, + 3416, + 3417, + 3418, + 3419, + 3420, + 3421, + 3422, + 3423, + 3424, + 3425, + 3426, + 3427, + 3428, + 3429, + 3430, + 3431, + 3432, + 3433, + 3434, + 3435, + 3436, + 3437, + 3438, + 3439, + 3440, + 3441, + 3442, + 3443, + 3444, + 3445, + 3446, + 3447, + 3448, + 3449, + 3450, + 3451, + 3452, + 3453, + 3454, + 3455, + 3456, + 3457, + 3458, + 3459, + 3460, + 3461, + 3462, + 3463, + 3464, + 3465, + 3466, + 3467, + 3468, + 3469 + ], + "spine1": [ + 598, + 599, + 600, + 601, + 610, + 611, + 612, + 613, + 614, + 615, + 616, + 617, + 618, + 619, + 620, + 621, + 642, + 645, + 646, + 647, + 652, + 653, + 658, + 659, + 660, + 661, + 668, + 669, + 670, + 671, + 684, + 685, + 686, + 687, + 688, + 689, + 690, + 691, + 692, + 722, + 723, + 724, + 725, + 736, + 750, + 751, + 761, + 764, + 766, + 767, + 794, + 795, + 891, + 892, + 893, + 894, + 925, + 926, + 927, + 928, + 929, + 940, + 941, + 942, + 943, + 1190, + 1191, + 1192, + 1193, + 1194, + 1195, + 1196, + 1197, + 1200, + 1201, + 1202, + 1212, + 1236, + 1252, + 1253, + 1254, + 1255, + 1268, + 1269, + 1270, + 1329, + 1330, + 1348, + 1349, + 1351, + 1420, + 1421, + 1423, + 1424, + 1425, + 1426, + 1436, + 1437, + 1756, + 1757, + 1758, + 2839, + 2840, + 2841, + 2842, + 2843, + 2844, + 2845, + 2846, + 2847, + 2848, + 2849, + 2850, + 2851, + 2870, + 2871, + 2883, + 2906, + 2908, + 3014, + 3017, + 3025, + 3030, + 3033, + 3034, + 3037, + 3039, + 3040, + 3041, + 3042, + 3043, + 3044, + 3076, + 3077, + 3079, + 3480, + 3505, + 3511, + 4086, + 4087, + 4088, + 4089, + 4098, + 4099, + 4100, + 4101, + 4102, + 4103, + 4104, + 4105, + 4106, + 4107, + 4108, + 4109, + 4130, + 4131, + 4134, + 4135, + 4140, + 4141, + 4146, + 4147, + 4148, + 4149, + 4156, + 4157, + 4158, + 4159, + 4172, + 4173, + 4174, + 4175, + 4176, + 4177, + 4178, + 4179, + 4180, + 4210, + 4211, + 4212, + 4213, + 4225, + 4239, + 4240, + 4249, + 4250, + 4255, + 4256, + 4282, + 4283, + 4377, + 4378, + 4379, + 4380, + 4411, + 4412, + 4413, + 4414, + 4415, + 4426, + 4427, + 4428, + 4429, + 4676, + 4677, + 4678, + 4679, + 4680, + 4681, + 4682, + 4683, + 4686, + 4687, + 4688, + 4695, + 4719, + 4735, + 4736, + 4737, + 4740, + 4751, + 4752, + 4753, + 4824, + 4825, + 4828, + 4893, + 4894, + 4895, + 4897, + 4898, + 4899, + 4908, + 4909, + 5223, + 5224, + 5225, + 6300, + 6301, + 6302, + 6303, + 6304, + 6305, + 6306, + 6307, + 6308, + 6309, + 6310, + 6311, + 6312, + 6331, + 6332, + 6342, + 6366, + 6367, + 6475, + 6477, + 6478, + 6481, + 6482, + 6485, + 6487, + 6488, + 6489, + 6490, + 6491, + 6878 + ], + "spine2": [ + 570, + 571, + 572, + 573, + 584, + 585, + 586, + 587, + 588, + 589, + 590, + 591, + 592, + 593, + 594, + 595, + 596, + 597, + 602, + 603, + 604, + 605, + 606, + 607, + 608, + 609, + 622, + 623, + 624, + 625, + 638, + 639, + 640, + 641, + 643, + 644, + 648, + 649, + 650, + 651, + 666, + 667, + 672, + 673, + 674, + 675, + 680, + 681, + 682, + 683, + 693, + 694, + 695, + 696, + 697, + 698, + 699, + 700, + 701, + 702, + 703, + 704, + 713, + 714, + 715, + 716, + 717, + 726, + 727, + 728, + 729, + 730, + 731, + 732, + 733, + 735, + 737, + 738, + 739, + 740, + 741, + 742, + 743, + 744, + 745, + 746, + 747, + 748, + 749, + 752, + 753, + 754, + 755, + 756, + 757, + 758, + 759, + 760, + 762, + 763, + 803, + 804, + 805, + 806, + 811, + 812, + 813, + 814, + 817, + 818, + 819, + 820, + 821, + 824, + 825, + 826, + 827, + 828, + 895, + 896, + 930, + 931, + 1198, + 1199, + 1213, + 1214, + 1215, + 1216, + 1217, + 1218, + 1219, + 1220, + 1235, + 1237, + 1256, + 1257, + 1271, + 1272, + 1273, + 1279, + 1280, + 1283, + 1284, + 1285, + 1286, + 1287, + 1288, + 1289, + 1290, + 1291, + 1292, + 1293, + 1294, + 1295, + 1296, + 1297, + 1298, + 1299, + 1300, + 1301, + 1302, + 1303, + 1304, + 1305, + 1306, + 1307, + 1308, + 1309, + 1312, + 1313, + 1319, + 1320, + 1346, + 1347, + 1350, + 1352, + 1401, + 1417, + 1418, + 1419, + 1422, + 1427, + 1434, + 1435, + 1503, + 1504, + 1536, + 1537, + 1544, + 1545, + 1753, + 1754, + 1755, + 1759, + 1760, + 1761, + 1762, + 1763, + 1808, + 1809, + 1810, + 1811, + 1816, + 1817, + 1818, + 1819, + 1820, + 1834, + 1835, + 1836, + 1837, + 1838, + 1839, + 1868, + 1879, + 1880, + 2812, + 2813, + 2852, + 2853, + 2854, + 2855, + 2856, + 2857, + 2858, + 2859, + 2860, + 2861, + 2862, + 2863, + 2864, + 2865, + 2866, + 2867, + 2868, + 2869, + 2872, + 2875, + 2876, + 2877, + 2878, + 2881, + 2882, + 2884, + 2885, + 2886, + 2904, + 2905, + 2907, + 2931, + 2932, + 2933, + 2934, + 2935, + 2936, + 2937, + 2941, + 2950, + 2951, + 2952, + 2953, + 2954, + 2955, + 2956, + 2957, + 2958, + 2959, + 2960, + 2961, + 2962, + 2963, + 2964, + 2965, + 2966, + 2967, + 2968, + 2969, + 2970, + 2971, + 2972, + 2973, + 2997, + 2998, + 3006, + 3007, + 3012, + 3015, + 3026, + 3027, + 3028, + 3029, + 3031, + 3032, + 3035, + 3036, + 3038, + 3059, + 3060, + 3061, + 3062, + 3063, + 3064, + 3065, + 3066, + 3067, + 3073, + 3074, + 3075, + 3078, + 3168, + 3169, + 3171, + 3470, + 3471, + 3482, + 3483, + 3495, + 3496, + 3497, + 3498, + 3506, + 3508, + 4058, + 4059, + 4060, + 4061, + 4072, + 4073, + 4074, + 4075, + 4076, + 4077, + 4078, + 4079, + 4080, + 4081, + 4082, + 4083, + 4084, + 4085, + 4090, + 4091, + 4092, + 4093, + 4094, + 4095, + 4096, + 4097, + 4110, + 4111, + 4112, + 4113, + 4126, + 4127, + 4128, + 4129, + 4132, + 4133, + 4136, + 4137, + 4138, + 4139, + 4154, + 4155, + 4160, + 4161, + 4162, + 4163, + 4168, + 4169, + 4170, + 4171, + 4181, + 4182, + 4183, + 4184, + 4185, + 4186, + 4187, + 4188, + 4189, + 4190, + 4191, + 4192, + 4201, + 4202, + 4203, + 4204, + 4207, + 4214, + 4215, + 4216, + 4217, + 4218, + 4219, + 4220, + 4221, + 4223, + 4224, + 4226, + 4227, + 4228, + 4229, + 4230, + 4231, + 4232, + 4233, + 4234, + 4235, + 4236, + 4237, + 4238, + 4241, + 4242, + 4243, + 4244, + 4245, + 4246, + 4247, + 4248, + 4251, + 4252, + 4291, + 4292, + 4293, + 4294, + 4299, + 4300, + 4301, + 4302, + 4305, + 4306, + 4307, + 4308, + 4309, + 4312, + 4313, + 4314, + 4315, + 4381, + 4382, + 4416, + 4417, + 4684, + 4685, + 4696, + 4697, + 4698, + 4699, + 4700, + 4701, + 4702, + 4703, + 4718, + 4720, + 4738, + 4739, + 4754, + 4755, + 4756, + 4761, + 4762, + 4765, + 4766, + 4767, + 4768, + 4769, + 4770, + 4771, + 4772, + 4773, + 4774, + 4775, + 4776, + 4777, + 4778, + 4779, + 4780, + 4781, + 4782, + 4783, + 4784, + 4785, + 4786, + 4787, + 4788, + 4789, + 4792, + 4793, + 4799, + 4800, + 4822, + 4823, + 4826, + 4827, + 4874, + 4890, + 4891, + 4892, + 4896, + 4900, + 4907, + 4910, + 4975, + 4976, + 5007, + 5008, + 5013, + 5014, + 5222, + 5226, + 5227, + 5228, + 5229, + 5230, + 5269, + 5270, + 5271, + 5272, + 5277, + 5278, + 5279, + 5280, + 5281, + 5295, + 5296, + 5297, + 5298, + 5299, + 5300, + 5329, + 5340, + 5341, + 6273, + 6274, + 6313, + 6314, + 6315, + 6316, + 6317, + 6318, + 6319, + 6320, + 6321, + 6322, + 6323, + 6324, + 6325, + 6326, + 6327, + 6328, + 6329, + 6330, + 6333, + 6336, + 6337, + 6340, + 6341, + 6343, + 6344, + 6345, + 6363, + 6364, + 6365, + 6390, + 6391, + 6392, + 6393, + 6394, + 6395, + 6396, + 6398, + 6409, + 6410, + 6411, + 6412, + 6413, + 6414, + 6415, + 6416, + 6417, + 6418, + 6419, + 6420, + 6421, + 6422, + 6423, + 6424, + 6425, + 6426, + 6427, + 6428, + 6429, + 6430, + 6431, + 6432, + 6456, + 6457, + 6465, + 6466, + 6476, + 6479, + 6480, + 6483, + 6484, + 6486, + 6496, + 6497, + 6498, + 6499, + 6500, + 6501, + 6502, + 6503, + 6879 + ], + "leftShoulder": [ + 591, + 604, + 605, + 606, + 609, + 634, + 635, + 636, + 637, + 674, + 706, + 707, + 708, + 709, + 710, + 711, + 712, + 713, + 715, + 717, + 730, + 733, + 734, + 735, + 781, + 782, + 783, + 1238, + 1239, + 1240, + 1241, + 1242, + 1243, + 1244, + 1245, + 1290, + 1291, + 1294, + 1316, + 1317, + 1318, + 1401, + 1402, + 1403, + 1404, + 1509, + 1535, + 1545, + 1808, + 1810, + 1811, + 1812, + 1813, + 1814, + 1815, + 1818, + 1819, + 1821, + 1822, + 1823, + 1824, + 1825, + 1826, + 1827, + 1828, + 1829, + 1830, + 1831, + 1832, + 1833, + 1837, + 1840, + 1841, + 1842, + 1843, + 1844, + 1845, + 1846, + 1847, + 1848, + 1849, + 1850, + 1851, + 1852, + 1853, + 1854, + 1855, + 1856, + 1857, + 1858, + 1859, + 1861, + 1862, + 1863, + 1864, + 1872, + 1873, + 1880, + 1881, + 1884, + 1885, + 1886, + 1887, + 1890, + 1891, + 1893, + 1894, + 1895, + 1896, + 1897, + 1898, + 1899, + 2879, + 2880, + 2881, + 2886, + 2887, + 2888, + 2889, + 2890, + 2891, + 2892, + 2893, + 2894, + 2903, + 2938, + 2939, + 2940, + 2941, + 2942, + 2943, + 2944, + 2945, + 2946, + 2947, + 2948, + 2949, + 2965, + 2967, + 2969, + 2999, + 3000, + 3001, + 3002, + 3003, + 3004, + 3005, + 3008, + 3009, + 3010, + 3011 + ], + "rightShoulder": [ + 4077, + 4091, + 4092, + 4094, + 4095, + 4122, + 4123, + 4124, + 4125, + 4162, + 4194, + 4195, + 4196, + 4197, + 4198, + 4199, + 4200, + 4201, + 4203, + 4207, + 4218, + 4219, + 4222, + 4223, + 4269, + 4270, + 4271, + 4721, + 4722, + 4723, + 4724, + 4725, + 4726, + 4727, + 4728, + 4773, + 4774, + 4778, + 4796, + 4797, + 4798, + 4874, + 4875, + 4876, + 4877, + 4982, + 5006, + 5014, + 5269, + 5271, + 5272, + 5273, + 5274, + 5275, + 5276, + 5279, + 5281, + 5282, + 5283, + 5284, + 5285, + 5286, + 5287, + 5288, + 5289, + 5290, + 5291, + 5292, + 5293, + 5294, + 5298, + 5301, + 5302, + 5303, + 5304, + 5305, + 5306, + 5307, + 5308, + 5309, + 5310, + 5311, + 5312, + 5313, + 5314, + 5315, + 5316, + 5317, + 5318, + 5319, + 5320, + 5322, + 5323, + 5324, + 5325, + 5333, + 5334, + 5341, + 5342, + 5345, + 5346, + 5347, + 5348, + 5351, + 5352, + 5354, + 5355, + 5356, + 5357, + 5358, + 5359, + 5360, + 6338, + 6339, + 6340, + 6345, + 6346, + 6347, + 6348, + 6349, + 6350, + 6351, + 6352, + 6353, + 6362, + 6397, + 6398, + 6399, + 6400, + 6401, + 6402, + 6403, + 6404, + 6405, + 6406, + 6407, + 6408, + 6424, + 6425, + 6428, + 6458, + 6459, + 6460, + 6461, + 6462, + 6463, + 6464, + 6467, + 6468, + 6469, + 6470 + ], + "rightFoot": [ + 6727, + 6728, + 6729, + 6730, + 6731, + 6732, + 6733, + 6734, + 6735, + 6736, + 6737, + 6738, + 6739, + 6740, + 6741, + 6742, + 6743, + 6744, + 6745, + 6746, + 6747, + 6748, + 6749, + 6750, + 6751, + 6752, + 6753, + 6754, + 6755, + 6756, + 6757, + 6758, + 6759, + 6760, + 6761, + 6762, + 6763, + 6764, + 6765, + 6766, + 6767, + 6768, + 6769, + 6770, + 6771, + 6772, + 6773, + 6774, + 6775, + 6776, + 6777, + 6778, + 6779, + 6780, + 6781, + 6782, + 6783, + 6784, + 6785, + 6786, + 6787, + 6788, + 6789, + 6790, + 6791, + 6792, + 6793, + 6794, + 6795, + 6796, + 6797, + 6798, + 6799, + 6800, + 6801, + 6802, + 6803, + 6804, + 6805, + 6806, + 6807, + 6808, + 6809, + 6810, + 6811, + 6812, + 6813, + 6814, + 6815, + 6816, + 6817, + 6818, + 6819, + 6820, + 6821, + 6822, + 6823, + 6824, + 6825, + 6826, + 6827, + 6828, + 6829, + 6830, + 6831, + 6832, + 6833, + 6834, + 6835, + 6836, + 6837, + 6838, + 6839, + 6840, + 6841, + 6842, + 6843, + 6844, + 6845, + 6846, + 6847, + 6848, + 6849, + 6850, + 6851, + 6852, + 6853, + 6854, + 6855, + 6856, + 6857, + 6858, + 6859, + 6860, + 6861, + 6862, + 6863, + 6864, + 6865, + 6866, + 6867, + 6868, + 6869 + ], + "head": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63, + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 83, + 84, + 85, + 86, + 87, + 88, + 89, + 90, + 91, + 92, + 93, + 94, + 95, + 96, + 97, + 98, + 99, + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109, + 110, + 111, + 112, + 113, + 114, + 115, + 116, + 117, + 118, + 119, + 120, + 121, + 122, + 123, + 124, + 125, + 126, + 127, + 128, + 129, + 130, + 131, + 132, + 133, + 134, + 135, + 136, + 137, + 138, + 139, + 140, + 141, + 142, + 143, + 144, + 145, + 146, + 147, + 148, + 149, + 154, + 155, + 156, + 157, + 158, + 159, + 160, + 161, + 162, + 163, + 164, + 165, + 166, + 167, + 168, + 169, + 170, + 171, + 172, + 173, + 176, + 177, + 178, + 179, + 180, + 181, + 182, + 183, + 184, + 185, + 186, + 187, + 188, + 189, + 190, + 191, + 192, + 193, + 194, + 195, + 196, + 197, + 198, + 199, + 200, + 201, + 202, + 203, + 204, + 205, + 220, + 221, + 225, + 226, + 227, + 228, + 229, + 230, + 231, + 232, + 233, + 234, + 235, + 236, + 237, + 238, + 239, + 240, + 241, + 242, + 243, + 244, + 245, + 246, + 247, + 248, + 249, + 250, + 251, + 252, + 253, + 254, + 255, + 258, + 259, + 260, + 261, + 262, + 263, + 264, + 265, + 266, + 267, + 268, + 269, + 270, + 271, + 272, + 273, + 274, + 275, + 276, + 277, + 278, + 279, + 280, + 281, + 282, + 283, + 286, + 287, + 288, + 289, + 290, + 291, + 292, + 293, + 294, + 295, + 303, + 304, + 306, + 307, + 310, + 311, + 312, + 313, + 314, + 315, + 316, + 317, + 318, + 319, + 320, + 321, + 322, + 323, + 324, + 325, + 326, + 327, + 328, + 329, + 330, + 331, + 332, + 335, + 336, + 337, + 338, + 339, + 340, + 341, + 342, + 343, + 344, + 345, + 346, + 347, + 348, + 349, + 350, + 351, + 352, + 353, + 354, + 355, + 356, + 357, + 358, + 359, + 360, + 361, + 362, + 363, + 364, + 365, + 366, + 367, + 368, + 369, + 370, + 371, + 372, + 373, + 374, + 375, + 376, + 377, + 378, + 379, + 380, + 381, + 382, + 383, + 384, + 385, + 386, + 387, + 388, + 389, + 390, + 391, + 392, + 393, + 394, + 395, + 396, + 397, + 398, + 399, + 400, + 401, + 402, + 403, + 404, + 405, + 406, + 407, + 408, + 409, + 410, + 411, + 412, + 413, + 414, + 415, + 416, + 417, + 418, + 419, + 420, + 421, + 422, + 427, + 428, + 429, + 430, + 431, + 432, + 433, + 434, + 435, + 436, + 437, + 438, + 439, + 442, + 443, + 444, + 445, + 446, + 447, + 448, + 449, + 450, + 454, + 455, + 456, + 457, + 458, + 459, + 461, + 462, + 463, + 464, + 465, + 466, + 467, + 468, + 469, + 470, + 471, + 472, + 473, + 474, + 475, + 476, + 477, + 478, + 479, + 480, + 481, + 482, + 483, + 484, + 485, + 486, + 487, + 488, + 489, + 490, + 491, + 492, + 493, + 494, + 495, + 496, + 497, + 498, + 499, + 500, + 501, + 502, + 503, + 504, + 505, + 506, + 507, + 508, + 509, + 510, + 511, + 512, + 513, + 514, + 515, + 516, + 517, + 518, + 519, + 520, + 521, + 522, + 523, + 524, + 525, + 526, + 527, + 528, + 529, + 530, + 531, + 532, + 533, + 534, + 535, + 536, + 537, + 538, + 539, + 540, + 541, + 542, + 543, + 544, + 545, + 546, + 547, + 548, + 549, + 550, + 551, + 552, + 553, + 554, + 555, + 556, + 557, + 558, + 559, + 560, + 561, + 562, + 563, + 564, + 565, + 566, + 567, + 568, + 569, + 574, + 575, + 576, + 577, + 578, + 579, + 580, + 581, + 582, + 583, + 1764, + 1765, + 1766, + 1770, + 1771, + 1772, + 1773, + 1774, + 1775, + 1776, + 1777, + 1778, + 1905, + 1906, + 1907, + 1908, + 2779, + 2780, + 2781, + 2782, + 2783, + 2784, + 2785, + 2786, + 2787, + 2788, + 2789, + 2790, + 2791, + 2792, + 2793, + 2794, + 2795, + 2796, + 2797, + 2798, + 2799, + 2800, + 2801, + 2802, + 2803, + 2804, + 2805, + 2806, + 2807, + 2808, + 2809, + 2810, + 2811, + 2814, + 2815, + 2816, + 2817, + 2818, + 3045, + 3046, + 3047, + 3048, + 3051, + 3052, + 3053, + 3054, + 3055, + 3056, + 3058, + 3069, + 3070, + 3071, + 3072, + 3161, + 3162, + 3163, + 3165, + 3166, + 3167, + 3485, + 3486, + 3487, + 3488, + 3489, + 3490, + 3491, + 3492, + 3493, + 3494, + 3499, + 3512, + 3513, + 3514, + 3515, + 3516, + 3517, + 3518, + 3519, + 3520, + 3521, + 3522, + 3523, + 3524, + 3525, + 3526, + 3527, + 3528, + 3529, + 3530, + 3531, + 3532, + 3533, + 3534, + 3535, + 3536, + 3537, + 3538, + 3539, + 3540, + 3541, + 3542, + 3543, + 3544, + 3545, + 3546, + 3547, + 3548, + 3549, + 3550, + 3551, + 3552, + 3553, + 3554, + 3555, + 3556, + 3557, + 3558, + 3559, + 3560, + 3561, + 3562, + 3563, + 3564, + 3565, + 3566, + 3567, + 3568, + 3569, + 3570, + 3571, + 3572, + 3573, + 3574, + 3575, + 3576, + 3577, + 3578, + 3579, + 3580, + 3581, + 3582, + 3583, + 3584, + 3585, + 3586, + 3587, + 3588, + 3589, + 3590, + 3591, + 3592, + 3593, + 3594, + 3595, + 3596, + 3597, + 3598, + 3599, + 3600, + 3601, + 3602, + 3603, + 3604, + 3605, + 3606, + 3607, + 3608, + 3609, + 3610, + 3611, + 3612, + 3613, + 3614, + 3615, + 3616, + 3617, + 3618, + 3619, + 3620, + 3621, + 3622, + 3623, + 3624, + 3625, + 3626, + 3627, + 3628, + 3629, + 3630, + 3631, + 3632, + 3633, + 3634, + 3635, + 3636, + 3637, + 3638, + 3639, + 3640, + 3641, + 3642, + 3643, + 3644, + 3645, + 3646, + 3647, + 3648, + 3649, + 3650, + 3651, + 3652, + 3653, + 3654, + 3655, + 3656, + 3657, + 3658, + 3659, + 3660, + 3661, + 3666, + 3667, + 3668, + 3669, + 3670, + 3671, + 3672, + 3673, + 3674, + 3675, + 3676, + 3677, + 3678, + 3679, + 3680, + 3681, + 3682, + 3683, + 3684, + 3685, + 3688, + 3689, + 3690, + 3691, + 3692, + 3693, + 3694, + 3695, + 3696, + 3697, + 3698, + 3699, + 3700, + 3701, + 3702, + 3703, + 3704, + 3705, + 3706, + 3707, + 3708, + 3709, + 3710, + 3711, + 3712, + 3713, + 3714, + 3715, + 3716, + 3717, + 3732, + 3733, + 3737, + 3738, + 3739, + 3740, + 3741, + 3742, + 3743, + 3744, + 3745, + 3746, + 3747, + 3748, + 3749, + 3750, + 3751, + 3752, + 3753, + 3754, + 3755, + 3756, + 3757, + 3758, + 3759, + 3760, + 3761, + 3762, + 3763, + 3764, + 3765, + 3766, + 3767, + 3770, + 3771, + 3772, + 3773, + 3774, + 3775, + 3776, + 3777, + 3778, + 3779, + 3780, + 3781, + 3782, + 3783, + 3784, + 3785, + 3786, + 3787, + 3788, + 3789, + 3790, + 3791, + 3792, + 3793, + 3794, + 3795, + 3798, + 3799, + 3800, + 3801, + 3802, + 3803, + 3804, + 3805, + 3806, + 3807, + 3815, + 3816, + 3819, + 3820, + 3821, + 3822, + 3823, + 3824, + 3825, + 3826, + 3827, + 3828, + 3829, + 3830, + 3831, + 3832, + 3833, + 3834, + 3835, + 3836, + 3837, + 3838, + 3841, + 3842, + 3843, + 3844, + 3845, + 3846, + 3847, + 3848, + 3849, + 3850, + 3851, + 3852, + 3853, + 3854, + 3855, + 3856, + 3857, + 3858, + 3859, + 3860, + 3861, + 3862, + 3863, + 3864, + 3865, + 3866, + 3867, + 3868, + 3869, + 3870, + 3871, + 3872, + 3873, + 3874, + 3875, + 3876, + 3877, + 3878, + 3879, + 3880, + 3881, + 3882, + 3883, + 3884, + 3885, + 3886, + 3887, + 3888, + 3889, + 3890, + 3891, + 3892, + 3893, + 3894, + 3895, + 3896, + 3897, + 3898, + 3899, + 3900, + 3901, + 3902, + 3903, + 3904, + 3905, + 3906, + 3907, + 3908, + 3909, + 3910, + 3911, + 3912, + 3913, + 3914, + 3915, + 3916, + 3917, + 3922, + 3923, + 3924, + 3925, + 3926, + 3927, + 3928, + 3929, + 3930, + 3931, + 3932, + 3933, + 3936, + 3937, + 3938, + 3939, + 3940, + 3941, + 3945, + 3946, + 3947, + 3948, + 3949, + 3950, + 3951, + 3952, + 3953, + 3954, + 3955, + 3956, + 3957, + 3958, + 3959, + 3960, + 3961, + 3962, + 3963, + 3964, + 3965, + 3966, + 3967, + 3968, + 3969, + 3970, + 3971, + 3972, + 3973, + 3974, + 3975, + 3976, + 3977, + 3978, + 3979, + 3980, + 3981, + 3982, + 3983, + 3984, + 3985, + 3986, + 3987, + 3988, + 3989, + 3990, + 3991, + 3992, + 3993, + 3994, + 3995, + 3996, + 3997, + 3998, + 3999, + 4000, + 4001, + 4002, + 4003, + 4004, + 4005, + 4006, + 4007, + 4008, + 4009, + 4010, + 4011, + 4012, + 4013, + 4014, + 4015, + 4016, + 4017, + 4018, + 4019, + 4020, + 4021, + 4022, + 4023, + 4024, + 4025, + 4026, + 4027, + 4028, + 4029, + 4030, + 4031, + 4032, + 4033, + 4034, + 4035, + 4036, + 4037, + 4038, + 4039, + 4040, + 4041, + 4042, + 4043, + 4044, + 4045, + 4046, + 4047, + 4048, + 4049, + 4050, + 4051, + 4052, + 4053, + 4054, + 4055, + 4056, + 4057, + 4062, + 4063, + 4064, + 4065, + 4066, + 4067, + 4068, + 4069, + 4070, + 4071, + 5231, + 5232, + 5233, + 5235, + 5236, + 5237, + 5238, + 5239, + 5240, + 5241, + 5242, + 5243, + 5366, + 5367, + 5368, + 5369, + 6240, + 6241, + 6242, + 6243, + 6244, + 6245, + 6246, + 6247, + 6248, + 6249, + 6250, + 6251, + 6252, + 6253, + 6254, + 6255, + 6256, + 6257, + 6258, + 6259, + 6260, + 6261, + 6262, + 6263, + 6264, + 6265, + 6266, + 6267, + 6268, + 6269, + 6270, + 6271, + 6272, + 6275, + 6276, + 6277, + 6278, + 6279, + 6492, + 6493, + 6494, + 6495, + 6880, + 6881, + 6882, + 6883, + 6884, + 6885, + 6886, + 6887, + 6888, + 6889 + ], + "rightArm": [ + 4114, + 4115, + 4116, + 4117, + 4122, + 4125, + 4168, + 4171, + 4204, + 4205, + 4206, + 4207, + 4257, + 4258, + 4259, + 4260, + 4261, + 4262, + 4263, + 4264, + 4265, + 4266, + 4267, + 4268, + 4272, + 4273, + 4274, + 4275, + 4276, + 4277, + 4278, + 4279, + 4280, + 4281, + 4714, + 4715, + 4716, + 4717, + 4741, + 4742, + 4743, + 4744, + 4756, + 4763, + 4764, + 4790, + 4791, + 4794, + 4795, + 4816, + 4817, + 4818, + 4819, + 4830, + 4831, + 4832, + 4833, + 4849, + 4850, + 4851, + 4852, + 4853, + 4854, + 4855, + 4856, + 4857, + 4858, + 4859, + 4860, + 4861, + 4862, + 4863, + 4864, + 4865, + 4866, + 4867, + 4868, + 4869, + 4870, + 4871, + 4872, + 4873, + 4876, + 4877, + 4878, + 4879, + 4880, + 4881, + 4882, + 4883, + 4884, + 4885, + 4886, + 4887, + 4888, + 4889, + 4901, + 4902, + 4903, + 4904, + 4905, + 4906, + 4911, + 4912, + 4913, + 4914, + 4915, + 4916, + 4917, + 4918, + 4974, + 4977, + 4978, + 4979, + 4980, + 4981, + 4982, + 5009, + 5010, + 5011, + 5012, + 5014, + 5088, + 5089, + 5090, + 5091, + 5100, + 5101, + 5102, + 5103, + 5104, + 5105, + 5106, + 5107, + 5108, + 5109, + 5110, + 5111, + 5114, + 5115, + 5116, + 5117, + 5118, + 5119, + 5120, + 5121, + 5122, + 5123, + 5124, + 5125, + 5128, + 5129, + 5130, + 5131, + 5134, + 5135, + 5136, + 5137, + 5138, + 5139, + 5140, + 5141, + 5142, + 5143, + 5144, + 5145, + 5146, + 5147, + 5148, + 5149, + 5150, + 5151, + 5152, + 5153, + 5165, + 5166, + 5167, + 5172, + 5173, + 5174, + 5175, + 5176, + 5177, + 5178, + 5179, + 5180, + 5181, + 5182, + 5183, + 5184, + 5185, + 5186, + 5187, + 5188, + 5189, + 5194, + 5200, + 5201, + 5202, + 5203, + 5204, + 5206, + 5208, + 5209, + 5214, + 5215, + 5216, + 5217, + 5218, + 5220, + 5229, + 5292, + 5293, + 5303, + 5306, + 5309, + 5311, + 5314, + 5315, + 5318, + 5319, + 5321, + 5326, + 5327, + 5328, + 5330, + 5331, + 5332, + 5335, + 5336, + 5337, + 5338, + 5339, + 5343, + 5344, + 5349, + 5350, + 5353, + 5361, + 5362, + 5363, + 5364, + 5365, + 5370, + 6280, + 6281, + 6282, + 6283, + 6354, + 6355, + 6356, + 6357, + 6358, + 6359, + 6360, + 6361, + 6362, + 6404, + 6405, + 6433, + 6434, + 6435, + 6436, + 6437, + 6438, + 6439, + 6440, + 6441, + 6442, + 6443, + 6444, + 6445, + 6446, + 6447, + 6448, + 6449, + 6450, + 6451, + 6452, + 6453, + 6454, + 6455, + 6461, + 6471 + ], + "leftHandIndex1": [ + 2027, + 2028, + 2029, + 2030, + 2037, + 2038, + 2039, + 2040, + 2057, + 2067, + 2068, + 2123, + 2124, + 2125, + 2126, + 2127, + 2128, + 2129, + 2130, + 2132, + 2145, + 2146, + 2152, + 2153, + 2154, + 2156, + 2157, + 2158, + 2159, + 2160, + 2161, + 2162, + 2163, + 2164, + 2165, + 2166, + 2167, + 2168, + 2169, + 2177, + 2178, + 2179, + 2181, + 2186, + 2187, + 2190, + 2191, + 2204, + 2205, + 2215, + 2216, + 2217, + 2218, + 2219, + 2220, + 2232, + 2233, + 2245, + 2246, + 2247, + 2258, + 2259, + 2261, + 2262, + 2263, + 2269, + 2270, + 2272, + 2273, + 2274, + 2276, + 2277, + 2280, + 2281, + 2282, + 2283, + 2291, + 2292, + 2293, + 2294, + 2295, + 2296, + 2297, + 2298, + 2299, + 2300, + 2301, + 2302, + 2303, + 2304, + 2305, + 2306, + 2307, + 2308, + 2309, + 2310, + 2311, + 2312, + 2313, + 2314, + 2315, + 2316, + 2317, + 2318, + 2319, + 2320, + 2321, + 2322, + 2323, + 2324, + 2325, + 2326, + 2327, + 2328, + 2329, + 2330, + 2331, + 2332, + 2333, + 2334, + 2335, + 2336, + 2337, + 2338, + 2339, + 2340, + 2341, + 2342, + 2343, + 2344, + 2345, + 2346, + 2347, + 2348, + 2349, + 2350, + 2351, + 2352, + 2353, + 2354, + 2355, + 2356, + 2357, + 2358, + 2359, + 2360, + 2361, + 2362, + 2363, + 2364, + 2365, + 2366, + 2367, + 2368, + 2369, + 2370, + 2371, + 2372, + 2373, + 2374, + 2375, + 2376, + 2377, + 2378, + 2379, + 2380, + 2381, + 2382, + 2383, + 2384, + 2385, + 2386, + 2387, + 2388, + 2389, + 2390, + 2391, + 2392, + 2393, + 2394, + 2395, + 2396, + 2397, + 2398, + 2399, + 2400, + 2401, + 2402, + 2403, + 2404, + 2405, + 2406, + 2407, + 2408, + 2409, + 2410, + 2411, + 2412, + 2413, + 2414, + 2415, + 2416, + 2417, + 2418, + 2419, + 2420, + 2421, + 2422, + 2423, + 2424, + 2425, + 2426, + 2427, + 2428, + 2429, + 2430, + 2431, + 2432, + 2433, + 2434, + 2435, + 2436, + 2437, + 2438, + 2439, + 2440, + 2441, + 2442, + 2443, + 2444, + 2445, + 2446, + 2447, + 2448, + 2449, + 2450, + 2451, + 2452, + 2453, + 2454, + 2455, + 2456, + 2457, + 2458, + 2459, + 2460, + 2461, + 2462, + 2463, + 2464, + 2465, + 2466, + 2467, + 2468, + 2469, + 2470, + 2471, + 2472, + 2473, + 2474, + 2475, + 2476, + 2477, + 2478, + 2479, + 2480, + 2481, + 2482, + 2483, + 2484, + 2485, + 2486, + 2487, + 2488, + 2489, + 2490, + 2491, + 2492, + 2493, + 2494, + 2495, + 2496, + 2497, + 2498, + 2499, + 2500, + 2501, + 2502, + 2503, + 2504, + 2505, + 2506, + 2507, + 2508, + 2509, + 2510, + 2511, + 2512, + 2513, + 2514, + 2515, + 2516, + 2517, + 2518, + 2519, + 2520, + 2521, + 2522, + 2523, + 2524, + 2525, + 2526, + 2527, + 2528, + 2529, + 2530, + 2531, + 2532, + 2533, + 2534, + 2535, + 2536, + 2537, + 2538, + 2539, + 2540, + 2541, + 2542, + 2543, + 2544, + 2545, + 2546, + 2547, + 2548, + 2549, + 2550, + 2551, + 2552, + 2553, + 2554, + 2555, + 2556, + 2557, + 2558, + 2559, + 2560, + 2561, + 2562, + 2563, + 2564, + 2565, + 2566, + 2567, + 2568, + 2569, + 2570, + 2571, + 2572, + 2573, + 2574, + 2575, + 2576, + 2577, + 2578, + 2579, + 2580, + 2581, + 2582, + 2583, + 2584, + 2585, + 2586, + 2587, + 2588, + 2589, + 2590, + 2591, + 2592, + 2593, + 2594, + 2596, + 2597, + 2599, + 2600, + 2601, + 2602, + 2603, + 2604, + 2606, + 2607, + 2609, + 2610, + 2611, + 2612, + 2613, + 2614, + 2615, + 2616, + 2617, + 2618, + 2619, + 2620, + 2621, + 2622, + 2623, + 2624, + 2625, + 2626, + 2627, + 2628, + 2629, + 2630, + 2631, + 2632, + 2633, + 2634, + 2635, + 2636, + 2637, + 2638, + 2639, + 2640, + 2641, + 2642, + 2643, + 2644, + 2645, + 2646, + 2647, + 2648, + 2649, + 2650, + 2651, + 2652, + 2653, + 2654, + 2655, + 2656, + 2657, + 2658, + 2659, + 2660, + 2661, + 2662, + 2663, + 2664, + 2665, + 2666, + 2667, + 2668, + 2669, + 2670, + 2671, + 2672, + 2673, + 2674, + 2675, + 2676, + 2677, + 2678, + 2679, + 2680, + 2681, + 2682, + 2683, + 2684, + 2685, + 2686, + 2687, + 2688, + 2689, + 2690, + 2691, + 2692, + 2693, + 2694, + 2695, + 2696 + ], + "rightLeg": [ + 4481, + 4482, + 4485, + 4486, + 4491, + 4492, + 4493, + 4495, + 4498, + 4500, + 4501, + 4505, + 4506, + 4529, + 4532, + 4533, + 4534, + 4535, + 4536, + 4537, + 4538, + 4539, + 4540, + 4541, + 4542, + 4543, + 4544, + 4545, + 4546, + 4547, + 4548, + 4549, + 4550, + 4551, + 4552, + 4553, + 4554, + 4555, + 4556, + 4557, + 4558, + 4559, + 4560, + 4561, + 4562, + 4563, + 4564, + 4565, + 4566, + 4567, + 4568, + 4569, + 4570, + 4571, + 4572, + 4573, + 4574, + 4575, + 4576, + 4577, + 4578, + 4579, + 4580, + 4581, + 4582, + 4583, + 4584, + 4585, + 4586, + 4587, + 4588, + 4589, + 4590, + 4591, + 4592, + 4593, + 4594, + 4595, + 4596, + 4597, + 4598, + 4599, + 4600, + 4601, + 4602, + 4603, + 4604, + 4605, + 4606, + 4607, + 4608, + 4609, + 4610, + 4611, + 4612, + 4613, + 4614, + 4615, + 4616, + 4617, + 4618, + 4619, + 4620, + 4621, + 4622, + 4634, + 4635, + 4636, + 4637, + 4638, + 4639, + 4640, + 4641, + 4642, + 4643, + 4644, + 4661, + 4662, + 4663, + 4664, + 4665, + 4666, + 4667, + 4668, + 4669, + 4842, + 4843, + 4844, + 4845, + 4846, + 4847, + 4848, + 4937, + 4938, + 4939, + 4940, + 4941, + 4942, + 4943, + 4944, + 4945, + 4946, + 4947, + 4993, + 4994, + 4995, + 4996, + 4997, + 4998, + 4999, + 5000, + 5001, + 5002, + 5003, + 6574, + 6575, + 6576, + 6577, + 6578, + 6579, + 6580, + 6581, + 6582, + 6583, + 6584, + 6585, + 6586, + 6587, + 6588, + 6589, + 6590, + 6591, + 6592, + 6593, + 6594, + 6595, + 6596, + 6597, + 6598, + 6599, + 6600, + 6601, + 6602, + 6603, + 6604, + 6605, + 6606, + 6607, + 6608, + 6609, + 6610, + 6719, + 6720, + 6721, + 6722, + 6723, + 6724, + 6725, + 6726, + 6727, + 6728, + 6729, + 6730, + 6731, + 6732, + 6733, + 6734, + 6735, + 6832, + 6833, + 6834, + 6835, + 6836, + 6869, + 6870, + 6871, + 6872 + ], + "rightHandIndex1": [ + 5488, + 5489, + 5490, + 5491, + 5498, + 5499, + 5500, + 5501, + 5518, + 5528, + 5529, + 5584, + 5585, + 5586, + 5587, + 5588, + 5589, + 5590, + 5591, + 5592, + 5606, + 5607, + 5613, + 5615, + 5616, + 5617, + 5618, + 5619, + 5620, + 5621, + 5622, + 5623, + 5624, + 5625, + 5626, + 5627, + 5628, + 5629, + 5630, + 5638, + 5639, + 5640, + 5642, + 5647, + 5648, + 5650, + 5651, + 5665, + 5666, + 5676, + 5677, + 5678, + 5679, + 5680, + 5681, + 5693, + 5694, + 5706, + 5707, + 5708, + 5719, + 5721, + 5722, + 5723, + 5724, + 5730, + 5731, + 5733, + 5734, + 5735, + 5737, + 5738, + 5741, + 5742, + 5743, + 5744, + 5752, + 5753, + 5754, + 5755, + 5756, + 5757, + 5758, + 5759, + 5760, + 5761, + 5762, + 5763, + 5764, + 5765, + 5766, + 5767, + 5768, + 5769, + 5770, + 5771, + 5772, + 5773, + 5774, + 5775, + 5776, + 5777, + 5778, + 5779, + 5780, + 5781, + 5782, + 5783, + 5784, + 5785, + 5786, + 5787, + 5788, + 5789, + 5790, + 5791, + 5792, + 5793, + 5794, + 5795, + 5796, + 5797, + 5798, + 5799, + 5800, + 5801, + 5802, + 5803, + 5804, + 5805, + 5806, + 5807, + 5808, + 5809, + 5810, + 5811, + 5812, + 5813, + 5814, + 5815, + 5816, + 5817, + 5818, + 5819, + 5820, + 5821, + 5822, + 5823, + 5824, + 5825, + 5826, + 5827, + 5828, + 5829, + 5830, + 5831, + 5832, + 5833, + 5834, + 5835, + 5836, + 5837, + 5838, + 5839, + 5840, + 5841, + 5842, + 5843, + 5844, + 5845, + 5846, + 5847, + 5848, + 5849, + 5850, + 5851, + 5852, + 5853, + 5854, + 5855, + 5856, + 5857, + 5858, + 5859, + 5860, + 5861, + 5862, + 5863, + 5864, + 5865, + 5866, + 5867, + 5868, + 5869, + 5870, + 5871, + 5872, + 5873, + 5874, + 5875, + 5876, + 5877, + 5878, + 5879, + 5880, + 5881, + 5882, + 5883, + 5884, + 5885, + 5886, + 5887, + 5888, + 5889, + 5890, + 5891, + 5892, + 5893, + 5894, + 5895, + 5896, + 5897, + 5898, + 5899, + 5900, + 5901, + 5902, + 5903, + 5904, + 5905, + 5906, + 5907, + 5908, + 5909, + 5910, + 5911, + 5912, + 5913, + 5914, + 5915, + 5916, + 5917, + 5918, + 5919, + 5920, + 5921, + 5922, + 5923, + 5924, + 5925, + 5926, + 5927, + 5928, + 5929, + 5930, + 5931, + 5932, + 5933, + 5934, + 5935, + 5936, + 5937, + 5938, + 5939, + 5940, + 5941, + 5942, + 5943, + 5944, + 5945, + 5946, + 5947, + 5948, + 5949, + 5950, + 5951, + 5952, + 5953, + 5954, + 5955, + 5956, + 5957, + 5958, + 5959, + 5960, + 5961, + 5962, + 5963, + 5964, + 5965, + 5966, + 5967, + 5968, + 5969, + 5970, + 5971, + 5972, + 5973, + 5974, + 5975, + 5976, + 5977, + 5978, + 5979, + 5980, + 5981, + 5982, + 5983, + 5984, + 5985, + 5986, + 5987, + 5988, + 5989, + 5990, + 5991, + 5992, + 5993, + 5994, + 5995, + 5996, + 5997, + 5998, + 5999, + 6000, + 6001, + 6002, + 6003, + 6004, + 6005, + 6006, + 6007, + 6008, + 6009, + 6010, + 6011, + 6012, + 6013, + 6014, + 6015, + 6016, + 6017, + 6018, + 6019, + 6020, + 6021, + 6022, + 6023, + 6024, + 6025, + 6026, + 6027, + 6028, + 6029, + 6030, + 6031, + 6032, + 6033, + 6034, + 6035, + 6036, + 6037, + 6038, + 6039, + 6040, + 6041, + 6042, + 6043, + 6044, + 6045, + 6046, + 6047, + 6048, + 6049, + 6050, + 6051, + 6052, + 6053, + 6054, + 6055, + 6058, + 6059, + 6060, + 6061, + 6062, + 6063, + 6064, + 6065, + 6068, + 6069, + 6070, + 6071, + 6072, + 6073, + 6074, + 6075, + 6076, + 6077, + 6078, + 6079, + 6080, + 6081, + 6082, + 6083, + 6084, + 6085, + 6086, + 6087, + 6088, + 6089, + 6090, + 6091, + 6092, + 6093, + 6094, + 6095, + 6096, + 6097, + 6098, + 6099, + 6100, + 6101, + 6102, + 6103, + 6104, + 6105, + 6106, + 6107, + 6108, + 6109, + 6110, + 6111, + 6112, + 6113, + 6114, + 6115, + 6116, + 6117, + 6118, + 6119, + 6120, + 6121, + 6122, + 6123, + 6124, + 6125, + 6126, + 6127, + 6128, + 6129, + 6130, + 6131, + 6132, + 6133, + 6134, + 6135, + 6136, + 6137, + 6138, + 6139, + 6140, + 6141, + 6142, + 6143, + 6144, + 6145, + 6146, + 6147, + 6148, + 6149, + 6150, + 6151, + 6152, + 6153, + 6154, + 6155, + 6156, + 6157 + ], + "leftForeArm": [ + 1546, + 1547, + 1548, + 1549, + 1550, + 1551, + 1552, + 1553, + 1554, + 1555, + 1556, + 1557, + 1558, + 1559, + 1560, + 1561, + 1562, + 1563, + 1564, + 1565, + 1566, + 1567, + 1568, + 1569, + 1570, + 1571, + 1572, + 1573, + 1574, + 1575, + 1576, + 1577, + 1578, + 1579, + 1580, + 1581, + 1582, + 1583, + 1584, + 1585, + 1586, + 1587, + 1588, + 1589, + 1590, + 1591, + 1592, + 1593, + 1594, + 1595, + 1596, + 1597, + 1598, + 1599, + 1600, + 1601, + 1602, + 1603, + 1604, + 1605, + 1606, + 1607, + 1608, + 1609, + 1610, + 1611, + 1612, + 1613, + 1614, + 1615, + 1616, + 1617, + 1618, + 1620, + 1621, + 1623, + 1624, + 1625, + 1626, + 1627, + 1628, + 1629, + 1630, + 1643, + 1644, + 1646, + 1647, + 1650, + 1651, + 1654, + 1655, + 1657, + 1658, + 1659, + 1660, + 1661, + 1662, + 1663, + 1664, + 1665, + 1666, + 1685, + 1686, + 1687, + 1688, + 1689, + 1690, + 1691, + 1692, + 1693, + 1694, + 1695, + 1699, + 1700, + 1701, + 1702, + 1721, + 1722, + 1723, + 1724, + 1725, + 1726, + 1727, + 1728, + 1729, + 1730, + 1732, + 1736, + 1738, + 1741, + 1742, + 1743, + 1744, + 1750, + 1752, + 1900, + 1909, + 1910, + 1911, + 1912, + 1913, + 1914, + 1915, + 1916, + 1917, + 1918, + 1919, + 1920, + 1921, + 1922, + 1923, + 1924, + 1925, + 1926, + 1927, + 1928, + 1929, + 1930, + 1931, + 1932, + 1933, + 1934, + 1935, + 1936, + 1937, + 1938, + 1939, + 1940, + 1941, + 1942, + 1943, + 1944, + 1945, + 1946, + 1947, + 1948, + 1949, + 1950, + 1951, + 1952, + 1953, + 1954, + 1955, + 1956, + 1957, + 1958, + 1959, + 1960, + 1961, + 1962, + 1963, + 1964, + 1965, + 1966, + 1967, + 1968, + 1969, + 1970, + 1971, + 1972, + 1973, + 1974, + 1975, + 1976, + 1977, + 1978, + 1979, + 1980, + 2019, + 2059, + 2060, + 2073, + 2089, + 2098, + 2099, + 2100, + 2101, + 2102, + 2103, + 2104, + 2105, + 2106, + 2107, + 2108, + 2109, + 2110, + 2111, + 2112, + 2147, + 2148, + 2206, + 2207, + 2208, + 2209, + 2228, + 2230, + 2234, + 2235, + 2241, + 2242, + 2243, + 2244, + 2279, + 2286, + 2873, + 2874 + ], + "rightForeArm": [ + 5015, + 5016, + 5017, + 5018, + 5019, + 5020, + 5021, + 5022, + 5023, + 5024, + 5025, + 5026, + 5027, + 5028, + 5029, + 5030, + 5031, + 5032, + 5033, + 5034, + 5035, + 5036, + 5037, + 5038, + 5039, + 5040, + 5041, + 5042, + 5043, + 5044, + 5045, + 5046, + 5047, + 5048, + 5049, + 5050, + 5051, + 5052, + 5053, + 5054, + 5055, + 5056, + 5057, + 5058, + 5059, + 5060, + 5061, + 5062, + 5063, + 5064, + 5065, + 5066, + 5067, + 5068, + 5069, + 5070, + 5071, + 5072, + 5073, + 5074, + 5075, + 5076, + 5077, + 5078, + 5079, + 5080, + 5081, + 5082, + 5083, + 5084, + 5085, + 5086, + 5087, + 5090, + 5091, + 5092, + 5093, + 5094, + 5095, + 5096, + 5097, + 5098, + 5099, + 5112, + 5113, + 5116, + 5117, + 5120, + 5121, + 5124, + 5125, + 5126, + 5127, + 5128, + 5129, + 5130, + 5131, + 5132, + 5133, + 5134, + 5135, + 5154, + 5155, + 5156, + 5157, + 5158, + 5159, + 5160, + 5161, + 5162, + 5163, + 5164, + 5168, + 5169, + 5170, + 5171, + 5190, + 5191, + 5192, + 5193, + 5194, + 5195, + 5196, + 5197, + 5198, + 5199, + 5202, + 5205, + 5207, + 5210, + 5211, + 5212, + 5213, + 5219, + 5221, + 5361, + 5370, + 5371, + 5372, + 5373, + 5374, + 5375, + 5376, + 5377, + 5378, + 5379, + 5380, + 5381, + 5382, + 5383, + 5384, + 5385, + 5386, + 5387, + 5388, + 5389, + 5390, + 5391, + 5392, + 5393, + 5394, + 5395, + 5396, + 5397, + 5398, + 5399, + 5400, + 5401, + 5402, + 5403, + 5404, + 5405, + 5406, + 5407, + 5408, + 5409, + 5410, + 5411, + 5412, + 5413, + 5414, + 5415, + 5416, + 5417, + 5418, + 5419, + 5420, + 5421, + 5422, + 5423, + 5424, + 5425, + 5426, + 5427, + 5428, + 5429, + 5430, + 5431, + 5432, + 5433, + 5434, + 5435, + 5436, + 5437, + 5438, + 5439, + 5440, + 5441, + 5480, + 5520, + 5521, + 5534, + 5550, + 5559, + 5560, + 5561, + 5562, + 5563, + 5564, + 5565, + 5566, + 5567, + 5568, + 5569, + 5570, + 5571, + 5572, + 5573, + 5608, + 5609, + 5667, + 5668, + 5669, + 5670, + 5689, + 5691, + 5695, + 5696, + 5702, + 5703, + 5704, + 5705, + 5740, + 5747, + 6334, + 6335 + ], + "neck": [ + 148, + 150, + 151, + 152, + 153, + 172, + 174, + 175, + 201, + 202, + 204, + 205, + 206, + 207, + 208, + 209, + 210, + 211, + 212, + 213, + 214, + 215, + 216, + 217, + 218, + 219, + 222, + 223, + 224, + 225, + 256, + 257, + 284, + 285, + 295, + 296, + 297, + 298, + 299, + 300, + 301, + 302, + 303, + 304, + 305, + 306, + 307, + 308, + 309, + 333, + 334, + 423, + 424, + 425, + 426, + 440, + 441, + 451, + 452, + 453, + 460, + 461, + 571, + 572, + 824, + 825, + 826, + 827, + 828, + 829, + 1279, + 1280, + 1312, + 1313, + 1319, + 1320, + 1331, + 3049, + 3050, + 3057, + 3058, + 3059, + 3068, + 3164, + 3661, + 3662, + 3663, + 3664, + 3665, + 3685, + 3686, + 3687, + 3714, + 3715, + 3716, + 3717, + 3718, + 3719, + 3720, + 3721, + 3722, + 3723, + 3724, + 3725, + 3726, + 3727, + 3728, + 3729, + 3730, + 3731, + 3734, + 3735, + 3736, + 3737, + 3768, + 3769, + 3796, + 3797, + 3807, + 3808, + 3809, + 3810, + 3811, + 3812, + 3813, + 3814, + 3815, + 3816, + 3817, + 3818, + 3819, + 3839, + 3840, + 3918, + 3919, + 3920, + 3921, + 3934, + 3935, + 3942, + 3943, + 3944, + 3950, + 4060, + 4061, + 4312, + 4313, + 4314, + 4315, + 4761, + 4762, + 4792, + 4793, + 4799, + 4800, + 4807 + ], + "rightToeBase": [ + 6611, + 6612, + 6613, + 6614, + 6615, + 6616, + 6617, + 6618, + 6619, + 6620, + 6621, + 6622, + 6623, + 6624, + 6625, + 6626, + 6627, + 6628, + 6629, + 6630, + 6631, + 6632, + 6633, + 6634, + 6635, + 6636, + 6637, + 6638, + 6639, + 6640, + 6641, + 6642, + 6643, + 6644, + 6645, + 6646, + 6647, + 6648, + 6649, + 6650, + 6651, + 6652, + 6653, + 6654, + 6655, + 6656, + 6657, + 6658, + 6659, + 6660, + 6661, + 6662, + 6663, + 6664, + 6665, + 6666, + 6667, + 6668, + 6669, + 6670, + 6671, + 6672, + 6673, + 6674, + 6675, + 6676, + 6677, + 6678, + 6679, + 6680, + 6681, + 6682, + 6683, + 6684, + 6685, + 6686, + 6687, + 6688, + 6689, + 6690, + 6691, + 6692, + 6693, + 6694, + 6695, + 6696, + 6697, + 6698, + 6699, + 6700, + 6701, + 6702, + 6703, + 6704, + 6705, + 6706, + 6707, + 6708, + 6709, + 6710, + 6711, + 6712, + 6713, + 6714, + 6715, + 6716, + 6717, + 6718, + 6736, + 6739, + 6741, + 6743, + 6745, + 6747, + 6749, + 6750, + 6752, + 6754, + 6757, + 6758, + 6760, + 6762 + ], + "spine": [ + 616, + 617, + 630, + 631, + 632, + 633, + 654, + 655, + 656, + 657, + 662, + 663, + 664, + 665, + 720, + 721, + 765, + 766, + 767, + 768, + 796, + 797, + 798, + 799, + 889, + 890, + 916, + 917, + 918, + 919, + 921, + 922, + 923, + 924, + 925, + 926, + 1188, + 1189, + 1211, + 1212, + 1248, + 1249, + 1250, + 1251, + 1264, + 1265, + 1266, + 1267, + 1323, + 1324, + 1325, + 1326, + 1327, + 1328, + 1332, + 1333, + 1334, + 1335, + 1336, + 1344, + 1345, + 1481, + 1482, + 1483, + 1484, + 1485, + 1486, + 1487, + 1488, + 1489, + 1490, + 1491, + 1492, + 1493, + 1494, + 1495, + 1496, + 1767, + 2823, + 2824, + 2825, + 2826, + 2827, + 2828, + 2829, + 2830, + 2831, + 2832, + 2833, + 2834, + 2835, + 2836, + 2837, + 2838, + 2839, + 2840, + 2841, + 2842, + 2843, + 2844, + 2845, + 2847, + 2848, + 2851, + 3016, + 3017, + 3018, + 3019, + 3020, + 3023, + 3024, + 3124, + 3173, + 3476, + 3477, + 3478, + 3480, + 3500, + 3501, + 3502, + 3504, + 3509, + 3511, + 4103, + 4104, + 4118, + 4119, + 4120, + 4121, + 4142, + 4143, + 4144, + 4145, + 4150, + 4151, + 4152, + 4153, + 4208, + 4209, + 4253, + 4254, + 4255, + 4256, + 4284, + 4285, + 4286, + 4287, + 4375, + 4376, + 4402, + 4403, + 4405, + 4406, + 4407, + 4408, + 4409, + 4410, + 4411, + 4412, + 4674, + 4675, + 4694, + 4695, + 4731, + 4732, + 4733, + 4734, + 4747, + 4748, + 4749, + 4750, + 4803, + 4804, + 4805, + 4806, + 4808, + 4809, + 4810, + 4811, + 4812, + 4820, + 4821, + 4953, + 4954, + 4955, + 4956, + 4957, + 4958, + 4959, + 4960, + 4961, + 4962, + 4963, + 4964, + 4965, + 4966, + 4967, + 4968, + 5234, + 6284, + 6285, + 6286, + 6287, + 6288, + 6289, + 6290, + 6291, + 6292, + 6293, + 6294, + 6295, + 6296, + 6297, + 6298, + 6299, + 6300, + 6301, + 6302, + 6303, + 6304, + 6305, + 6306, + 6308, + 6309, + 6312, + 6472, + 6473, + 6474, + 6545, + 6874, + 6875, + 6876, + 6878 + ], + "leftUpLeg": [ + 833, + 834, + 838, + 839, + 847, + 848, + 849, + 850, + 851, + 852, + 853, + 854, + 870, + 871, + 872, + 873, + 874, + 875, + 876, + 877, + 878, + 879, + 880, + 881, + 897, + 898, + 899, + 900, + 901, + 902, + 903, + 904, + 905, + 906, + 907, + 908, + 909, + 910, + 911, + 912, + 913, + 914, + 915, + 933, + 934, + 935, + 936, + 944, + 945, + 946, + 947, + 948, + 949, + 950, + 951, + 952, + 953, + 954, + 955, + 956, + 957, + 958, + 959, + 960, + 961, + 962, + 963, + 964, + 965, + 966, + 967, + 968, + 969, + 970, + 971, + 972, + 973, + 974, + 975, + 976, + 977, + 978, + 979, + 980, + 981, + 982, + 983, + 984, + 985, + 986, + 987, + 988, + 989, + 990, + 991, + 992, + 993, + 994, + 995, + 996, + 997, + 998, + 999, + 1000, + 1001, + 1002, + 1003, + 1004, + 1005, + 1006, + 1007, + 1008, + 1009, + 1010, + 1011, + 1012, + 1013, + 1014, + 1015, + 1016, + 1017, + 1018, + 1019, + 1020, + 1021, + 1022, + 1023, + 1024, + 1025, + 1026, + 1027, + 1028, + 1029, + 1030, + 1031, + 1032, + 1033, + 1034, + 1035, + 1036, + 1037, + 1038, + 1039, + 1040, + 1041, + 1042, + 1043, + 1044, + 1045, + 1046, + 1137, + 1138, + 1139, + 1140, + 1141, + 1142, + 1143, + 1144, + 1145, + 1146, + 1147, + 1148, + 1159, + 1160, + 1161, + 1162, + 1163, + 1164, + 1165, + 1166, + 1167, + 1168, + 1169, + 1170, + 1171, + 1172, + 1173, + 1174, + 1184, + 1185, + 1186, + 1187, + 1221, + 1222, + 1223, + 1224, + 1225, + 1226, + 1227, + 1228, + 1229, + 1230, + 1262, + 1263, + 1274, + 1275, + 1276, + 1277, + 1321, + 1322, + 1354, + 1359, + 1360, + 1361, + 1362, + 1365, + 1366, + 1367, + 1368, + 1451, + 1452, + 1453, + 1455, + 1456, + 1457, + 1458, + 1459, + 1460, + 1461, + 1462, + 1463, + 1475, + 1477, + 1478, + 1479, + 1480, + 1498, + 1499, + 1500, + 1501, + 1511, + 1512, + 1513, + 1514, + 1516, + 1517, + 1518, + 1519, + 1520, + 1521, + 1522, + 1533, + 1534, + 3125, + 3126, + 3127, + 3128, + 3131, + 3132, + 3133, + 3134, + 3135, + 3475, + 3479 + ], + "leftHand": [ + 1981, + 1982, + 1983, + 1984, + 1985, + 1986, + 1987, + 1988, + 1989, + 1990, + 1991, + 1992, + 1993, + 1994, + 1995, + 1996, + 1997, + 1998, + 1999, + 2000, + 2001, + 2002, + 2003, + 2004, + 2005, + 2006, + 2007, + 2008, + 2009, + 2010, + 2011, + 2012, + 2013, + 2014, + 2015, + 2016, + 2017, + 2018, + 2019, + 2020, + 2021, + 2022, + 2023, + 2024, + 2025, + 2026, + 2031, + 2032, + 2033, + 2034, + 2035, + 2036, + 2041, + 2042, + 2043, + 2044, + 2045, + 2046, + 2047, + 2048, + 2049, + 2050, + 2051, + 2052, + 2053, + 2054, + 2055, + 2056, + 2057, + 2058, + 2059, + 2060, + 2061, + 2062, + 2063, + 2064, + 2065, + 2066, + 2069, + 2070, + 2071, + 2072, + 2073, + 2074, + 2075, + 2076, + 2077, + 2078, + 2079, + 2080, + 2081, + 2082, + 2083, + 2084, + 2085, + 2086, + 2087, + 2088, + 2089, + 2090, + 2091, + 2092, + 2093, + 2094, + 2095, + 2096, + 2097, + 2098, + 2099, + 2100, + 2101, + 2107, + 2111, + 2113, + 2114, + 2115, + 2116, + 2117, + 2118, + 2119, + 2120, + 2121, + 2122, + 2127, + 2130, + 2131, + 2132, + 2133, + 2134, + 2135, + 2136, + 2137, + 2138, + 2139, + 2140, + 2141, + 2142, + 2143, + 2144, + 2149, + 2150, + 2151, + 2152, + 2155, + 2160, + 2163, + 2164, + 2170, + 2171, + 2172, + 2173, + 2174, + 2175, + 2176, + 2177, + 2178, + 2179, + 2180, + 2182, + 2183, + 2184, + 2185, + 2188, + 2189, + 2191, + 2192, + 2193, + 2194, + 2195, + 2196, + 2197, + 2198, + 2199, + 2200, + 2201, + 2202, + 2203, + 2207, + 2209, + 2210, + 2211, + 2212, + 2213, + 2214, + 2221, + 2222, + 2223, + 2224, + 2225, + 2226, + 2227, + 2228, + 2229, + 2231, + 2234, + 2236, + 2237, + 2238, + 2239, + 2240, + 2246, + 2247, + 2248, + 2249, + 2250, + 2251, + 2252, + 2253, + 2254, + 2255, + 2256, + 2257, + 2258, + 2259, + 2260, + 2262, + 2263, + 2264, + 2265, + 2266, + 2267, + 2268, + 2269, + 2270, + 2271, + 2274, + 2275, + 2276, + 2277, + 2278, + 2279, + 2284, + 2285, + 2287, + 2288, + 2289, + 2290, + 2293, + 2595, + 2598, + 2605, + 2608, + 2697, + 2698, + 2699, + 2700, + 2701, + 2702, + 2703, + 2704, + 2705, + 2706, + 2707, + 2708, + 2709, + 2710, + 2711, + 2712, + 2713, + 2714, + 2715, + 2716, + 2717, + 2718, + 2719, + 2720, + 2721, + 2722, + 2723, + 2724, + 2725, + 2726, + 2727, + 2728, + 2729, + 2730, + 2731, + 2732, + 2733, + 2734, + 2735, + 2736, + 2737, + 2738, + 2739, + 2740, + 2741, + 2742, + 2743, + 2744, + 2745, + 2746, + 2747, + 2748, + 2749, + 2750, + 2751, + 2752, + 2753, + 2754, + 2755, + 2756, + 2757, + 2758, + 2759, + 2760, + 2761, + 2762, + 2763, + 2764, + 2765, + 2766, + 2767, + 2768, + 2769, + 2770, + 2771, + 2772, + 2773, + 2774, + 2775, + 2776, + 2777, + 2778 + ], + "hips": [ + 631, + 632, + 654, + 657, + 662, + 665, + 676, + 677, + 678, + 679, + 705, + 720, + 796, + 799, + 800, + 801, + 802, + 807, + 808, + 809, + 810, + 815, + 816, + 822, + 823, + 830, + 831, + 832, + 833, + 834, + 835, + 836, + 837, + 838, + 839, + 840, + 841, + 842, + 843, + 844, + 845, + 846, + 855, + 856, + 857, + 858, + 859, + 860, + 861, + 862, + 863, + 864, + 865, + 866, + 867, + 868, + 869, + 871, + 878, + 881, + 882, + 883, + 884, + 885, + 886, + 887, + 888, + 889, + 890, + 912, + 915, + 916, + 917, + 918, + 919, + 920, + 932, + 937, + 938, + 939, + 1163, + 1166, + 1203, + 1204, + 1205, + 1206, + 1207, + 1208, + 1209, + 1210, + 1246, + 1247, + 1262, + 1263, + 1276, + 1277, + 1278, + 1321, + 1336, + 1337, + 1338, + 1339, + 1353, + 1354, + 1361, + 1362, + 1363, + 1364, + 1446, + 1447, + 1448, + 1449, + 1450, + 1454, + 1476, + 1497, + 1511, + 1513, + 1514, + 1515, + 1533, + 1534, + 1539, + 1540, + 1768, + 1769, + 1779, + 1780, + 1781, + 1782, + 1783, + 1784, + 1785, + 1786, + 1787, + 1788, + 1789, + 1790, + 1791, + 1792, + 1793, + 1794, + 1795, + 1796, + 1797, + 1798, + 1799, + 1800, + 1801, + 1802, + 1803, + 1804, + 1805, + 1806, + 1807, + 2909, + 2910, + 2911, + 2912, + 2913, + 2914, + 2915, + 2916, + 2917, + 2918, + 2919, + 2920, + 2921, + 2922, + 2923, + 2924, + 2925, + 2926, + 2927, + 2928, + 2929, + 2930, + 3018, + 3019, + 3021, + 3022, + 3080, + 3081, + 3082, + 3083, + 3084, + 3085, + 3086, + 3087, + 3088, + 3089, + 3090, + 3091, + 3092, + 3093, + 3094, + 3095, + 3096, + 3097, + 3098, + 3099, + 3100, + 3101, + 3102, + 3103, + 3104, + 3105, + 3106, + 3107, + 3108, + 3109, + 3110, + 3111, + 3112, + 3113, + 3114, + 3115, + 3116, + 3117, + 3118, + 3119, + 3120, + 3121, + 3122, + 3123, + 3124, + 3128, + 3129, + 3130, + 3136, + 3137, + 3138, + 3139, + 3140, + 3141, + 3142, + 3143, + 3144, + 3145, + 3146, + 3147, + 3148, + 3149, + 3150, + 3151, + 3152, + 3153, + 3154, + 3155, + 3156, + 3157, + 3158, + 3159, + 3160, + 3170, + 3172, + 3481, + 3484, + 3500, + 3502, + 3503, + 3507, + 3510, + 4120, + 4121, + 4142, + 4143, + 4150, + 4151, + 4164, + 4165, + 4166, + 4167, + 4193, + 4208, + 4284, + 4285, + 4288, + 4289, + 4290, + 4295, + 4296, + 4297, + 4298, + 4303, + 4304, + 4310, + 4311, + 4316, + 4317, + 4318, + 4319, + 4320, + 4321, + 4322, + 4323, + 4324, + 4325, + 4326, + 4327, + 4328, + 4329, + 4330, + 4331, + 4332, + 4341, + 4342, + 4343, + 4344, + 4345, + 4346, + 4347, + 4348, + 4349, + 4350, + 4351, + 4352, + 4353, + 4354, + 4355, + 4356, + 4364, + 4365, + 4368, + 4369, + 4370, + 4371, + 4372, + 4373, + 4374, + 4375, + 4376, + 4398, + 4399, + 4402, + 4403, + 4404, + 4405, + 4406, + 4418, + 4423, + 4424, + 4425, + 4649, + 4650, + 4689, + 4690, + 4691, + 4692, + 4693, + 4729, + 4730, + 4745, + 4746, + 4759, + 4760, + 4801, + 4812, + 4813, + 4814, + 4815, + 4829, + 4836, + 4837, + 4919, + 4920, + 4921, + 4922, + 4923, + 4927, + 4969, + 4983, + 4984, + 4986, + 5004, + 5005, + 5244, + 5245, + 5246, + 5247, + 5248, + 5249, + 5250, + 5251, + 5252, + 5253, + 5254, + 5255, + 5256, + 5257, + 5258, + 5259, + 5260, + 5261, + 5262, + 5263, + 5264, + 5265, + 5266, + 5267, + 5268, + 6368, + 6369, + 6370, + 6371, + 6372, + 6373, + 6374, + 6375, + 6376, + 6377, + 6378, + 6379, + 6380, + 6381, + 6382, + 6383, + 6384, + 6385, + 6386, + 6387, + 6388, + 6389, + 6473, + 6474, + 6504, + 6505, + 6506, + 6507, + 6508, + 6509, + 6510, + 6511, + 6512, + 6513, + 6514, + 6515, + 6516, + 6517, + 6518, + 6519, + 6520, + 6521, + 6522, + 6523, + 6524, + 6525, + 6526, + 6527, + 6528, + 6529, + 6530, + 6531, + 6532, + 6533, + 6534, + 6535, + 6536, + 6537, + 6538, + 6539, + 6540, + 6541, + 6542, + 6543, + 6544, + 6545, + 6549, + 6550, + 6551, + 6557, + 6558, + 6559, + 6560, + 6561, + 6562, + 6563, + 6564, + 6565, + 6566, + 6567, + 6568, + 6569, + 6570, + 6571, + 6572, + 6573 + ] +} \ No newline at end of file diff --git a/lib/common/train_util.py b/lib/common/train_util.py new file mode 100644 index 0000000000000000000000000000000000000000..eb1b690c3abdfef9e0190a2d99109ac54ca2bc16 --- /dev/null +++ b/lib/common/train_util.py @@ -0,0 +1,597 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import yaml +import os.path as osp +import torch +import numpy as np +import torch.nn.functional as F +from ..dataset.mesh_util import * +from ..net.geometry import orthogonal +from pytorch3d.renderer.mesh import rasterize_meshes +from .render_utils import Pytorch3dRasterizer +from pytorch3d.structures import Meshes +import cv2 +from PIL import Image +from tqdm import tqdm +import os +from termcolor import colored + + +def reshape_sample_tensor(sample_tensor, num_views): + if num_views == 1: + return sample_tensor + # Need to repeat sample_tensor along the batch dim num_views times + sample_tensor = sample_tensor.unsqueeze(dim=1) + sample_tensor = sample_tensor.repeat(1, num_views, 1, 1) + sample_tensor = sample_tensor.view( + sample_tensor.shape[0] * sample_tensor.shape[1], + sample_tensor.shape[2], sample_tensor.shape[3]) + return sample_tensor + + +def gen_mesh_eval(opt, net, cuda, data, resolution=None): + resolution = opt.resolution if resolution is None else resolution + image_tensor = data['img'].to(device=cuda) + calib_tensor = data['calib'].to(device=cuda) + + net.filter(image_tensor) + + b_min = data['b_min'] + b_max = data['b_max'] + try: + verts, faces, _, _ = reconstruction_faster(net, + cuda, + calib_tensor, + resolution, + b_min, + b_max, + use_octree=False) + + except Exception as e: + print(e) + print('Can not create marching cubes at this time.') + verts, faces = None, None + return verts, faces + + +def gen_mesh(opt, net, cuda, data, save_path, resolution=None): + resolution = opt.resolution if resolution is None else resolution + image_tensor = data['img'].to(device=cuda) + calib_tensor = data['calib'].to(device=cuda) + + net.filter(image_tensor) + + b_min = data['b_min'] + b_max = data['b_max'] + try: + save_img_path = save_path[:-4] + '.png' + save_img_list = [] + for v in range(image_tensor.shape[0]): + save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), + (1, 2, 0)) * 0.5 + + 0.5)[:, :, ::-1] * 255.0 + save_img_list.append(save_img) + save_img = np.concatenate(save_img_list, axis=1) + Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path) + + verts, faces, _, _ = reconstruction_faster(net, cuda, calib_tensor, + resolution, b_min, b_max) + verts_tensor = torch.from_numpy( + verts.T).unsqueeze(0).to(device=cuda).float() + xyz_tensor = net.projection(verts_tensor, calib_tensor[:1]) + uv = xyz_tensor[:, :2, :] + color = netG.index(image_tensor[:1], uv).detach().cpu().numpy()[0].T + color = color * 0.5 + 0.5 + save_obj_mesh_with_color(save_path, verts, faces, color) + except Exception as e: + print(e) + print('Can not create marching cubes at this time.') + verts, faces, color = None, None, None + return verts, faces, color + + +def gen_mesh_color(opt, netG, netC, cuda, data, save_path, use_octree=True): + image_tensor = data['img'].to(device=cuda) + calib_tensor = data['calib'].to(device=cuda) + + netG.filter(image_tensor) + netC.filter(image_tensor) + netC.attach(netG.get_im_feat()) + + b_min = data['b_min'] + b_max = data['b_max'] + try: + save_img_path = save_path[:-4] + '.png' + save_img_list = [] + for v in range(image_tensor.shape[0]): + save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), + (1, 2, 0)) * 0.5 + + 0.5)[:, :, ::-1] * 255.0 + save_img_list.append(save_img) + save_img = np.concatenate(save_img_list, axis=1) + Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path) + + verts, faces, _, _ = reconstruction_faster(netG, + cuda, + calib_tensor, + opt.resolution, + b_min, + b_max, + use_octree=use_octree) + + # Now Getting colors + verts_tensor = torch.from_numpy( + verts.T).unsqueeze(0).to(device=cuda).float() + verts_tensor = reshape_sample_tensor(verts_tensor, opt.num_views) + color = np.zeros(verts.shape) + interval = 10000 + for i in range(len(color) // interval): + left = i * interval + right = i * interval + interval + if i == len(color) // interval - 1: + right = -1 + netC.query(verts_tensor[:, :, left:right], calib_tensor) + rgb = netC.get_preds()[0].detach().cpu().numpy() * 0.5 + 0.5 + color[left:right] = rgb.T + + save_obj_mesh_with_color(save_path, verts, faces, color) + except Exception as e: + print(e) + print('Can not create marching cubes at this time.') + verts, faces, color = None, None, None + return verts, faces, color + + +def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma): + """Sets the learning rate to the initial LR decayed by schedule""" + if epoch in schedule: + lr *= gamma + for param_group in optimizer.param_groups: + param_group['lr'] = lr + return lr + + +def compute_acc(pred, gt, thresh=0.5): + ''' + return: + IOU, precision, and recall + ''' + with torch.no_grad(): + vol_pred = pred > thresh + vol_gt = gt > thresh + + union = vol_pred | vol_gt + inter = vol_pred & vol_gt + + true_pos = inter.sum().float() + + union = union.sum().float() + if union == 0: + union = 1 + vol_pred = vol_pred.sum().float() + if vol_pred == 0: + vol_pred = 1 + vol_gt = vol_gt.sum().float() + if vol_gt == 0: + vol_gt = 1 + return true_pos / union, true_pos / vol_pred, true_pos / vol_gt + + +# def calc_metrics(opt, net, cuda, dataset, num_tests, +# resolution=128, sampled_points=1000, use_kaolin=True): +# if num_tests > len(dataset): +# num_tests = len(dataset) +# with torch.no_grad(): +# chamfer_arr, p2s_arr = [], [] +# for idx in tqdm(range(num_tests)): +# data = dataset[idx * len(dataset) // num_tests] + +# verts, faces = gen_mesh_eval(opt, net, cuda, data, resolution) +# if verts is None: +# continue + +# mesh_gt = trimesh.load(data['mesh_path']) +# mesh_gt = mesh_gt.split(only_watertight=False) +# comp_num = [mesh.vertices.shape[0] for mesh in mesh_gt] +# mesh_gt = mesh_gt[comp_num.index(max(comp_num))] + +# mesh_pred = trimesh.Trimesh(verts, faces) + +# gt_surface_pts, _ = trimesh.sample.sample_surface_even( +# mesh_gt, sampled_points) +# pred_surface_pts, _ = trimesh.sample.sample_surface_even( +# mesh_pred, sampled_points) + +# if use_kaolin and has_kaolin: +# kal_mesh_gt = kal.rep.TriangleMesh.from_tensors( +# torch.tensor(mesh_gt.vertices).float().to(device=cuda), +# torch.tensor(mesh_gt.faces).long().to(device=cuda)) +# kal_mesh_pred = kal.rep.TriangleMesh.from_tensors( +# torch.tensor(mesh_pred.vertices).float().to(device=cuda), +# torch.tensor(mesh_pred.faces).long().to(device=cuda)) + +# kal_distance_0 = kal.metrics.mesh.point_to_surface( +# torch.tensor(pred_surface_pts).float().to(device=cuda), kal_mesh_gt) +# kal_distance_1 = kal.metrics.mesh.point_to_surface( +# torch.tensor(gt_surface_pts).float().to(device=cuda), kal_mesh_pred) + +# dist_gt_pred = torch.sqrt(kal_distance_0).cpu().numpy() +# dist_pred_gt = torch.sqrt(kal_distance_1).cpu().numpy() +# else: +# try: +# _, dist_pred_gt, _ = trimesh.proximity.closest_point(mesh_pred, gt_surface_pts) +# _, dist_gt_pred, _ = trimesh.proximity.closest_point(mesh_gt, pred_surface_pts) +# except Exception as e: +# print (e) +# continue + +# chamfer_dist = 0.5 * (dist_pred_gt.mean() + dist_gt_pred.mean()) +# p2s_dist = dist_pred_gt.mean() + +# chamfer_arr.append(chamfer_dist) +# p2s_arr.append(p2s_dist) + +# return np.average(chamfer_arr), np.average(p2s_arr) + + +def calc_error(opt, net, cuda, dataset, num_tests): + if num_tests > len(dataset): + num_tests = len(dataset) + with torch.no_grad(): + erorr_arr, IOU_arr, prec_arr, recall_arr = [], [], [], [] + for idx in tqdm(range(num_tests)): + data = dataset[idx * len(dataset) // num_tests] + # retrieve the data + image_tensor = data['img'].to(device=cuda) + calib_tensor = data['calib'].to(device=cuda) + sample_tensor = data['samples'].to(device=cuda).unsqueeze(0) + if opt.num_views > 1: + sample_tensor = reshape_sample_tensor(sample_tensor, + opt.num_views) + label_tensor = data['labels'].to(device=cuda).unsqueeze(0) + + res, error = net.forward(image_tensor, + sample_tensor, + calib_tensor, + labels=label_tensor) + + IOU, prec, recall = compute_acc(res, label_tensor) + + # print( + # '{0}/{1} | Error: {2:06f} IOU: {3:06f} prec: {4:06f} recall: {5:06f}' + # .format(idx, num_tests, error.item(), IOU.item(), prec.item(), recall.item())) + erorr_arr.append(error.item()) + IOU_arr.append(IOU.item()) + prec_arr.append(prec.item()) + recall_arr.append(recall.item()) + + return np.average(erorr_arr), np.average(IOU_arr), np.average( + prec_arr), np.average(recall_arr) + + +def calc_error_color(opt, netG, netC, cuda, dataset, num_tests): + if num_tests > len(dataset): + num_tests = len(dataset) + with torch.no_grad(): + error_color_arr = [] + + for idx in tqdm(range(num_tests)): + data = dataset[idx * len(dataset) // num_tests] + # retrieve the data + image_tensor = data['img'].to(device=cuda) + calib_tensor = data['calib'].to(device=cuda) + color_sample_tensor = data['color_samples'].to( + device=cuda).unsqueeze(0) + + if opt.num_views > 1: + color_sample_tensor = reshape_sample_tensor( + color_sample_tensor, opt.num_views) + + rgb_tensor = data['rgbs'].to(device=cuda).unsqueeze(0) + + netG.filter(image_tensor) + _, errorC = netC.forward(image_tensor, + netG.get_im_feat(), + color_sample_tensor, + calib_tensor, + labels=rgb_tensor) + + # print('{0}/{1} | Error inout: {2:06f} | Error color: {3:06f}' + # .format(idx, num_tests, errorG.item(), errorC.item())) + error_color_arr.append(errorC.item()) + + return np.average(error_color_arr) + + +# pytorch lightning training related fucntions + + +def query_func(opt, netG, features, points, proj_matrix=None): + ''' + - points: size of (bz, N, 3) + - proj_matrix: size of (bz, 4, 4) + return: size of (bz, 1, N) + ''' + assert len(points) == 1 + samples = points.repeat(opt.num_views, 1, 1) + samples = samples.permute(0, 2, 1) # [bz, 3, N] + + # view specific query + if proj_matrix is not None: + samples = orthogonal(samples, proj_matrix) + + calib_tensor = torch.stack([torch.eye(4).float()], dim=0).type_as(samples) + + preds = netG.query(features=features, + points=samples, + calibs=calib_tensor, + regressor=netG.if_regressor) + + if type(preds) is list: + preds = preds[0] + + return preds + + +def isin(ar1, ar2): + return (ar1[..., None] == ar2).any(-1) + + +def in1d(ar1, ar2): + mask = ar2.new_zeros((max(ar1.max(), ar2.max()) + 1, ), dtype=torch.bool) + mask[ar2.unique()] = True + return mask[ar1] + + +def get_visibility(xy, z, faces): + """get the visibility of vertices + + Args: + xy (torch.tensor): [N,2] + z (torch.tensor): [N,1] + faces (torch.tensor): [N,3] + size (int): resolution of rendered image + """ + + xyz = torch.cat((xy, -z), dim=1) + xyz = (xyz + 1.0) / 2.0 + faces = faces.long() + + rasterizer = Pytorch3dRasterizer(image_size=2**12) + meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...]) + raster_settings = rasterizer.raster_settings + + pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( + meshes_screen, + image_size=raster_settings.image_size, + blur_radius=raster_settings.blur_radius, + faces_per_pixel=raster_settings.faces_per_pixel, + bin_size=raster_settings.bin_size, + max_faces_per_bin=raster_settings.max_faces_per_bin, + perspective_correct=raster_settings.perspective_correct, + cull_backfaces=raster_settings.cull_backfaces, + ) + + vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :]) + vis_mask = torch.zeros(size=(z.shape[0], 1)) + vis_mask[vis_vertices_id] = 1.0 + + # print("------------------------\n") + # print(f"keep points : {vis_mask.sum()/len(vis_mask)}") + + return vis_mask + + +def batch_mean(res, key): + # recursive mean for multilevel dicts + return torch.stack([ + x[key] if isinstance(x, dict) else batch_mean(x, key) for x in res + ]).mean() + + +def tf_log_convert(log_dict): + new_log_dict = log_dict.copy() + for k, v in log_dict.items(): + new_log_dict[k.replace("_", "/")] = v + del new_log_dict[k] + + return new_log_dict + + +def bar_log_convert(log_dict, name=None, rot=None): + from decimal import Decimal + + new_log_dict = {} + + if name is not None: + new_log_dict['name'] = name[0] + if rot is not None: + new_log_dict['rot'] = rot[0] + + for k, v in log_dict.items(): + color = "yellow" + if 'loss' in k: + color = "red" + k = k.replace("loss", "L") + elif 'acc' in k: + color = "green" + k = k.replace("acc", "A") + elif 'iou' in k: + color = "green" + k = k.replace("iou", "I") + elif 'prec' in k: + color = "green" + k = k.replace("prec", "P") + elif 'recall' in k: + color = "green" + k = k.replace("recall", "R") + + if 'lr' not in k: + new_log_dict[colored(k.split("_")[1], + color)] = colored(f"{v:.3f}", color) + else: + new_log_dict[colored(k.split("_")[1], + color)] = colored(f"{Decimal(str(v)):.1E}", + color) + + if 'loss' in new_log_dict.keys(): + del new_log_dict['loss'] + + return new_log_dict + + +def accumulate(outputs, rot_num, split): + + hparam_log_dict = {} + + metrics = outputs[0].keys() + datasets = split.keys() + + for dataset in datasets: + for metric in metrics: + keyword = f"hparam/{dataset}-{metric}" + if keyword not in hparam_log_dict.keys(): + hparam_log_dict[keyword] = 0 + for idx in range(split[dataset][0] * rot_num, + split[dataset][1] * rot_num): + hparam_log_dict[keyword] += outputs[idx][metric] + hparam_log_dict[keyword] /= (split[dataset][1] - + split[dataset][0]) * rot_num + + print(colored(hparam_log_dict, "green")) + + return hparam_log_dict + + +def calc_error_N(outputs, targets): + """calculate the error of normal (IGR) + + Args: + outputs (torch.tensor): [B, 3, N] + target (torch.tensor): [B, N, 3] + + # manifold loss and grad_loss in IGR paper + grad_loss = ((nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean() + normals_loss = ((mnfld_grad - normals).abs()).norm(2, dim=1).mean() + + Returns: + torch.tensor: error of valid normals on the surface + """ + # outputs = torch.tanh(-outputs.permute(0,2,1).reshape(-1,3)) + outputs = -outputs.permute(0, 2, 1).reshape(-1, 1) + targets = targets.reshape(-1, 3)[:, 2:3] + with_normals = targets.sum(dim=1).abs() > 0.0 + + # eikonal loss + grad_loss = ((outputs[with_normals].norm(2, dim=-1) - 1)**2).mean() + # normals loss + normal_loss = (outputs - targets)[with_normals].abs().norm(2, dim=1).mean() + + return grad_loss * 0.0 + normal_loss + + +def calc_knn_acc(preds, carn_verts, labels, pick_num): + """calculate knn accuracy + + Args: + preds (torch.tensor): [B, 3, N] + carn_verts (torch.tensor): [SMPLX_V_num, 3] + labels (torch.tensor): [B, N_knn, N] + """ + N_knn_full = labels.shape[1] + preds = preds.permute(0, 2, 1).reshape(-1, 3) + labels = labels.permute(0, 2, 1).reshape(-1, N_knn_full) # [BxN, num_knn] + labels = labels[:, :pick_num] + + dist = torch.cdist(preds, carn_verts, p=2) # [BxN, SMPL_V_num] + knn = dist.topk(k=pick_num, dim=1, largest=False)[1] # [BxN, num_knn] + cat_mat = torch.sort(torch.cat((knn, labels), dim=1))[0] + bool_col = torch.zeros_like(cat_mat)[:, 0] + for i in range(pick_num * 2 - 1): + bool_col += cat_mat[:, i] == cat_mat[:, i + 1] + acc = (bool_col > 0).sum() / len(bool_col) + + return acc + + +def calc_acc_seg(output, target, num_multiseg): + from pytorch_lightning.metrics import Accuracy + return Accuracy()(output.reshape(-1, num_multiseg).cpu(), + target.flatten().cpu()) + + +def add_watermark(imgs, titles): + + # Write some Text + + font = cv2.FONT_HERSHEY_SIMPLEX + bottomLeftCornerOfText = (350, 50) + bottomRightCornerOfText = (800, 50) + fontScale = 1 + fontColor = (1.0, 1.0, 1.0) + lineType = 2 + + for i in range(len(imgs)): + + title = titles[i + 1] + cv2.putText(imgs[i], title, bottomLeftCornerOfText, font, fontScale, + fontColor, lineType) + + if i == 0: + cv2.putText(imgs[i], str(titles[i][0]), bottomRightCornerOfText, + font, fontScale, fontColor, lineType) + + result = np.concatenate(imgs, axis=0).transpose(2, 0, 1) + + return result + + +def make_test_gif(img_dir): + + if img_dir is not None and len(os.listdir(img_dir)) > 0: + for dataset in os.listdir(img_dir): + for subject in sorted(os.listdir(osp.join(img_dir, dataset))): + img_lst = [] + im1 = None + for file in sorted( + os.listdir(osp.join(img_dir, dataset, subject))): + if file[-3:] not in ['obj', 'gif']: + img_path = os.path.join(img_dir, dataset, subject, + file) + if im1 == None: + im1 = Image.open(img_path) + else: + img_lst.append(Image.open(img_path)) + + print(os.path.join(img_dir, dataset, subject, "out.gif")) + im1.save(os.path.join(img_dir, dataset, subject, "out.gif"), + save_all=True, + append_images=img_lst, + duration=500, + loop=0) + + +def export_cfg(logger, cfg): + + cfg_export_file = osp.join(logger.save_dir, logger.name, + f"version_{logger.version}", "cfg.yaml") + + if not osp.exists(cfg_export_file): + os.makedirs(osp.dirname(cfg_export_file), exist_ok=True) + with open(cfg_export_file, "w+") as file: + _ = yaml.dump(cfg, file) diff --git a/lib/dataloader_demo.py b/lib/dataloader_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..d9856826485b582e3932f66a9b1afe0c3ddefbbe --- /dev/null +++ b/lib/dataloader_demo.py @@ -0,0 +1,58 @@ +import argparse +from lib.common.config import get_cfg_defaults +from lib.dataset.PIFuDataset import PIFuDataset + +if __name__ == '__main__': + + parser = argparse.ArgumentParser() + parser.add_argument('-v', + '--show', + action='store_true', + help='vis sampler 3D') + parser.add_argument('-s', + '--speed', + action='store_true', + help='vis sampler 3D') + parser.add_argument('-l', + '--list', + action='store_true', + help='vis sampler 3D') + parser.add_argument('-c', + '--config', + default='./configs/train/icon-filter.yaml', + help='vis sampler 3D') + args_c = parser.parse_args() + + args = get_cfg_defaults() + args.merge_from_file(args_c.config) + + dataset = PIFuDataset(args, split='train', vis=args_c.show) + print(f"Number of subjects :{len(dataset.subject_list)}") + data_dict = dataset[0] + + if args_c.list: + for k in data_dict.keys(): + if not hasattr(data_dict[k], "shape"): + print(f"{k}: {data_dict[k]}") + else: + print(f"{k}: {data_dict[k].shape}") + + if args_c.show: + # for item in dataset: + item = dataset[0] + dataset.visualize_sampling3D(item, mode='occ') + + if args_c.speed: + # original: 2 it/s + # smpl online compute: 2 it/s + # normal online compute: 1.5 it/s + from tqdm import tqdm + for item in tqdm(dataset): + # pass + for k in item.keys(): + if 'voxel' in k: + if not hasattr(item[k], "shape"): + print(f"{k}: {item[k]}") + else: + print(f"{k}: {item[k].shape}") + print("--------------------") diff --git a/lib/dataset/Evaluator.py b/lib/dataset/Evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..42aff71b4161c028e64aec2388f515dcc7c916cf --- /dev/null +++ b/lib/dataset/Evaluator.py @@ -0,0 +1,264 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.renderer.gl.normal_render import NormalRender +from lib.dataset.mesh_util import projection +from lib.common.render import Render +from PIL import Image +import os +import numpy as np +import torch +from torch import nn +import trimesh +import os.path as osp +from PIL import Image + + +class Evaluator: + + _normal_render = None + + @staticmethod + def init_gl(): + Evaluator._normal_render = NormalRender(width=512, height=512) + + def __init__(self, device): + self.device = device + self.render = Render(size=512, device=self.device) + self.error_term = nn.MSELoss() + + self.offset = 0.0 + self.scale_factor = None + + def set_mesh(self, result_dict, scale_factor=1.0, offset=0.0): + + for key in result_dict.keys(): + if torch.is_tensor(result_dict[key]): + result_dict[key] = result_dict[key].detach().cpu().numpy() + + for k, v in result_dict.items(): + setattr(self, k, v) + + self.scale_factor = scale_factor + self.offset = offset + + def _render_normal(self, mesh, deg, norms=None): + view_mat = np.identity(4) + rz = deg / 180.0 * np.pi + model_mat = np.identity(4) + model_mat[:3, :3] = self._normal_render.euler_to_rot_mat(0, rz, 0) + model_mat[1, 3] = self.offset + view_mat[2, 2] *= -1 + + self._normal_render.set_matrices(view_mat, model_mat) + if norms is None: + norms = mesh.vertex_normals + self._normal_render.set_normal_mesh(self.scale_factor * mesh.vertices, + mesh.faces, norms, mesh.faces) + self._normal_render.draw() + normal_img = self._normal_render.get_color() + return normal_img + + def render_mesh_list(self, mesh_lst): + + self.offset = 0.0 + self.scale_factor = 1.0 + + full_list = [] + for mesh in mesh_lst: + row_lst = [] + for deg in np.arange(0, 360, 90): + normal = self._render_normal(mesh, deg) + row_lst.append(normal) + full_list.append(np.concatenate(row_lst, axis=1)) + + res_array = np.concatenate(full_list, axis=0) + + return res_array + + def _get_reproj_normal_error(self, deg): + + tgt_normal = self._render_normal(self.tgt_mesh, deg) + src_normal = self._render_normal(self.src_mesh, deg) + error = (((src_normal[:, :, :3] - + tgt_normal[:, :, :3])**2).sum(axis=2).mean(axis=(0, 1))) + + return error, [src_normal, tgt_normal] + + def render_normal(self, verts, faces): + + verts = verts[0].detach().cpu().numpy() + faces = faces[0].detach().cpu().numpy() + + mesh_F = trimesh.Trimesh(verts * np.array([1.0, -1.0, 1.0]), faces) + mesh_B = trimesh.Trimesh(verts * np.array([1.0, -1.0, -1.0]), faces) + + self.scale_factor = 1.0 + + normal_F = self._render_normal(mesh_F, 0) + normal_B = self._render_normal(mesh_B, + 0, + norms=mesh_B.vertex_normals * + np.array([-1.0, -1.0, 1.0])) + + mask = normal_F[:, :, 3:4] + normal_F = (torch.as_tensor(2.0 * (normal_F - 0.5) * mask).permute( + 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) + normal_B = (torch.as_tensor(2.0 * (normal_B - 0.5) * mask).permute( + 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) + + return {"T_normal_F": normal_F, "T_normal_B": normal_B} + + def calculate_normal_consist( + self, + frontal=True, + back=True, + left=True, + right=True, + save_demo_img=None, + return_demo=False, + ): + + # reproj error + # if save_demo_img is not None, save a visualization at the given path (etc, "./test.png") + if self._normal_render is None: + print( + "In order to use normal render, " + "you have to call init_gl() before initialing any evaluator objects." + ) + return -1 + + side_cnt = 0 + total_error = 0 + demo_list = [] + + if frontal: + side_cnt += 1 + error, normal_lst = self._get_reproj_normal_error(0) + total_error += error + demo_list.append(np.concatenate(normal_lst, axis=0)) + if back: + side_cnt += 1 + error, normal_lst = self._get_reproj_normal_error(180) + total_error += error + demo_list.append(np.concatenate(normal_lst, axis=0)) + if left: + side_cnt += 1 + error, normal_lst = self._get_reproj_normal_error(90) + total_error += error + demo_list.append(np.concatenate(normal_lst, axis=0)) + if right: + side_cnt += 1 + error, normal_lst = self._get_reproj_normal_error(270) + total_error += error + demo_list.append(np.concatenate(normal_lst, axis=0)) + if save_demo_img is not None: + res_array = np.concatenate(demo_list, axis=1) + res_img = Image.fromarray((res_array * 255).astype(np.uint8)) + res_img.save(save_demo_img) + + if return_demo: + res_array = np.concatenate(demo_list, axis=1) + return res_array + else: + return total_error + + def space_transfer(self): + + # convert from GT to SDF + self.verts_pr -= self.recon_size / 2.0 + self.verts_pr /= self.recon_size / 2.0 + + self.verts_gt = projection(self.verts_gt, self.calib) + self.verts_gt[:, 1] *= -1 + + self.tgt_mesh = trimesh.Trimesh(self.verts_gt, self.faces_gt) + self.src_mesh = trimesh.Trimesh(self.verts_pr, self.faces_pr) + + # (self.tgt_mesh+self.src_mesh).show() + + def export_mesh(self, dir, name): + self.tgt_mesh.visual.vertex_colors = np.array([255, 0, 0]) + self.src_mesh.visual.vertex_colors = np.array([0, 255, 0]) + + (self.tgt_mesh + self.src_mesh).export( + osp.join(dir, f"{name}_gt_pr.obj")) + + def calculate_chamfer_p2s(self, sampled_points=1000): + """calculate the geometry metrics [chamfer, p2s, chamfer_H, p2s_H] + + Args: + verts_gt (torch.cuda.tensor): [N, 3] + faces_gt (torch.cuda.tensor): [M, 3] + verts_pr (torch.cuda.tensor): [N', 3] + faces_pr (torch.cuda.tensor): [M', 3] + sampled_points (int, optional): use smaller number for faster testing. Defaults to 1000. + + Returns: + tuple: chamfer, p2s, chamfer_H, p2s_H + """ + + gt_surface_pts, _ = trimesh.sample.sample_surface_even( + self.tgt_mesh, sampled_points) + pred_surface_pts, _ = trimesh.sample.sample_surface_even( + self.src_mesh, sampled_points) + + _, dist_pred_gt, _ = trimesh.proximity.closest_point( + self.src_mesh, gt_surface_pts) + _, dist_gt_pred, _ = trimesh.proximity.closest_point( + self.tgt_mesh, pred_surface_pts) + + dist_pred_gt[np.isnan(dist_pred_gt)] = 0 + dist_gt_pred[np.isnan(dist_gt_pred)] = 0 + chamfer_dist = 0.5 * (dist_pred_gt.mean() + + dist_gt_pred.mean()).item() * 100 + p2s_dist = dist_pred_gt.mean().item() * 100 + + return chamfer_dist, p2s_dist + + def calc_acc(self, output, target, thres=0.5, use_sdf=False): + + # # remove the surface points with thres + # non_surf_ids = (target != thres) + # output = output[non_surf_ids] + # target = target[non_surf_ids] + + with torch.no_grad(): + output = output.masked_fill(output < thres, 0.0) + output = output.masked_fill(output > thres, 1.0) + + if use_sdf: + target = target.masked_fill(target < thres, 0.0) + target = target.masked_fill(target > thres, 1.0) + + acc = output.eq(target).float().mean() + + # iou, precison, recall + output = output > thres + target = target > thres + + union = output | target + inter = output & target + + _max = torch.tensor(1.0).to(output.device) + + union = max(union.sum().float(), _max) + true_pos = max(inter.sum().float(), _max) + vol_pred = max(output.sum().float(), _max) + vol_gt = max(target.sum().float(), _max) + + return acc, true_pos / union, true_pos / vol_pred, true_pos / vol_gt diff --git a/lib/dataset/NormalDataset.py b/lib/dataset/NormalDataset.py new file mode 100644 index 0000000000000000000000000000000000000000..636d00c6952c002d3bab2dbcfe52ce80506d42ef --- /dev/null +++ b/lib/dataset/NormalDataset.py @@ -0,0 +1,212 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import os.path as osp +import numpy as np +from PIL import Image +import torchvision.transforms as transforms + + +class NormalDataset(): + def __init__(self, cfg, split='train'): + + self.split = split + self.root = cfg.root + self.overfit = cfg.overfit + + self.opt = cfg.dataset + self.datasets = self.opt.types + self.input_size = self.opt.input_size + self.set_splits = self.opt.set_splits + self.scales = self.opt.scales + self.pifu = self.opt.pifu + + # input data types and dimensions + self.in_nml = [item[0] for item in cfg.net.in_nml] + self.in_nml_dim = [item[1] for item in cfg.net.in_nml] + self.in_total = self.in_nml + ['normal_F', 'normal_B'] + self.in_total_dim = self.in_nml_dim + [3, 3] + + if self.split != 'train': + self.rotations = range(0, 360, 120) + else: + self.rotations = np.arange(0, 360, 360 / + self.opt.rotation_num).astype(np.int) + + self.datasets_dict = {} + for dataset_id, dataset in enumerate(self.datasets): + dataset_dir = osp.join(self.root, dataset, "smplx") + self.datasets_dict[dataset] = { + "subjects": + np.loadtxt(osp.join(self.root, dataset, "all.txt"), dtype=str), + "path": + dataset_dir, + "scale": + self.scales[dataset_id] + } + + self.subject_list = self.get_subject_list(split) + + # PIL to tensor + self.image_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + + # PIL to tensor + self.mask_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.0, ), (1.0, )) + ]) + + def get_subject_list(self, split): + + subject_list = [] + + for dataset in self.datasets: + + if self.pifu: + txt = osp.join(self.root, dataset, f'{split}_pifu.txt') + else: + txt = osp.join(self.root, dataset, f'{split}.txt') + + if osp.exists(txt): + print(f"load from {txt}") + subject_list += sorted(np.loadtxt(txt, dtype=str).tolist()) + + if self.pifu: + miss_pifu = sorted( + np.loadtxt(osp.join(self.root, dataset, + "miss_pifu.txt"), + dtype=str).tolist()) + subject_list = [ + subject for subject in subject_list + if subject not in miss_pifu + ] + subject_list = [ + "renderpeople/" + subject for subject in subject_list + ] + + else: + train_txt = osp.join(self.root, dataset, 'train.txt') + val_txt = osp.join(self.root, dataset, 'val.txt') + test_txt = osp.join(self.root, dataset, 'test.txt') + + print( + f"generate lists of [train, val, test] \n {train_txt} \n {val_txt} \n {test_txt} \n" + ) + + split_txt = osp.join(self.root, dataset, f'{split}.txt') + + subjects = self.datasets_dict[dataset]['subjects'] + train_split = int(len(subjects) * self.set_splits[0]) + val_split = int( + len(subjects) * self.set_splits[1]) + train_split + + with open(train_txt, "w") as f: + f.write("\n".join(dataset + "/" + item + for item in subjects[:train_split])) + with open(val_txt, "w") as f: + f.write("\n".join( + dataset + "/" + item + for item in subjects[train_split:val_split])) + with open(test_txt, "w") as f: + f.write("\n".join(dataset + "/" + item + for item in subjects[val_split:])) + + subject_list += sorted( + np.loadtxt(split_txt, dtype=str).tolist()) + + bug_list = sorted( + np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist()) + + subject_list = [ + subject for subject in subject_list if (subject not in bug_list) + ] + + return subject_list + + def __len__(self): + return len(self.subject_list) * len(self.rotations) + + def __getitem__(self, index): + + # only pick the first data if overfitting + if self.overfit: + index = 0 + + rid = index % len(self.rotations) + mid = index // len(self.rotations) + + rotation = self.rotations[rid] + + # choose specific test sets + subject = self.subject_list[mid] + + subject_render = "/".join( + [subject.split("/")[0] + "_12views", + subject.split("/")[1]]) + + # setup paths + data_dict = { + 'dataset': + subject.split("/")[0], + 'subject': + subject, + 'rotation': + rotation, + 'image_path': + osp.join(self.root, subject_render, 'render', + f'{rotation:03d}.png') + } + + # image/normal/depth loader + for name, channel in zip(self.in_total, self.in_total_dim): + + if name != 'image': + data_dict.update({ + f'{name}_path': + osp.join(self.root, subject_render, name, + f'{rotation:03d}.png') + }) + data_dict.update({ + name: + self.imagepath2tensor(data_dict[f'{name}_path'], + channel, + inv='depth_B' in name) + }) + + path_keys = [ + key for key in data_dict.keys() if '_path' in key or '_dir' in key + ] + for key in path_keys: + del data_dict[key] + + return data_dict + + def imagepath2tensor(self, path, channel=3, inv=False): + + rgba = Image.open(path).convert('RGBA') + mask = rgba.split()[-1] + image = rgba.convert('RGB') + image = self.image_to_tensor(image) + mask = self.mask_to_tensor(mask) + image = (image * mask)[:channel] + + return (image * (0.5 - inv) * 2.0).float() diff --git a/lib/dataset/NormalModule.py b/lib/dataset/NormalModule.py new file mode 100644 index 0000000000000000000000000000000000000000..6b87ec2ff2d1c29f2811bc57fe69b151d59084aa --- /dev/null +++ b/lib/dataset/NormalModule.py @@ -0,0 +1,94 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np +from torch.utils.data import DataLoader +from .NormalDataset import NormalDataset + +# pytorch lightning related libs +import pytorch_lightning as pl + + +class NormalModule(pl.LightningDataModule): + def __init__(self, cfg): + super(NormalModule, self).__init__() + self.cfg = cfg + self.overfit = self.cfg.overfit + + if self.overfit: + self.batch_size = 1 + else: + self.batch_size = self.cfg.batch_size + + self.data_size = {} + + def prepare_data(self): + + pass + + @staticmethod + def worker_init_fn(worker_id): + np.random.seed(np.random.get_state()[1][0] + worker_id) + + def setup(self, stage): + + if stage == 'fit' or stage is None: + self.train_dataset = NormalDataset(cfg=self.cfg, split="train") + self.val_dataset = NormalDataset(cfg=self.cfg, split="val") + self.data_size = { + 'train': len(self.train_dataset), + 'val': len(self.val_dataset) + } + + if stage == 'test' or stage is None: + self.test_dataset = NormalDataset(cfg=self.cfg, split="test") + + def train_dataloader(self): + + train_data_loader = DataLoader(self.train_dataset, + batch_size=self.batch_size, + shuffle=not self.overfit, + num_workers=self.cfg.num_threads, + pin_memory=True, + worker_init_fn=self.worker_init_fn) + + return train_data_loader + + def val_dataloader(self): + + if self.overfit: + current_dataset = self.train_dataset + else: + current_dataset = self.val_dataset + + val_data_loader = DataLoader(current_dataset, + batch_size=self.batch_size, + shuffle=False, + num_workers=self.cfg.num_threads, + pin_memory=True) + + return val_data_loader + + def test_dataloader(self): + + test_data_loader = DataLoader(self.test_dataset, + batch_size=1, + shuffle=False, + num_workers=self.cfg.num_threads, + pin_memory=True) + + return test_data_loader diff --git a/lib/dataset/PIFuDataModule.py b/lib/dataset/PIFuDataModule.py new file mode 100644 index 0000000000000000000000000000000000000000..8793b507c57862415645ead0bfaf4f9b1f6cc7d6 --- /dev/null +++ b/lib/dataset/PIFuDataModule.py @@ -0,0 +1,71 @@ +import numpy as np +from torch.utils.data import DataLoader +from .PIFuDataset import PIFuDataset +import pytorch_lightning as pl + + +class PIFuDataModule(pl.LightningDataModule): + def __init__(self, cfg): + super(PIFuDataModule, self).__init__() + self.cfg = cfg + self.overfit = self.cfg.overfit + + if self.overfit: + self.batch_size = 1 + else: + self.batch_size = self.cfg.batch_size + + self.data_size = {} + + def prepare_data(self): + + pass + + @staticmethod + def worker_init_fn(worker_id): + np.random.seed(np.random.get_state()[1][0] + worker_id) + + def setup(self, stage): + + if stage == 'fit': + self.train_dataset = PIFuDataset(cfg=self.cfg, split="train") + self.val_dataset = PIFuDataset(cfg=self.cfg, split="val") + self.data_size = {'train': len(self.train_dataset), + 'val': len(self.val_dataset)} + + if stage == 'test': + self.test_dataset = PIFuDataset(cfg=self.cfg, split="test") + + def train_dataloader(self): + + train_data_loader = DataLoader( + self.train_dataset, + batch_size=self.batch_size, shuffle=True, + num_workers=self.cfg.num_threads, pin_memory=True, + worker_init_fn=self.worker_init_fn) + + return train_data_loader + + def val_dataloader(self): + + if self.overfit: + current_dataset = self.train_dataset + else: + current_dataset = self.val_dataset + + val_data_loader = DataLoader( + current_dataset, + batch_size=1, shuffle=False, + num_workers=self.cfg.num_threads, pin_memory=True, + worker_init_fn=self.worker_init_fn) + + return val_data_loader + + def test_dataloader(self): + + test_data_loader = DataLoader( + self.test_dataset, + batch_size=1, shuffle=False, + num_workers=self.cfg.num_threads, pin_memory=True) + + return test_data_loader diff --git a/lib/dataset/PIFuDataset.py b/lib/dataset/PIFuDataset.py new file mode 100644 index 0000000000000000000000000000000000000000..f9709d3586f431b9fd33b9d45e8d8f8c32e7f4eb --- /dev/null +++ b/lib/dataset/PIFuDataset.py @@ -0,0 +1,662 @@ +from lib.renderer.mesh import load_fit_body +from lib.dataset.hoppeMesh import HoppeMesh +from lib.dataset.body_model import TetraSMPLModel +from lib.common.render import Render +from lib.dataset.mesh_util import SMPLX, projection, cal_sdf_batch, get_visibility +from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis +from termcolor import colored +import os.path as osp +import numpy as np +from PIL import Image +import random +import trimesh +import torch +import vedo +from kaolin.ops.mesh import check_sign +import torchvision.transforms as transforms +from ipdb import set_trace + + +class PIFuDataset(): + def __init__(self, cfg, split='train', vis=False): + + self.split = split + self.root = cfg.root + self.bsize = cfg.batch_size + self.overfit = cfg.overfit + + # for debug, only used in visualize_sampling3D + self.vis = vis + + self.opt = cfg.dataset + self.datasets = self.opt.types + self.input_size = self.opt.input_size + self.scales = self.opt.scales + self.workers = cfg.num_threads + self.prior_type = cfg.net.prior_type + + self.noise_type = self.opt.noise_type + self.noise_scale = self.opt.noise_scale + + noise_joints = [4, 5, 7, 8, 13, 14, 16, 17, 18, 19, 20, 21] + + self.noise_smpl_idx = [] + self.noise_smplx_idx = [] + + for idx in noise_joints: + self.noise_smpl_idx.append(idx * 3) + self.noise_smpl_idx.append(idx * 3 + 1) + self.noise_smpl_idx.append(idx * 3 + 2) + + self.noise_smplx_idx.append((idx-1) * 3) + self.noise_smplx_idx.append((idx-1) * 3 + 1) + self.noise_smplx_idx.append((idx-1) * 3 + 2) + + self.use_sdf = cfg.sdf + self.sdf_clip = cfg.sdf_clip + + # [(feat_name, channel_num),...] + self.in_geo = [item[0] for item in cfg.net.in_geo] + self.in_nml = [item[0] for item in cfg.net.in_nml] + + self.in_geo_dim = [item[1] for item in cfg.net.in_geo] + self.in_nml_dim = [item[1] for item in cfg.net.in_nml] + + self.in_total = self.in_geo + self.in_nml + self.in_total_dim = self.in_geo_dim + self.in_nml_dim + + if self.split == 'train': + self.rotations = np.arange( + 0, 360, 360 / self.opt.rotation_num).astype(np.int32) + else: + self.rotations = range(0, 360, 120) + + self.datasets_dict = {} + + for dataset_id, dataset in enumerate(self.datasets): + + mesh_dir = None + smplx_dir = None + + dataset_dir = osp.join(self.root, dataset) + + if dataset in ['thuman2']: + mesh_dir = osp.join(dataset_dir, "scans") + smplx_dir = osp.join(dataset_dir, "fits") + smpl_dir = osp.join(dataset_dir, "smpl") + + self.datasets_dict[dataset] = { + "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str), + "smplx_dir": smplx_dir, + "smpl_dir": smpl_dir, + "mesh_dir": mesh_dir, + "scale": self.scales[dataset_id] + } + + self.subject_list = self.get_subject_list(split) + self.smplx = SMPLX() + + # PIL to tensor + self.image_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + + # PIL to tensor + self.mask_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.0, ), (1.0, )) + ]) + + self.device = torch.device(f"cuda:{cfg.gpus[0]}") + self.render = Render(size=512, device=self.device) + + def render_normal(self, verts, faces): + + # render optimized mesh (normal, T_normal, image [-1,1]) + self.render.load_meshes(verts, faces) + return self.render.get_rgb_image() + + def get_subject_list(self, split): + + subject_list = [] + + for dataset in self.datasets: + + split_txt = osp.join(self.root, dataset, f'{split}.txt') + + if osp.exists(split_txt): + print(f"load from {split_txt}") + subject_list += np.loadtxt(split_txt, dtype=str).tolist() + else: + full_txt = osp.join(self.root, dataset, 'all.txt') + print(f"split {full_txt} into train/val/test") + + full_lst = np.loadtxt(full_txt, dtype=str) + full_lst = [dataset+"/"+item for item in full_lst] + [train_lst, test_lst, val_lst] = np.split( + full_lst, [500, 500+5, ]) + + np.savetxt(full_txt.replace( + "all", "train"), train_lst, fmt="%s") + np.savetxt(full_txt.replace("all", "test"), test_lst, fmt="%s") + np.savetxt(full_txt.replace("all", "val"), val_lst, fmt="%s") + + print(f"load from {split_txt}") + subject_list += np.loadtxt(split_txt, dtype=str).tolist() + + if self.split != 'test': + subject_list += subject_list[:self.bsize - + len(subject_list) % self.bsize] + print(colored(f"total: {len(subject_list)}", "yellow")) + random.shuffle(subject_list) + + # subject_list = ["thuman2/0008"] + return subject_list + + def __len__(self): + return len(self.subject_list) * len(self.rotations) + + def __getitem__(self, index): + + # only pick the first data if overfitting + if self.overfit: + index = 0 + + rid = index % len(self.rotations) + mid = index // len(self.rotations) + + rotation = self.rotations[rid] + subject = self.subject_list[mid].split("/")[1] + dataset = self.subject_list[mid].split("/")[0] + render_folder = "/".join([dataset + + f"_{self.opt.rotation_num}views", subject]) + + # setup paths + data_dict = { + 'dataset': dataset, + 'subject': subject, + 'rotation': rotation, + 'scale': self.datasets_dict[dataset]["scale"], + 'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}/{subject}.obj"), + 'smplx_path': osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}/smplx_param.pkl"), + 'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"), + 'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'), + 'vis_path': osp.join(self.root, render_folder, 'vis', f'{rotation:03d}.pt'), + 'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png') + } + + # load training data + data_dict.update(self.load_calib(data_dict)) + + # image/normal/depth loader + for name, channel in zip(self.in_total, self.in_total_dim): + + if f'{name}_path' not in data_dict.keys(): + data_dict.update({ + f'{name}_path': osp.join(self.root, render_folder, name, f'{rotation:03d}.png') + }) + + # tensor update + data_dict.update({ + name: self.imagepath2tensor( + data_dict[f'{name}_path'], channel, inv=False) + }) + + data_dict.update(self.load_mesh(data_dict)) + data_dict.update(self.get_sampling_geo( + data_dict, is_valid=self.split == "val", is_sdf=self.use_sdf)) + data_dict.update(self.load_smpl(data_dict, self.vis)) + + if self.prior_type == 'pamir': + data_dict.update(self.load_smpl_voxel(data_dict)) + + if (self.split != 'test') and (not self.vis): + + del data_dict['verts'] + del data_dict['faces'] + + if not self.vis: + del data_dict['mesh'] + + path_keys = [ + key for key in data_dict.keys() if '_path' in key or '_dir' in key + ] + for key in path_keys: + del data_dict[key] + + return data_dict + + def imagepath2tensor(self, path, channel=3, inv=False): + + rgba = Image.open(path).convert('RGBA') + mask = rgba.split()[-1] + image = rgba.convert('RGB') + image = self.image_to_tensor(image) + mask = self.mask_to_tensor(mask) + image = (image * mask)[:channel] + + return (image * (0.5 - inv) * 2.0).float() + + def load_calib(self, data_dict): + calib_data = np.loadtxt(data_dict['calib_path'], dtype=float) + extrinsic = calib_data[:4, :4] + intrinsic = calib_data[4:8, :4] + calib_mat = np.matmul(intrinsic, extrinsic) + calib_mat = torch.from_numpy(calib_mat).float() + return {'calib': calib_mat} + + def load_mesh(self, data_dict): + mesh_path = data_dict['mesh_path'] + scale = data_dict['scale'] + + mesh_ori = trimesh.load(mesh_path, + skip_materials=True, + process=False, + maintain_order=True) + verts = mesh_ori.vertices * scale + faces = mesh_ori.faces + + vert_normals = np.array(mesh_ori.vertex_normals) + face_normals = np.array(mesh_ori.face_normals) + + mesh = HoppeMesh(verts, faces, vert_normals, face_normals) + + return { + 'mesh': mesh, + 'verts': torch.as_tensor(mesh.verts).float(), + 'faces': torch.as_tensor(mesh.faces).long() + } + + def add_noise(self, + beta_num, + smpl_pose, + smpl_betas, + noise_type, + noise_scale, + type, + hashcode): + + np.random.seed(hashcode) + + if type == 'smplx': + noise_idx = self.noise_smplx_idx + else: + noise_idx = self.noise_smpl_idx + + if 'beta' in noise_type and noise_scale[noise_type.index("beta")] > 0.0: + smpl_betas += (np.random.rand(beta_num) - + 0.5) * 2.0 * noise_scale[noise_type.index("beta")] + smpl_betas = smpl_betas.astype(np.float32) + + if 'pose' in noise_type and noise_scale[noise_type.index("pose")] > 0.0: + smpl_pose[noise_idx] += ( + np.random.rand(len(noise_idx)) - + 0.5) * 2.0 * np.pi * noise_scale[noise_type.index("pose")] + smpl_pose = smpl_pose.astype(np.float32) + if type == 'smplx': + return torch.as_tensor(smpl_pose[None, ...]), torch.as_tensor(smpl_betas[None, ...]) + else: + return smpl_pose, smpl_betas + + def compute_smpl_verts(self, data_dict, noise_type=None, noise_scale=None): + + dataset = data_dict['dataset'] + smplx_dict = {} + + smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True) + smplx_pose = smplx_param["body_pose"] # [1,63] + smplx_betas = smplx_param["betas"] # [1,10] + smplx_pose, smplx_betas = self.add_noise( + smplx_betas.shape[1], + smplx_pose[0], + smplx_betas[0], + noise_type, + noise_scale, + type='smplx', + hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8)) + + smplx_out, _ = load_fit_body(fitted_path=data_dict['smplx_path'], + scale=self.datasets_dict[dataset]['scale'], + smpl_type='smplx', + smpl_gender='male', + noise_dict=dict(betas=smplx_betas, body_pose=smplx_pose)) + + smplx_dict.update({"type": "smplx", + "gender": 'male', + "body_pose": torch.as_tensor(smplx_pose), + "betas": torch.as_tensor(smplx_betas)}) + + return smplx_out.vertices, smplx_dict + + def compute_voxel_verts(self, + data_dict, + noise_type=None, + noise_scale=None): + + smpl_param = np.load(data_dict['smpl_path'], allow_pickle=True) + smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True) + + smpl_pose = rotation_matrix_to_angle_axis( + torch.as_tensor(smpl_param['full_pose'][0])).numpy() + smpl_betas = smpl_param["betas"] + + smpl_path = osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl") + tetra_path = osp.join(self.smplx.tedra_dir, + "tetra_male_adult_smpl.npz") + + smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult') + + smpl_pose, smpl_betas = self.add_noise( + smpl_model.beta_shape[0], + smpl_pose.flatten(), + smpl_betas[0], + noise_type, + noise_scale, + type='smpl', + hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8)) + + smpl_model.set_params(pose=smpl_pose.reshape(-1, 3), + beta=smpl_betas, + trans=smpl_param["transl"]) + + verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added], + axis=0) * smplx_param["scale"] + smplx_param["translation"] + ) * self.datasets_dict[data_dict['dataset']]['scale'] + faces = np.loadtxt(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"), + dtype=np.int32) - 1 + + pad_v_num = int(8000 - verts.shape[0]) + pad_f_num = int(25100 - faces.shape[0]) + + verts = np.pad(verts, ((0, pad_v_num), (0, 0)), + mode='constant', + constant_values=0.0).astype(np.float32) + faces = np.pad(faces, ((0, pad_f_num), (0, 0)), + mode='constant', + constant_values=0.0).astype(np.int32) + + + return verts, faces, pad_v_num, pad_f_num + + def load_smpl(self, data_dict, vis=False): + + smplx_verts, smplx_dict = self.compute_smpl_verts( + data_dict, self.noise_type, + self.noise_scale) # compute using smpl model + + smplx_verts = projection(smplx_verts, data_dict['calib']).float() + smplx_faces = torch.as_tensor(self.smplx.faces).long() + smplx_vis = torch.load(data_dict['vis_path']).float() + smplx_cmap = torch.as_tensor( + np.load(self.smplx.cmap_vert_path)).float() + + # get smpl_signs + query_points = projection(data_dict['samples_geo'], + data_dict['calib']).float() + + pts_signs = 2.0 * (check_sign(smplx_verts.unsqueeze(0), + smplx_faces, + query_points.unsqueeze(0)).float() - 0.5).squeeze(0) + + return_dict = { + 'smpl_verts': smplx_verts, + 'smpl_faces': smplx_faces, + 'smpl_vis': smplx_vis, + 'smpl_cmap': smplx_cmap, + 'pts_signs': pts_signs + } + if smplx_dict is not None: + return_dict.update(smplx_dict) + + if vis: + + (xy, z) = torch.as_tensor(smplx_verts).to( + self.device).split([2, 1], dim=1) + smplx_vis = get_visibility(xy, z, torch.as_tensor( + smplx_faces).to(self.device).long()) + + T_normal_F, T_normal_B = self.render_normal( + (smplx_verts*torch.tensor([1.0, -1.0, 1.0])).to(self.device), + smplx_faces.to(self.device)) + + return_dict.update({"T_normal_F": T_normal_F.squeeze(0), + "T_normal_B": T_normal_B.squeeze(0)}) + query_points = projection(data_dict['samples_geo'], + data_dict['calib']).float() + + smplx_sdf, smplx_norm, smplx_cmap, smplx_vis = cal_sdf_batch( + smplx_verts.unsqueeze(0).to(self.device), + smplx_faces.unsqueeze(0).to(self.device), + smplx_cmap.unsqueeze(0).to(self.device), + smplx_vis.unsqueeze(0).to(self.device), + query_points.unsqueeze(0).contiguous().to(self.device)) + + return_dict.update({ + 'smpl_feat': + torch.cat( + (smplx_sdf[0].detach().cpu(), + smplx_cmap[0].detach().cpu(), + smplx_norm[0].detach().cpu(), + smplx_vis[0].detach().cpu()), + dim=1) + }) + + return return_dict + + def load_smpl_voxel(self, data_dict): + + smpl_verts, smpl_faces, pad_v_num, pad_f_num = self.compute_voxel_verts( + data_dict, self.noise_type, + self.noise_scale) # compute using smpl model + smpl_verts = projection(smpl_verts, data_dict['calib']) + + smpl_verts *= 0.5 + + return { + 'voxel_verts': smpl_verts, + 'voxel_faces': smpl_faces, + 'pad_v_num': pad_v_num, + 'pad_f_num': pad_f_num + } + + def get_sampling_geo(self, data_dict, is_valid=False, is_sdf=False): + + mesh = data_dict['mesh'] + calib = data_dict['calib'] + + # Samples are around the true surface with an offset + n_samples_surface = 4 * self.opt.num_sample_geo + vert_ids = np.arange(mesh.verts.shape[0]) + thickness_sample_ratio = np.ones_like(vert_ids).astype(np.float32) + + thickness_sample_ratio /= thickness_sample_ratio.sum() + + samples_surface_ids = np.random.choice(vert_ids, + n_samples_surface, + replace=True, + p=thickness_sample_ratio) + + samples_normal_ids = np.random.choice(vert_ids, + self.opt.num_sample_geo // 2, + replace=False, + p=thickness_sample_ratio) + + surf_samples = mesh.verts[samples_normal_ids, :] + surf_normals = mesh.vert_normals[samples_normal_ids, :] + + samples_surface = mesh.verts[samples_surface_ids, :] + + # Sampling offsets are random noise with constant scale (15cm - 20cm) + offset = np.random.normal(scale=self.opt.sigma_geo, + size=(n_samples_surface, 1)) + samples_surface += mesh.vert_normals[samples_surface_ids, :] * offset + + # Uniform samples in [-1, 1] + calib_inv = np.linalg.inv(calib) + n_samples_space = self.opt.num_sample_geo // 4 + samples_space_img = 2.0 * np.random.rand(n_samples_space, 3) - 1.0 + samples_space = projection(samples_space_img, calib_inv) + + # z-ray direction samples + if self.opt.zray_type and not is_valid: + n_samples_rayz = self.opt.ray_sample_num + samples_surface_cube = projection(samples_surface, calib) + samples_surface_cube_repeat = np.repeat(samples_surface_cube, + n_samples_rayz, + axis=0) + + thickness_repeat = np.repeat(0.5 * + np.ones_like(samples_surface_ids), + n_samples_rayz, + axis=0) + + noise_repeat = np.random.normal(scale=0.40, + size=(n_samples_surface * + n_samples_rayz, )) + samples_surface_cube_repeat[:, + -1] += thickness_repeat * noise_repeat + samples_surface_rayz = projection(samples_surface_cube_repeat, + calib_inv) + + samples = np.concatenate( + [samples_surface, samples_space, samples_surface_rayz], 0) + else: + samples = np.concatenate([samples_surface, samples_space], 0) + + np.random.shuffle(samples) + + # labels: in->1.0; out->0.0. + if is_sdf: + sdfs = mesh.get_sdf(samples) + inside_samples = samples[sdfs < 0] + outside_samples = samples[sdfs >= 0] + + inside_sdfs = sdfs[sdfs < 0] + outside_sdfs = sdfs[sdfs >= 0] + else: + inside = mesh.contains(samples) + inside_samples = samples[inside >= 0.5] + outside_samples = samples[inside < 0.5] + + nin = inside_samples.shape[0] + + if nin > self.opt.num_sample_geo // 2: + inside_samples = inside_samples[:self.opt.num_sample_geo // 2] + outside_samples = outside_samples[:self.opt.num_sample_geo // 2] + if is_sdf: + inside_sdfs = inside_sdfs[:self.opt.num_sample_geo // 2] + outside_sdfs = outside_sdfs[:self.opt.num_sample_geo // 2] + else: + outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)] + if is_sdf: + outside_sdfs = outside_sdfs[:(self.opt.num_sample_geo - nin)] + + if is_sdf: + samples = np.concatenate( + [inside_samples, outside_samples, surf_samples], 0) + + labels = np.concatenate([ + inside_sdfs, outside_sdfs, 0.0 * np.ones(surf_samples.shape[0]) + ]) + + normals = np.zeros_like(samples) + normals[-self.opt.num_sample_geo // 2:, :] = surf_normals + + # convert sdf from [-14, 130] to [0, 1] + # outside: 0, inside: 1 + # Note: Marching cubes is defined on occupancy space (inside=1.0, outside=0.0) + + labels = -labels.clip(min=-self.sdf_clip, max=self.sdf_clip) + labels += self.sdf_clip + labels /= (self.sdf_clip * 2) + + else: + samples = np.concatenate([inside_samples, outside_samples]) + labels = np.concatenate([ + np.ones(inside_samples.shape[0]), + np.zeros(outside_samples.shape[0]) + ]) + + normals = np.zeros_like(samples) + + samples = torch.from_numpy(samples).float() + labels = torch.from_numpy(labels).float() + normals = torch.from_numpy(normals).float() + + return {'samples_geo': samples, 'labels_geo': labels} + + def visualize_sampling3D(self, data_dict, mode='vis'): + + # create plot + vp = vedo.Plotter(title="", size=(1500, 1500), axes=0, bg='white') + vis_list = [] + + assert mode in ['vis', 'sdf', 'normal', 'cmap', 'occ'] + + # sdf-1 cmap-3 norm-3 vis-1 + if mode == 'vis': + labels = data_dict[f'smpl_feat'][:, [-1]] # visibility + colors = np.concatenate([labels, labels, labels], axis=1) + elif mode == 'occ': + labels = data_dict[f'labels_geo'][..., None] # occupancy + colors = np.concatenate([labels, labels, labels], axis=1) + elif mode == 'sdf': + labels = data_dict[f'smpl_feat'][:, [0]] # sdf + labels -= labels.min() + labels /= labels.max() + colors = np.concatenate([labels, labels, labels], axis=1) + elif mode == 'normal': + labels = data_dict[f'smpl_feat'][:, -4:-1] # normal + colors = (labels + 1.0) * 0.5 + elif mode == 'cmap': + labels = data_dict[f'smpl_feat'][:, -7:-4] # colormap + colors = np.array(labels) + + points = projection(data_dict['samples_geo'], data_dict['calib']) + verts = projection(data_dict['verts'], data_dict['calib']) + points[:, 1] *= -1 + verts[:, 1] *= -1 + + # create a mesh + mesh = trimesh.Trimesh(verts, data_dict['faces'], process=True) + mesh.visual.vertex_colors = [128.0, 128.0, 128.0, 255.0] + vis_list.append(mesh) + + if 'voxel_verts' in data_dict.keys(): + print(colored("voxel verts", "green")) + voxel_verts = data_dict['voxel_verts'] * 2.0 + voxel_faces = data_dict['voxel_faces'] + voxel_verts[:, 1] *= -1 + voxel = trimesh.Trimesh( + voxel_verts, voxel_faces[:, [0, 2, 1]], process=False, maintain_order=True) + voxel.visual.vertex_colors = [0.0, 128.0, 0.0, 255.0] + vis_list.append(voxel) + + if 'smpl_verts' in data_dict.keys(): + print(colored("smpl verts", "green")) + smplx_verts = data_dict['smpl_verts'] + smplx_faces = data_dict['smpl_faces'] + smplx_verts[:, 1] *= -1 + smplx = trimesh.Trimesh( + smplx_verts, smplx_faces[:, [0, 2, 1]], process=False, maintain_order=True) + smplx.visual.vertex_colors = [128.0, 128.0, 0.0, 255.0] + vis_list.append(smplx) + + # create a picure + img_pos = [1.0, 0.0, -1.0] + for img_id, img_key in enumerate(['normal_F', 'image', 'T_normal_B']): + image_arr = (data_dict[img_key].detach().cpu().permute( + 1, 2, 0).numpy() + 1.0) * 0.5 * 255.0 + image_dim = image_arr.shape[0] + image = vedo.Picture(image_arr).scale( + 2.0 / image_dim).pos(-1.0, -1.0, img_pos[img_id]) + vis_list.append(image) + + # create a pointcloud + pc = vedo.Points(points, r=15, c=np.float32(colors)) + vis_list.append(pc) + + vp.show(*vis_list, bg="white", axes=1.0, interactive=True) diff --git a/lib/dataset/TestDataset.py b/lib/dataset/TestDataset.py new file mode 100644 index 0000000000000000000000000000000000000000..716eb455a9dabea227dfc707f453dbcfbef56b8e --- /dev/null +++ b/lib/dataset/TestDataset.py @@ -0,0 +1,342 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import smplx +from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues +from lib.pymaf.utils.imutils import process_image +from lib.pymaf.core import path_config +from lib.pymaf.models import pymaf_net +from lib.common.config import cfg +from lib.common.render import Render +from lib.dataset.body_model import TetraSMPLModel +from lib.dataset.mesh_util import get_visibility, SMPLX +import os.path as osp +import os +import torch +import glob +import numpy as np +import random +import human_det +from termcolor import colored +from PIL import ImageFile + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +class TestDataset(): + def __init__(self, cfg, device): + + random.seed(1993) + + self.image_dir = cfg['image_dir'] + self.seg_dir = cfg['seg_dir'] + self.has_det = cfg['has_det'] + self.hps_type = cfg['hps_type'] + self.smpl_type = 'smpl' if cfg['hps_type'] != 'pixie' else 'smplx' + self.smpl_gender = 'neutral' + + self.device = device + + if self.has_det: + self.det = human_det.Detection() + else: + self.det = None + + keep_lst = sorted(glob.glob(f"{self.image_dir}/*")) + img_fmts = ['jpg', 'png', 'jpeg', "JPG", 'bmp'] + keep_lst = [ + item for item in keep_lst if item.split(".")[-1] in img_fmts + ] + + self.subject_list = sorted( + [item for item in keep_lst if item.split(".")[-1] in img_fmts]) + + # smpl related + self.smpl_data = SMPLX() + + self.get_smpl_model = lambda smpl_type, smpl_gender: smplx.create( + model_path=self.smpl_data.model_dir, + gender=smpl_gender, + model_type=smpl_type, + ext='npz') + + # Load SMPL model + self.smpl_model = self.get_smpl_model( + self.smpl_type, self.smpl_gender).to(self.device) + self.faces = self.smpl_model.faces + + self.hps = pymaf_net(path_config.SMPL_MEAN_PARAMS, + pretrained=True).to(self.device) + self.hps.load_state_dict(torch.load( + path_config.CHECKPOINT_FILE)['model'], + strict=True) + self.hps.eval() + + print(colored(f"Using {self.hps_type} as HPS Estimator\n", "green")) + + self.render = Render(size=512, device=device) + + def __len__(self): + return len(self.subject_list) + + def compute_vis_cmap(self, smpl_verts, smpl_faces): + + (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1) + smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long()) + if self.smpl_type == 'smpl': + smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0])) + else: + smplx_ind = np.arange(smpl_vis.shape[0]) + smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind) + + return { + 'smpl_vis': smpl_vis.unsqueeze(0).to(self.device), + 'smpl_cmap': smpl_cmap.unsqueeze(0).to(self.device), + 'smpl_verts': smpl_verts.unsqueeze(0) + } + + def compute_voxel_verts(self, body_pose, global_orient, betas, trans, + scale): + + smpl_path = osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl") + tetra_path = osp.join(self.smpl_data.tedra_dir, + 'tetra_neutral_adult_smpl.npz') + smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult') + + pose = torch.cat([global_orient[0], body_pose[0]], dim=0) + smpl_model.set_params(rotation_matrix_to_angle_axis(pose), + beta=betas[0]) + + verts = np.concatenate( + [smpl_model.verts, smpl_model.verts_added], + axis=0) * scale.item() + trans.detach().cpu().numpy() + faces = np.loadtxt(osp.join(self.smpl_data.tedra_dir, + 'tetrahedrons_neutral_adult.txt'), + dtype=np.int32) - 1 + + pad_v_num = int(8000 - verts.shape[0]) + pad_f_num = int(25100 - faces.shape[0]) + + verts = np.pad(verts, ((0, pad_v_num), (0, 0)), + mode='constant', + constant_values=0.0).astype(np.float32) * 0.5 + faces = np.pad(faces, ((0, pad_f_num), (0, 0)), + mode='constant', + constant_values=0.0).astype(np.int32) + + verts[:, 2] *= -1.0 + + voxel_dict = { + 'voxel_verts': + torch.from_numpy(verts).to(self.device).unsqueeze(0).float(), + 'voxel_faces': + torch.from_numpy(faces).to(self.device).unsqueeze(0).long(), + 'pad_v_num': + torch.tensor(pad_v_num).to(self.device).unsqueeze(0).long(), + 'pad_f_num': + torch.tensor(pad_f_num).to(self.device).unsqueeze(0).long() + } + + return voxel_dict + + def __getitem__(self, index): + + img_path = self.subject_list[index] + img_name = img_path.split("/")[-1].rsplit(".", 1)[0] + + if self.seg_dir is None: + img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image( + img_path, self.det, self.hps_type, 512, self.device) + + data_dict = { + 'name': img_name, + 'image': img_icon.to(self.device).unsqueeze(0), + 'ori_image': img_ori, + 'mask': img_mask, + 'uncrop_param': uncrop_param + } + + else: + img_icon, img_hps, img_ori, img_mask, uncrop_param, segmentations = process_image( + img_path, self.det, self.hps_type, 512, self.device, + seg_path=os.path.join(self.seg_dir, f'{img_name}.json')) + data_dict = { + 'name': img_name, + 'image': img_icon.to(self.device).unsqueeze(0), + 'ori_image': img_ori, + 'mask': img_mask, + 'uncrop_param': uncrop_param, + 'segmentations': segmentations + } + + with torch.no_grad(): + # import ipdb; ipdb.set_trace() + preds_dict = self.hps.forward(img_hps) + + data_dict['smpl_faces'] = torch.Tensor( + self.faces.astype(np.int16)).long().unsqueeze(0).to( + self.device) + + if self.hps_type == 'pymaf': + output = preds_dict['smpl_out'][-1] + scale, tranX, tranY = output['theta'][0, :3] + data_dict['betas'] = output['pred_shape'] + data_dict['body_pose'] = output['rotmat'][:, 1:] + data_dict['global_orient'] = output['rotmat'][:, 0:1] + data_dict['smpl_verts'] = output['verts'] + + elif self.hps_type == 'pare': + data_dict['body_pose'] = preds_dict['pred_pose'][:, 1:] + data_dict['global_orient'] = preds_dict['pred_pose'][:, 0:1] + data_dict['betas'] = preds_dict['pred_shape'] + data_dict['smpl_verts'] = preds_dict['smpl_vertices'] + scale, tranX, tranY = preds_dict['pred_cam'][0, :3] + + elif self.hps_type == 'pixie': + data_dict.update(preds_dict) + data_dict['body_pose'] = preds_dict['body_pose'] + data_dict['global_orient'] = preds_dict['global_pose'] + data_dict['betas'] = preds_dict['shape'] + data_dict['smpl_verts'] = preds_dict['vertices'] + scale, tranX, tranY = preds_dict['cam'][0, :3] + + elif self.hps_type == 'hybrik': + data_dict['body_pose'] = preds_dict['pred_theta_mats'][:, 1:] + data_dict['global_orient'] = preds_dict['pred_theta_mats'][:, [0]] + data_dict['betas'] = preds_dict['pred_shape'] + data_dict['smpl_verts'] = preds_dict['pred_vertices'] + scale, tranX, tranY = preds_dict['pred_camera'][0, :3] + scale = scale * 2 + + elif self.hps_type == 'bev': + data_dict['betas'] = torch.from_numpy(preds_dict['smpl_betas'])[ + [0], :10].to(self.device).float() + pred_thetas = batch_rodrigues(torch.from_numpy( + preds_dict['smpl_thetas'][0]).reshape(-1, 3)).float() + data_dict['body_pose'] = pred_thetas[1:][None].to(self.device) + data_dict['global_orient'] = pred_thetas[[0]][None].to(self.device) + data_dict['smpl_verts'] = torch.from_numpy( + preds_dict['verts'][[0]]).to(self.device).float() + tranX = preds_dict['cam_trans'][0, 0] + tranY = preds_dict['cam'][0, 1] + 0.28 + scale = preds_dict['cam'][0, 0] * 1.1 + + data_dict['scale'] = scale + data_dict['trans'] = torch.tensor( + [tranX, tranY, 0.0]).to(self.device).float() + + # data_dict info (key-shape): + # scale, tranX, tranY - tensor.float + # betas - [1,10] / [1, 200] + # body_pose - [1, 23, 3, 3] / [1, 21, 3, 3] + # global_orient - [1, 1, 3, 3] + # smpl_verts - [1, 6890, 3] / [1, 10475, 3] + + return data_dict + + def render_normal(self, verts, faces): + + # render optimized mesh (normal, T_normal, image [-1,1]) + self.render.load_meshes(verts, faces) + return self.render.get_rgb_image() + + def render_depth(self, verts, faces): + + # render optimized mesh (normal, T_normal, image [-1,1]) + self.render.load_meshes(verts, faces) + return self.render.get_depth_map(cam_ids=[0, 2]) + + def visualize_alignment(self, data): + + import vedo + import trimesh + + if self.hps_type != 'pixie': + smpl_out = self.smpl_model(betas=data['betas'], + body_pose=data['body_pose'], + global_orient=data['global_orient'], + pose2rot=False) + smpl_verts = ( + (smpl_out.vertices + data['trans']) * data['scale']).detach().cpu().numpy()[0] + else: + smpl_verts, _, _ = self.smpl_model(shape_params=data['betas'], + expression_params=data['exp'], + body_pose=data['body_pose'], + global_pose=data['global_orient'], + jaw_pose=data['jaw_pose'], + left_hand_pose=data['left_hand_pose'], + right_hand_pose=data['right_hand_pose']) + + smpl_verts = ( + (smpl_verts + data['trans']) * data['scale']).detach().cpu().numpy()[0] + + smpl_verts *= np.array([1.0, -1.0, -1.0]) + faces = data['smpl_faces'][0].detach().cpu().numpy() + + image_P = data['image'] + image_F, image_B = self.render_normal(smpl_verts, faces) + + # create plot + vp = vedo.Plotter(title="", size=(1500, 1500)) + vis_list = [] + + image_F = ( + 0.5 * (1.0 + image_F[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0) + image_B = ( + 0.5 * (1.0 + image_B[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0) + image_P = ( + 0.5 * (1.0 + image_P[0].permute(1, 2, 0).detach().cpu().numpy()) * 255.0) + + vis_list.append(vedo.Picture(image_P*0.5+image_F * + 0.5).scale(2.0/image_P.shape[0]).pos(-1.0, -1.0, 1.0)) + vis_list.append(vedo.Picture(image_F).scale( + 2.0/image_F.shape[0]).pos(-1.0, -1.0, -0.5)) + vis_list.append(vedo.Picture(image_B).scale( + 2.0/image_B.shape[0]).pos(-1.0, -1.0, -1.0)) + + # create a mesh + mesh = trimesh.Trimesh(smpl_verts, faces, process=False) + mesh.visual.vertex_colors = [200, 200, 0] + vis_list.append(mesh) + + vp.show(*vis_list, bg="white", axes=1, interactive=True) + + +if __name__ == '__main__': + + cfg.merge_from_file("./configs/icon-filter.yaml") + cfg.merge_from_file('./lib/pymaf/configs/pymaf_config.yaml') + + cfg_show_list = [ + 'test_gpus', ['0'], 'mcube_res', 512, 'clean_mesh', False + ] + + cfg.merge_from_list(cfg_show_list) + cfg.freeze() + + os.environ['CUDA_VISIBLE_DEVICES'] = "0" + device = torch.device('cuda:0') + + dataset = TestDataset( + { + 'image_dir': "./examples", + 'has_det': True, # w/ or w/o detection + 'hps_type': 'bev' # pymaf/pare/pixie/hybrik/bev + }, device) + + for i in range(len(dataset)): + dataset.visualize_alignment(dataset[i]) diff --git a/lib/dataset/__init__.py b/lib/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/dataset/body_model.py b/lib/dataset/body_model.py new file mode 100644 index 0000000000000000000000000000000000000000..5d25d4aeb1e2a877b1652e8022a0f089802609f8 --- /dev/null +++ b/lib/dataset/body_model.py @@ -0,0 +1,494 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np +import pickle +import torch +import os + + +class SMPLModel(): + def __init__(self, model_path, age): + """ + SMPL model. + + Parameter: + --------- + model_path: Path to the SMPL model parameters, pre-processed by + `preprocess.py`. + + """ + with open(model_path, 'rb') as f: + params = pickle.load(f, encoding='latin1') + + self.J_regressor = params['J_regressor'] + self.weights = np.asarray(params['weights']) + self.posedirs = np.asarray(params['posedirs']) + self.v_template = np.asarray(params['v_template']) + self.shapedirs = np.asarray(params['shapedirs']) + self.faces = np.asarray(params['f']) + self.kintree_table = np.asarray(params['kintree_table']) + + self.pose_shape = [24, 3] + self.beta_shape = [10] + self.trans_shape = [3] + + if age == 'kid': + v_template_smil = np.load( + os.path.join(os.path.dirname(model_path), + "smpl/smpl_kid_template.npy")) + v_template_smil -= np.mean(v_template_smil, axis=0) + v_template_diff = np.expand_dims(v_template_smil - self.v_template, + axis=2) + self.shapedirs = np.concatenate( + (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), + axis=2) + self.beta_shape[0] += 1 + + id_to_col = { + self.kintree_table[1, i]: i + for i in range(self.kintree_table.shape[1]) + } + self.parent = { + i: id_to_col[self.kintree_table[0, i]] + for i in range(1, self.kintree_table.shape[1]) + } + + self.pose = np.zeros(self.pose_shape) + self.beta = np.zeros(self.beta_shape) + self.trans = np.zeros(self.trans_shape) + + self.verts = None + self.J = None + self.R = None + self.G = None + + self.update() + + def set_params(self, pose=None, beta=None, trans=None): + """ + Set pose, shape, and/or translation parameters of SMPL model. Verices of the + model will be updated and returned. + + Prameters: + --------- + pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation + relative to parent joint. For root joint it's global orientation. + Represented in a axis-angle format. + + beta: Parameter for model shape. A vector of shape [10]. Coefficients for + PCA component. Only 10 components were released by MPI. + + trans: Global translation of shape [3]. + + Return: + ------ + Updated vertices. + + """ + if pose is not None: + self.pose = pose + if beta is not None: + self.beta = beta + if trans is not None: + self.trans = trans + self.update() + return self.verts + + def update(self): + """ + Called automatically when parameters are updated. + + """ + # how beta affect body shape + v_shaped = self.shapedirs.dot(self.beta) + self.v_template + # joints location + self.J = self.J_regressor.dot(v_shaped) + pose_cube = self.pose.reshape((-1, 1, 3)) + # rotation matrix for each joint + self.R = self.rodrigues(pose_cube) + I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), + (self.R.shape[0] - 1, 3, 3)) + lrotmin = (self.R[1:] - I_cube).ravel() + # how pose affect body shape in zero pose + v_posed = v_shaped + self.posedirs.dot(lrotmin) + # world transformation of each joint + G = np.empty((self.kintree_table.shape[1], 4, 4)) + G[0] = self.with_zeros( + np.hstack((self.R[0], self.J[0, :].reshape([3, 1])))) + for i in range(1, self.kintree_table.shape[1]): + G[i] = G[self.parent[i]].dot( + self.with_zeros( + np.hstack([ + self.R[i], + ((self.J[i, :] - self.J[self.parent[i], :]).reshape( + [3, 1])) + ]))) + # remove the transformation due to the rest pose + G = G - self.pack( + np.matmul( + G, + np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1]))) + # transformation of each vertex + T = np.tensordot(self.weights, G, axes=[[1], [0]]) + rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1]))) + v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, + 4])[:, :3] + self.verts = v + self.trans.reshape([1, 3]) + self.G = G + + def rodrigues(self, r): + """ + Rodrigues' rotation formula that turns axis-angle vector into rotation + matrix in a batch-ed manner. + + Parameter: + ---------- + r: Axis-angle rotation vector of shape [batch_size, 1, 3]. + + Return: + ------- + Rotation matrix of shape [batch_size, 3, 3]. + + """ + theta = np.linalg.norm(r, axis=(1, 2), keepdims=True) + # avoid zero divide + theta = np.maximum(theta, np.finfo(np.float64).tiny) + r_hat = r / theta + cos = np.cos(theta) + z_stick = np.zeros(theta.shape[0]) + m = np.dstack([ + z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, + -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick + ]).reshape([-1, 3, 3]) + i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), + [theta.shape[0], 3, 3]) + A = np.transpose(r_hat, axes=[0, 2, 1]) + B = r_hat + dot = np.matmul(A, B) + R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m + return R + + def with_zeros(self, x): + """ + Append a [0, 0, 0, 1] vector to a [3, 4] matrix. + + Parameter: + --------- + x: Matrix to be appended. + + Return: + ------ + Matrix after appending of shape [4,4] + + """ + return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]]))) + + def pack(self, x): + """ + Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched + manner. + + Parameter: + ---------- + x: Matrices to be appended of shape [batch_size, 4, 1] + + Return: + ------ + Matrix of shape [batch_size, 4, 4] after appending. + + """ + return np.dstack((np.zeros((x.shape[0], 4, 3)), x)) + + def save_to_obj(self, path): + """ + Save the SMPL model into .obj file. + + Parameter: + --------- + path: Path to save. + + """ + with open(path, 'w') as fp: + for v in self.verts: + fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) + for f in self.faces + 1: + fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) + + +class TetraSMPLModel(): + def __init__(self, + model_path, + model_addition_path, + age='adult', + v_template=None): + """ + SMPL model. + + Parameter: + --------- + model_path: Path to the SMPL model parameters, pre-processed by + `preprocess.py`. + + """ + with open(model_path, 'rb') as f: + params = pickle.load(f, encoding='latin1') + + self.J_regressor = params['J_regressor'] + self.weights = np.asarray(params['weights']) + self.posedirs = np.asarray(params['posedirs']) + + if v_template is not None: + self.v_template = v_template + else: + self.v_template = np.asarray(params['v_template']) + + self.shapedirs = np.asarray(params['shapedirs']) + self.faces = np.asarray(params['f']) + self.kintree_table = np.asarray(params['kintree_table']) + + params_added = np.load(model_addition_path) + self.v_template_added = params_added['v_template_added'] + self.weights_added = params_added['weights_added'] + self.shapedirs_added = params_added['shapedirs_added'] + self.posedirs_added = params_added['posedirs_added'] + self.tetrahedrons = params_added['tetrahedrons'] + + id_to_col = { + self.kintree_table[1, i]: i + for i in range(self.kintree_table.shape[1]) + } + self.parent = { + i: id_to_col[self.kintree_table[0, i]] + for i in range(1, self.kintree_table.shape[1]) + } + + self.pose_shape = [24, 3] + self.beta_shape = [10] + self.trans_shape = [3] + + if age == 'kid': + v_template_smil = np.load( + os.path.join(os.path.dirname(model_path), + "smpl/smpl_kid_template.npy")) + v_template_smil -= np.mean(v_template_smil, axis=0) + v_template_diff = np.expand_dims(v_template_smil - self.v_template, + axis=2) + self.shapedirs = np.concatenate( + (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), + axis=2) + self.beta_shape[0] += 1 + + self.pose = np.zeros(self.pose_shape) + self.beta = np.zeros(self.beta_shape) + self.trans = np.zeros(self.trans_shape) + + self.verts = None + self.verts_added = None + self.J = None + self.R = None + self.G = None + + self.update() + + def set_params(self, pose=None, beta=None, trans=None): + """ + Set pose, shape, and/or translation parameters of SMPL model. Verices of the + model will be updated and returned. + + Prameters: + --------- + pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation + relative to parent joint. For root joint it's global orientation. + Represented in a axis-angle format. + + beta: Parameter for model shape. A vector of shape [10]. Coefficients for + PCA component. Only 10 components were released by MPI. + + trans: Global translation of shape [3]. + + Return: + ------ + Updated vertices. + + """ + + if torch.is_tensor(pose): + pose = pose.detach().cpu().numpy() + if torch.is_tensor(beta): + beta = beta.detach().cpu().numpy() + + if pose is not None: + self.pose = pose + if beta is not None: + self.beta = beta + if trans is not None: + self.trans = trans + self.update() + return self.verts + + def update(self): + """ + Called automatically when parameters are updated. + + """ + # how beta affect body shape + v_shaped = self.shapedirs.dot(self.beta) + self.v_template + v_shaped_added = self.shapedirs_added.dot( + self.beta) + self.v_template_added + # joints location + self.J = self.J_regressor.dot(v_shaped) + pose_cube = self.pose.reshape((-1, 1, 3)) + # rotation matrix for each joint + self.R = self.rodrigues(pose_cube) + I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), + (self.R.shape[0] - 1, 3, 3)) + lrotmin = (self.R[1:] - I_cube).ravel() + # how pose affect body shape in zero pose + v_posed = v_shaped + self.posedirs.dot(lrotmin) + v_posed_added = v_shaped_added + self.posedirs_added.dot(lrotmin) + # world transformation of each joint + G = np.empty((self.kintree_table.shape[1], 4, 4)) + G[0] = self.with_zeros( + np.hstack((self.R[0], self.J[0, :].reshape([3, 1])))) + for i in range(1, self.kintree_table.shape[1]): + G[i] = G[self.parent[i]].dot( + self.with_zeros( + np.hstack([ + self.R[i], + ((self.J[i, :] - self.J[self.parent[i], :]).reshape( + [3, 1])) + ]))) + # remove the transformation due to the rest pose + G = G - self.pack( + np.matmul( + G, + np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1]))) + self.G = G + # transformation of each vertex + T = np.tensordot(self.weights, G, axes=[[1], [0]]) + rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1]))) + v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, + 4])[:, :3] + self.verts = v + self.trans.reshape([1, 3]) + T_added = np.tensordot(self.weights_added, G, axes=[[1], [0]]) + rest_shape_added_h = np.hstack( + (v_posed_added, np.ones([v_posed_added.shape[0], 1]))) + v_added = np.matmul(T_added, + rest_shape_added_h.reshape([-1, 4, + 1])).reshape([-1, 4 + ])[:, :3] + self.verts_added = v_added + self.trans.reshape([1, 3]) + + def rodrigues(self, r): + """ + Rodrigues' rotation formula that turns axis-angle vector into rotation + matrix in a batch-ed manner. + + Parameter: + ---------- + r: Axis-angle rotation vector of shape [batch_size, 1, 3]. + + Return: + ------- + Rotation matrix of shape [batch_size, 3, 3]. + + """ + theta = np.linalg.norm(r, axis=(1, 2), keepdims=True) + # avoid zero divide + theta = np.maximum(theta, np.finfo(np.float64).tiny) + r_hat = r / theta + cos = np.cos(theta) + z_stick = np.zeros(theta.shape[0]) + m = np.dstack([ + z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, + -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick + ]).reshape([-1, 3, 3]) + i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), + [theta.shape[0], 3, 3]) + A = np.transpose(r_hat, axes=[0, 2, 1]) + B = r_hat + dot = np.matmul(A, B) + R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m + return R + + def with_zeros(self, x): + """ + Append a [0, 0, 0, 1] vector to a [3, 4] matrix. + + Parameter: + --------- + x: Matrix to be appended. + + Return: + ------ + Matrix after appending of shape [4,4] + + """ + return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]]))) + + def pack(self, x): + """ + Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched + manner. + + Parameter: + ---------- + x: Matrices to be appended of shape [batch_size, 4, 1] + + Return: + ------ + Matrix of shape [batch_size, 4, 4] after appending. + + """ + return np.dstack((np.zeros((x.shape[0], 4, 3)), x)) + + def save_mesh_to_obj(self, path): + """ + Save the SMPL model into .obj file. + + Parameter: + --------- + path: Path to save. + + """ + with open(path, 'w') as fp: + for v in self.verts: + fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) + for f in self.faces + 1: + fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) + + def save_tetrahedron_to_obj(self, path): + """ + Save the tetrahedron SMPL model into .obj file. + + Parameter: + --------- + path: Path to save. + + """ + + with open(path, 'w') as fp: + for v in self.verts: + fp.write('v %f %f %f 1 0 0\n' % (v[0], v[1], v[2])) + for va in self.verts_added: + fp.write('v %f %f %f 0 0 1\n' % (va[0], va[1], va[2])) + for t in self.tetrahedrons + 1: + fp.write('f %d %d %d\n' % (t[0], t[2], t[1])) + fp.write('f %d %d %d\n' % (t[0], t[3], t[2])) + fp.write('f %d %d %d\n' % (t[0], t[1], t[3])) + fp.write('f %d %d %d\n' % (t[1], t[2], t[3])) diff --git a/lib/dataset/hoppeMesh.py b/lib/dataset/hoppeMesh.py new file mode 100644 index 0000000000000000000000000000000000000000..f671f8ee9adf5cc1649bdbaee15ac8db08caa67a --- /dev/null +++ b/lib/dataset/hoppeMesh.py @@ -0,0 +1,116 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np +from scipy.spatial import cKDTree +import trimesh + +import logging + +logging.getLogger("trimesh").setLevel(logging.ERROR) + + +def save_obj_mesh(mesh_path, verts, faces): + file = open(mesh_path, 'w') + for v in verts: + file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) + for f in faces: + f_plus = f + 1 + file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) + file.close() + + +def save_obj_mesh_with_color(mesh_path, verts, faces, colors): + file = open(mesh_path, 'w') + + for idx, v in enumerate(verts): + c = colors[idx] + file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % + (v[0], v[1], v[2], c[0], c[1], c[2])) + for f in faces: + f_plus = f + 1 + file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) + file.close() + + +def save_ply(mesh_path, points, rgb): + ''' + Save the visualization of sampling to a ply file. + Red points represent positive predictions. + Green points represent negative predictions. + :param mesh_path: File name to save + :param points: [N, 3] array of points + :param rgb: [N, 3] array of rgb values in the range [0~1] + :return: + ''' + to_save = np.concatenate([points, rgb * 255], axis=-1) + return np.savetxt( + mesh_path, + to_save, + fmt='%.6f %.6f %.6f %d %d %d', + comments='', + header=( + 'ply\nformat ascii 1.0\nelement vertex {:d}\n' + + 'property float x\nproperty float y\nproperty float z\n' + + 'property uchar red\nproperty uchar green\nproperty uchar blue\n' + + 'end_header').format(points.shape[0])) + + +class HoppeMesh: + def __init__(self, verts, faces, vert_normals, face_normals): + ''' + The HoppeSDF calculates signed distance towards a predefined oriented point cloud + http://hhoppe.com/recon.pdf + For clean and high-resolution pcl data, this is the fastest and accurate approximation of sdf + :param points: pts + :param normals: normals + ''' + self.verts = verts # [n, 3] + self.faces = faces # [m, 3] + self.vert_normals = vert_normals # [n, 3] + self.face_normals = face_normals # [m, 3] + + self.kd_tree = cKDTree(self.verts) + self.len = len(self.verts) + + def query(self, points): + dists, idx = self.kd_tree.query(points, n_jobs=1) + # FIXME: because the eyebows are removed, cKDTree around eyebows + # are not accurate. Cause a few false-inside labels here. + dirs = points - self.verts[idx] + signs = (dirs * self.vert_normals[idx]).sum(axis=1) + signs = (signs > 0) * 2 - 1 + return signs * dists + + def contains(self, points): + + labels = trimesh.Trimesh(vertices=self.verts, + faces=self.faces).contains(points) + return labels + + def export(self, path): + if self.colors is not None: + save_obj_mesh_with_color(path, self.verts, self.faces, + self.colors[:, 0:3] / 255.0) + else: + save_obj_mesh(path, self.verts, self.faces) + + def export_ply(self, path): + save_ply(path, self.verts, self.colors[:, 0:3] / 255.0) + + def triangles(self): + return self.verts[self.faces] # [n, 3, 3] diff --git a/lib/dataset/mesh_util.py b/lib/dataset/mesh_util.py new file mode 100644 index 0000000000000000000000000000000000000000..1e2ec31b52a1ebe41061ad2b5b7f5c6ee0c0e817 --- /dev/null +++ b/lib/dataset/mesh_util.py @@ -0,0 +1,899 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np +import cv2 +import pymeshlab +import torch +import torchvision +import trimesh +from pytorch3d.io import load_obj +import os +from termcolor import colored +import os.path as osp +from scipy.spatial import cKDTree + +from pytorch3d.structures import Meshes +import torch.nn.functional as F +from lib.pymaf.utils.imutils import uncrop +from lib.common.render_utils import Pytorch3dRasterizer, face_vertices + +from pytorch3d.renderer.mesh import rasterize_meshes +from PIL import Image, ImageFont, ImageDraw +from kaolin.ops.mesh import check_sign +from kaolin.metrics.trianglemesh import point_to_mesh_distance + +from pytorch3d.loss import ( + mesh_laplacian_smoothing, + mesh_normal_consistency +) + + +def tensor2variable(tensor, device): + # [1,23,3,3] + return torch.tensor(tensor, device=device, requires_grad=True) + + +def normal_loss(vec1, vec2): + + # vec1_mask = vec1.sum(dim=1) != 0.0 + # vec2_mask = vec2.sum(dim=1) != 0.0 + # union_mask = vec1_mask * vec2_mask + vec_sim = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(vec1, vec2) + # vec_diff = ((vec_sim-1.0)**2)[union_mask].mean() + vec_diff = ((vec_sim-1.0)**2).mean() + + return vec_diff + + +class GMoF(torch.nn.Module): + def __init__(self, rho=1): + super(GMoF, self).__init__() + self.rho = rho + + def extra_repr(self): + return 'rho = {}'.format(self.rho) + + def forward(self, residual): + dist = torch.div(residual, residual + self.rho ** 2) + return self.rho ** 2 * dist + + +def mesh_edge_loss(meshes, target_length: float = 0.0): + """ + Computes mesh edge length regularization loss averaged across all meshes + in a batch. Each mesh contributes equally to the final loss, regardless of + the number of edges per mesh in the batch by weighting each mesh with the + inverse number of edges. For example, if mesh 3 (out of N) has only E=4 + edges, then the loss for each edge in mesh 3 should be multiplied by 1/E to + contribute to the final loss. + + Args: + meshes: Meshes object with a batch of meshes. + target_length: Resting value for the edge length. + + Returns: + loss: Average loss across the batch. Returns 0 if meshes contains + no meshes or all empty meshes. + """ + if meshes.isempty(): + return torch.tensor( + [0.0], dtype=torch.float32, device=meshes.device, requires_grad=True + ) + + N = len(meshes) + edges_packed = meshes.edges_packed() # (sum(E_n), 3) + verts_packed = meshes.verts_packed() # (sum(V_n), 3) + edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx() # (sum(E_n), ) + num_edges_per_mesh = meshes.num_edges_per_mesh() # N + + # Determine the weight for each edge based on the number of edges in the + # mesh it corresponds to. + # TODO (nikhilar) Find a faster way of computing the weights for each edge + # as this is currently a bottleneck for meshes with a large number of faces. + weights = num_edges_per_mesh.gather(0, edge_to_mesh_idx) + weights = 1.0 / weights.float() + + verts_edges = verts_packed[edges_packed] + v0, v1 = verts_edges.unbind(1) + loss = ((v0 - v1).norm(dim=1, p=2) - target_length) ** 2.0 + loss_vertex = loss * weights + # loss_outlier = torch.topk(loss, 100)[0].mean() + # loss_all = (loss_vertex.sum() + loss_outlier.mean()) / N + loss_all = loss_vertex.sum() / N + + return loss_all + + +def remesh(obj_path, perc, device): + + ms = pymeshlab.MeshSet() + ms.load_new_mesh(obj_path) + ms.laplacian_smooth() + ms.remeshing_isotropic_explicit_remeshing( + targetlen=pymeshlab.Percentage(perc), adaptive=True) + ms.save_current_mesh(obj_path.replace("recon", "remesh")) + polished_mesh = trimesh.load_mesh(obj_path.replace("recon", "remesh")) + verts_pr = torch.tensor(polished_mesh.vertices).float().unsqueeze(0).to(device) + faces_pr = torch.tensor(polished_mesh.faces).long().unsqueeze(0).to(device) + + return verts_pr, faces_pr + + +def possion(mesh, obj_path): + + mesh.export(obj_path) + ms = pymeshlab.MeshSet() + ms.load_new_mesh(obj_path) + ms.surface_reconstruction_screened_poisson(depth=10) + ms.set_current_mesh(1) + ms.save_current_mesh(obj_path) + + return trimesh.load(obj_path) + + +def get_mask(tensor, dim): + + mask = torch.abs(tensor).sum(dim=dim, keepdims=True) > 0.0 + mask = mask.type_as(tensor) + + return mask + + +def blend_rgb_norm(rgb, norm, mask): + + # [0,0,0] or [127,127,127] should be marked as mask + final = rgb * (1-mask) + norm * (mask) + + return final.astype(np.uint8) + + +def unwrap(image, data): + + img_uncrop = uncrop(np.array(Image.fromarray(image).resize(data['uncrop_param']['box_shape'][:2])), + data['uncrop_param']['center'], + data['uncrop_param']['scale'], + data['uncrop_param']['crop_shape']) + + img_orig = cv2.warpAffine(img_uncrop, + np.linalg.inv(data['uncrop_param']['M'])[:2, :], + data['uncrop_param']['ori_shape'][::-1][1:], + flags=cv2.INTER_CUBIC) + + return img_orig + + +# Losses to smooth / regularize the mesh shape +def update_mesh_shape_prior_losses(mesh, losses): + + # and (b) the edge length of the predicted mesh + losses["edge"]['value'] = mesh_edge_loss(mesh) + # mesh normal consistency + losses["nc"]['value'] = mesh_normal_consistency(mesh) + # mesh laplacian smoothing + losses["laplacian"]['value'] = mesh_laplacian_smoothing( + mesh, method="uniform") + + +def rename(old_dict, old_name, new_name): + new_dict = {} + for key, value in zip(old_dict.keys(), old_dict.values()): + new_key = key if key != old_name else new_name + new_dict[new_key] = old_dict[key] + return new_dict + + +def load_checkpoint(model, cfg): + + model_dict = model.state_dict() + main_dict = {} + normal_dict = {} + + device = torch.device(f"cuda:{cfg['test_gpus'][0]}") + + if os.path.exists(cfg.resume_path) and cfg.resume_path.endswith("ckpt"): + main_dict = torch.load(cfg.resume_path, + map_location=device)['state_dict'] + + main_dict = { + k: v + for k, v in main_dict.items() + if k in model_dict and v.shape == model_dict[k].shape and ( + 'reconEngine' not in k) and ("normal_filter" not in k) and ( + 'voxelization' not in k) + } + print(colored(f"Resume MLP weights from {cfg.resume_path}", 'green')) + + if os.path.exists(cfg.normal_path) and cfg.normal_path.endswith("ckpt"): + normal_dict = torch.load(cfg.normal_path, + map_location=device)['state_dict'] + + for key in normal_dict.keys(): + normal_dict = rename(normal_dict, key, + key.replace("netG", "netG.normal_filter")) + + normal_dict = { + k: v + for k, v in normal_dict.items() + if k in model_dict and v.shape == model_dict[k].shape + } + print(colored(f"Resume normal model from {cfg.normal_path}", 'green')) + + model_dict.update(main_dict) + model_dict.update(normal_dict) + model.load_state_dict(model_dict) + + model.netG = model.netG.to(device) + model.reconEngine = model.reconEngine.to(device) + + model.netG.training = False + model.netG.eval() + + del main_dict + del normal_dict + del model_dict + + return model + + +def read_smpl_constants(folder): + """Load smpl vertex code""" + smpl_vtx_std = np.loadtxt(os.path.join(folder, 'vertices.txt')) + min_x = np.min(smpl_vtx_std[:, 0]) + max_x = np.max(smpl_vtx_std[:, 0]) + min_y = np.min(smpl_vtx_std[:, 1]) + max_y = np.max(smpl_vtx_std[:, 1]) + min_z = np.min(smpl_vtx_std[:, 2]) + max_z = np.max(smpl_vtx_std[:, 2]) + + smpl_vtx_std[:, 0] = (smpl_vtx_std[:, 0] - min_x) / (max_x - min_x) + smpl_vtx_std[:, 1] = (smpl_vtx_std[:, 1] - min_y) / (max_y - min_y) + smpl_vtx_std[:, 2] = (smpl_vtx_std[:, 2] - min_z) / (max_z - min_z) + smpl_vertex_code = np.float32(np.copy(smpl_vtx_std)) + """Load smpl faces & tetrahedrons""" + smpl_faces = np.loadtxt(os.path.join(folder, 'faces.txt'), + dtype=np.int32) - 1 + smpl_face_code = (smpl_vertex_code[smpl_faces[:, 0]] + + smpl_vertex_code[smpl_faces[:, 1]] + + smpl_vertex_code[smpl_faces[:, 2]]) / 3.0 + smpl_tetras = np.loadtxt(os.path.join(folder, 'tetrahedrons.txt'), + dtype=np.int32) - 1 + + return smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras + + +def feat_select(feat, select): + + # feat [B, featx2, N] + # select [B, 1, N] + # return [B, feat, N] + + dim = feat.shape[1] // 2 + idx = torch.tile((1-select), (1, dim, 1))*dim + \ + torch.arange(0, dim).unsqueeze(0).unsqueeze(2).type_as(select) + feat_select = torch.gather(feat, 1, idx.long()) + + return feat_select + + +def get_visibility(xy, z, faces): + """get the visibility of vertices + + Args: + xy (torch.tensor): [N,2] + z (torch.tensor): [N,1] + faces (torch.tensor): [N,3] + size (int): resolution of rendered image + """ + + xyz = torch.cat((xy, -z), dim=1) + xyz = (xyz + 1.0) / 2.0 + faces = faces.long() + + rasterizer = Pytorch3dRasterizer(image_size=2**12) + meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...]) + raster_settings = rasterizer.raster_settings + + pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( + meshes_screen, + image_size=raster_settings.image_size, + blur_radius=raster_settings.blur_radius, + faces_per_pixel=raster_settings.faces_per_pixel, + bin_size=raster_settings.bin_size, + max_faces_per_bin=raster_settings.max_faces_per_bin, + perspective_correct=raster_settings.perspective_correct, + cull_backfaces=raster_settings.cull_backfaces, + ) + + vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :]) + vis_mask = torch.zeros(size=(z.shape[0], 1)) + vis_mask[vis_vertices_id] = 1.0 + + # print("------------------------\n") + # print(f"keep points : {vis_mask.sum()/len(vis_mask)}") + + return vis_mask + + +def barycentric_coordinates_of_projection(points, vertices): + ''' https://github.com/MPI-IS/mesh/blob/master/mesh/geometry/barycentric_coordinates_of_projection.py + ''' + """Given a point, gives projected coords of that point to a triangle + in barycentric coordinates. + See + **Heidrich**, Computing the Barycentric Coordinates of a Projected Point, JGT 05 + at http://www.cs.ubc.ca/~heidrich/Papers/JGT.05.pdf + + :param p: point to project. [B, 3] + :param v0: first vertex of triangles. [B, 3] + :returns: barycentric coordinates of ``p``'s projection in triangle defined by ``q``, ``u``, ``v`` + vectorized so ``p``, ``q``, ``u``, ``v`` can all be ``3xN`` + """ + #(p, q, u, v) + v0, v1, v2 = vertices[:, 0], vertices[:, 1], vertices[:, 2] + p = points + + q = v0 + u = v1 - v0 + v = v2 - v0 + n = torch.cross(u, v) + s = torch.sum(n * n, dim=1) + # If the triangle edges are collinear, cross-product is zero, + # which makes "s" 0, which gives us divide by zero. So we + # make the arbitrary choice to set s to epsv (=numpy.spacing(1)), + # the closest thing to zero + s[s == 0] = 1e-6 + oneOver4ASquared = 1.0 / s + w = p - q + b2 = torch.sum(torch.cross(u, w) * n, dim=1) * oneOver4ASquared + b1 = torch.sum(torch.cross(w, v) * n, dim=1) * oneOver4ASquared + weights = torch.stack((1 - b1 - b2, b1, b2), dim=-1) + # check barycenric weights + # p_n = v0*weights[:,0:1] + v1*weights[:,1:2] + v2*weights[:,2:3] + return weights + + +def cal_sdf_batch(verts, faces, cmaps, vis, points): + + # verts [B, N_vert, 3] + # faces [B, N_face, 3] + # triangles [B, N_face, 3, 3] + # points [B, N_point, 3] + # cmaps [B, N_vert, 3] + + Bsize = points.shape[0] + + normals = Meshes(verts, faces).verts_normals_padded() + + triangles = face_vertices(verts, faces) + normals = face_vertices(normals, faces) + cmaps = face_vertices(cmaps, faces) + vis = face_vertices(vis, faces) + + residues, pts_ind, _ = point_to_mesh_distance(points, triangles) + closest_triangles = torch.gather( + triangles, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) + closest_normals = torch.gather( + normals, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) + closest_cmaps = torch.gather( + cmaps, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) + closest_vis = torch.gather( + vis, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 1)).view(-1, 3, 1) + bary_weights = barycentric_coordinates_of_projection( + points.view(-1, 3), closest_triangles) + + pts_cmap = (closest_cmaps*bary_weights[:, :, None]).sum(1).unsqueeze(0) + pts_vis = (closest_vis*bary_weights[:, + :, None]).sum(1).unsqueeze(0).ge(1e-1) + pts_norm = (closest_normals*bary_weights[:, :, None]).sum( + 1).unsqueeze(0) * torch.tensor([-1.0, 1.0, -1.0]).type_as(normals) + pts_dist = torch.sqrt(residues) / torch.sqrt(torch.tensor(3)) + + pts_signs = 2.0 * (check_sign(verts, faces[0], points).float() - 0.5) + pts_sdf = (pts_dist * pts_signs).unsqueeze(-1) + + return pts_sdf.view(Bsize, -1, 1), pts_norm.view(Bsize, -1, 3), pts_cmap.view(Bsize, -1, 3), pts_vis.view(Bsize, -1, 1) + + +def orthogonal(points, calibrations, transforms=None): + ''' + Compute the orthogonal projections of 3D points into the image plane by given projection matrix + :param points: [B, 3, N] Tensor of 3D points + :param calibrations: [B, 3, 4] Tensor of projection matrix + :param transforms: [B, 2, 3] Tensor of image transform matrix + :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane + ''' + rot = calibrations[:, :3, :3] + trans = calibrations[:, :3, 3:4] + pts = torch.baddbmm(trans, rot, points) # [B, 3, N] + if transforms is not None: + scale = transforms[:2, :2] + shift = transforms[:2, 2:3] + pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :]) + return pts + + +def projection(points, calib, format='numpy'): + if format == 'tensor': + return torch.mm(calib[:3, :3], points.T).T + calib[:3, 3] + else: + return np.matmul(calib[:3, :3], points.T).T + calib[:3, 3] + + +def load_calib(calib_path): + calib_data = np.loadtxt(calib_path, dtype=float) + extrinsic = calib_data[:4, :4] + intrinsic = calib_data[4:8, :4] + calib_mat = np.matmul(intrinsic, extrinsic) + calib_mat = torch.from_numpy(calib_mat).float() + return calib_mat + + +def load_obj_mesh_for_Hoppe(mesh_file): + vertex_data = [] + face_data = [] + + if isinstance(mesh_file, str): + f = open(mesh_file, "r") + else: + f = mesh_file + for line in f: + if isinstance(line, bytes): + line = line.decode("utf-8") + if line.startswith('#'): + continue + values = line.split() + if not values: + continue + + if values[0] == 'v': + v = list(map(float, values[1:4])) + vertex_data.append(v) + + elif values[0] == 'f': + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + f = list( + map(lambda x: int(x.split('/')[0]), + [values[3], values[4], values[1]])) + face_data.append(f) + # tri mesh + else: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + + vertices = np.array(vertex_data) + faces = np.array(face_data) + faces[faces > 0] -= 1 + + normals, _ = compute_normal(vertices, faces) + + return vertices, normals, faces + + +def load_obj_mesh_with_color(mesh_file): + vertex_data = [] + color_data = [] + face_data = [] + + if isinstance(mesh_file, str): + f = open(mesh_file, "r") + else: + f = mesh_file + for line in f: + if isinstance(line, bytes): + line = line.decode("utf-8") + if line.startswith('#'): + continue + values = line.split() + if not values: + continue + + if values[0] == 'v': + v = list(map(float, values[1:4])) + vertex_data.append(v) + c = list(map(float, values[4:7])) + color_data.append(c) + + elif values[0] == 'f': + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + f = list( + map(lambda x: int(x.split('/')[0]), + [values[3], values[4], values[1]])) + face_data.append(f) + # tri mesh + else: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + + vertices = np.array(vertex_data) + colors = np.array(color_data) + faces = np.array(face_data) + faces[faces > 0] -= 1 + + return vertices, colors, faces + + +def load_obj_mesh(mesh_file, with_normal=False, with_texture=False): + vertex_data = [] + norm_data = [] + uv_data = [] + + face_data = [] + face_norm_data = [] + face_uv_data = [] + + if isinstance(mesh_file, str): + f = open(mesh_file, "r") + else: + f = mesh_file + for line in f: + if isinstance(line, bytes): + line = line.decode("utf-8") + if line.startswith('#'): + continue + values = line.split() + if not values: + continue + + if values[0] == 'v': + v = list(map(float, values[1:4])) + vertex_data.append(v) + elif values[0] == 'vn': + vn = list(map(float, values[1:4])) + norm_data.append(vn) + elif values[0] == 'vt': + vt = list(map(float, values[1:3])) + uv_data.append(vt) + + elif values[0] == 'f': + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + f = list( + map(lambda x: int(x.split('/')[0]), + [values[3], values[4], values[1]])) + face_data.append(f) + # tri mesh + else: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + + # deal with texture + if len(values[1].split('/')) >= 2: + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) + face_uv_data.append(f) + f = list( + map(lambda x: int(x.split('/')[1]), + [values[3], values[4], values[1]])) + face_uv_data.append(f) + # tri mesh + elif len(values[1].split('/')[1]) != 0: + f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) + face_uv_data.append(f) + # deal with normal + if len(values[1].split('/')) == 3: + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) + face_norm_data.append(f) + f = list( + map(lambda x: int(x.split('/')[2]), + [values[3], values[4], values[1]])) + face_norm_data.append(f) + # tri mesh + elif len(values[1].split('/')[2]) != 0: + f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) + face_norm_data.append(f) + + vertices = np.array(vertex_data) + faces = np.array(face_data) + faces[faces > 0] -= 1 + + if with_texture and with_normal: + uvs = np.array(uv_data) + face_uvs = np.array(face_uv_data) + face_uvs[face_uvs > 0] -= 1 + norms = np.array(norm_data) + if norms.shape[0] == 0: + norms, _ = compute_normal(vertices, faces) + face_normals = faces + else: + norms = normalize_v3(norms) + face_normals = np.array(face_norm_data) + face_normals[face_normals > 0] -= 1 + return vertices, faces, norms, face_normals, uvs, face_uvs + + if with_texture: + uvs = np.array(uv_data) + face_uvs = np.array(face_uv_data) - 1 + return vertices, faces, uvs, face_uvs + + if with_normal: + norms = np.array(norm_data) + norms = normalize_v3(norms) + face_normals = np.array(face_norm_data) - 1 + return vertices, faces, norms, face_normals + + return vertices, faces + + +def normalize_v3(arr): + ''' Normalize a numpy array of 3 component vectors shape=(n,3) ''' + lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2) + eps = 0.00000001 + lens[lens < eps] = eps + arr[:, 0] /= lens + arr[:, 1] /= lens + arr[:, 2] /= lens + return arr + + +def compute_normal(vertices, faces): + # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal + vert_norms = np.zeros(vertices.shape, dtype=vertices.dtype) + # Create an indexed view into the vertex array using the array of three indices for triangles + tris = vertices[faces] + # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle + face_norms = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0]) + # n is now an array of normals per triangle. The length of each normal is dependent the vertices, + # we need to normalize these, so that our next step weights each normal equally. + normalize_v3(face_norms) + # now we have a normalized array of normals, one per triangle, i.e., per triangle normals. + # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle, + # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards. + # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array + vert_norms[faces[:, 0]] += face_norms + vert_norms[faces[:, 1]] += face_norms + vert_norms[faces[:, 2]] += face_norms + normalize_v3(vert_norms) + + return vert_norms, face_norms + + +def save_obj_mesh(mesh_path, verts, faces): + file = open(mesh_path, 'w') + for v in verts: + file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) + for f in faces: + f_plus = f + 1 + file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) + file.close() + + +def save_obj_mesh_with_color(mesh_path, verts, faces, colors): + file = open(mesh_path, 'w') + + for idx, v in enumerate(verts): + c = colors[idx] + file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % + (v[0], v[1], v[2], c[0], c[1], c[2])) + for f in faces: + f_plus = f + 1 + file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) + file.close() + + +def calculate_mIoU(outputs, labels): + + SMOOTH = 1e-6 + + outputs = outputs.int() + labels = labels.int() + + intersection = ( + outputs + & labels).float().sum() # Will be zero if Truth=0 or Prediction=0 + union = (outputs | labels).float().sum() # Will be zzero if both are 0 + + iou = (intersection + SMOOTH) / (union + SMOOTH + ) # We smooth our devision to avoid 0/0 + + thresholded = torch.clamp( + 20 * (iou - 0.5), 0, + 10).ceil() / 10 # This is equal to comparing with thresolds + + return thresholded.mean().detach().cpu().numpy( + ) # Or thresholded.mean() if you are interested in average across the batch + + +def mask_filter(mask, number=1000): + """only keep {number} True items within a mask + + Args: + mask (bool array): [N, ] + number (int, optional): total True item. Defaults to 1000. + """ + true_ids = np.where(mask)[0] + keep_ids = np.random.choice(true_ids, size=number) + filter_mask = np.isin(np.arange(len(mask)), keep_ids) + + return filter_mask + + +def query_mesh(path): + + verts, faces_idx, _ = load_obj(path) + + return verts, faces_idx.verts_idx + + +def add_alpha(colors, alpha=0.7): + + colors_pad = np.pad(colors, ((0, 0), (0, 1)), + mode='constant', + constant_values=alpha) + + return colors_pad + + +def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type='smpl'): + + font_path = os.path.join(os.path.dirname(__file__), "tbfo.ttf") + font = ImageFont.truetype(font_path, 30) + grid_img = torchvision.utils.make_grid(torch.cat(per_loop_lst, dim=0), + nrow=nrow) + grid_img = Image.fromarray( + ((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 * + 255.0).astype(np.uint8)) + + # add text + draw = ImageDraw.Draw(grid_img) + grid_size = 512 + if loss is not None: + draw.text((10, 5), f"error: {loss:.3f}", (255, 0, 0), font=font) + + if type == 'smpl': + for col_id, col_txt in enumerate( + ['image', 'smpl-norm(render)', 'cloth-norm(pred)', 'diff-norm', 'diff-mask']): + draw.text((10+(col_id*grid_size), 5), + col_txt, (255, 0, 0), font=font) + elif type == 'cloth': + for col_id, col_txt in enumerate( + ['image', 'cloth-norm(recon)', 'cloth-norm(pred)', 'diff-norm']): + draw.text((10+(col_id*grid_size), 5), + col_txt, (255, 0, 0), font=font) + for col_id, col_txt in enumerate( + ['0', '90', '180', '270']): + draw.text((10+(col_id*grid_size), grid_size*2+5), + col_txt, (255, 0, 0), font=font) + else: + print(f"{type} should be 'smpl' or 'cloth'") + + grid_img = grid_img.resize((grid_img.size[0], grid_img.size[1]), + Image.ANTIALIAS) + + return grid_img + + +def clean_mesh(verts, faces): + + device = verts.device + + mesh_lst = trimesh.Trimesh(verts.detach().cpu().numpy(), + faces.detach().cpu().numpy()) + mesh_lst = mesh_lst.split(only_watertight=False) + comp_num = [mesh.vertices.shape[0] for mesh in mesh_lst] + mesh_clean = mesh_lst[comp_num.index(max(comp_num))] + + final_verts = torch.as_tensor(mesh_clean.vertices).float().to(device) + final_faces = torch.as_tensor(mesh_clean.faces).int().to(device) + + return final_verts, final_faces + + +def merge_mesh(verts_A, faces_A, verts_B, faces_B, color=False): + + sep_mesh = trimesh.Trimesh(np.concatenate([verts_A, verts_B], axis=0), + np.concatenate( + [faces_A, faces_B + faces_A.max() + 1], + axis=0), + maintain_order=True, + process=False) + if color: + colors = np.ones_like(sep_mesh.vertices) + colors[:verts_A.shape[0]] *= np.array([255.0, 0.0, 0.0]) + colors[verts_A.shape[0]:] *= np.array([0.0, 255.0, 0.0]) + sep_mesh.visual.vertex_colors = colors + + # union_mesh = trimesh.boolean.union([trimesh.Trimesh(verts_A, faces_A), + # trimesh.Trimesh(verts_B, faces_B)], engine='blender') + + return sep_mesh + + +def mesh_move(mesh_lst, step, scale=1.0): + + trans = np.array([1.0, 0.0, 0.0]) * step + + resize_matrix = trimesh.transformations.scale_and_translate( + scale=(scale), translate=trans) + + results = [] + + for mesh in mesh_lst: + mesh.apply_transform(resize_matrix) + results.append(mesh) + + return results + + +class SMPLX(): + def __init__(self): + + self.current_dir = osp.join(osp.dirname(__file__), + "../../data/smpl_related") + + self.smpl_verts_path = osp.join(self.current_dir, + "smpl_data/smpl_verts.npy") + self.smplx_verts_path = osp.join(self.current_dir, + "smpl_data/smplx_verts.npy") + self.faces_path = osp.join(self.current_dir, + "smpl_data/smplx_faces.npy") + self.cmap_vert_path = osp.join(self.current_dir, + "smpl_data/smplx_cmap.npy") + + self.faces = np.load(self.faces_path) + self.verts = np.load(self.smplx_verts_path) + self.smpl_verts = np.load(self.smpl_verts_path) + + self.model_dir = osp.join(self.current_dir, "models") + self.tedra_dir = osp.join(self.current_dir, "../tedra_data") + + def get_smpl_mat(self, vert_ids): + + mat = torch.as_tensor(np.load(self.cmap_vert_path)).float() + return mat[vert_ids, :] + + def smpl2smplx(self, vert_ids=None): + """convert vert_ids in smpl to vert_ids in smplx + + Args: + vert_ids ([int.array]): [n, knn_num] + """ + smplx_tree = cKDTree(self.verts, leafsize=1) + _, ind = smplx_tree.query(self.smpl_verts, k=1) # ind: [smpl_num, 1] + + if vert_ids is not None: + smplx_vert_ids = ind[vert_ids] + else: + smplx_vert_ids = ind + + return smplx_vert_ids + + def smplx2smpl(self, vert_ids=None): + """convert vert_ids in smplx to vert_ids in smpl + + Args: + vert_ids ([int.array]): [n, knn_num] + """ + smpl_tree = cKDTree(self.smpl_verts, leafsize=1) + _, ind = smpl_tree.query(self.verts, k=1) # ind: [smplx_num, 1] + if vert_ids is not None: + smpl_vert_ids = ind[vert_ids] + else: + smpl_vert_ids = ind + + return smpl_vert_ids diff --git a/lib/dataset/tbfo.ttf b/lib/dataset/tbfo.ttf new file mode 100644 index 0000000000000000000000000000000000000000..6cc76fcd568a5a42edd71272a19b15214de0b0d5 Binary files /dev/null and b/lib/dataset/tbfo.ttf differ diff --git a/lib/net/BasePIFuNet.py b/lib/net/BasePIFuNet.py new file mode 100644 index 0000000000000000000000000000000000000000..d1e5986298a38fc993b697169ec95b1007b58755 --- /dev/null +++ b/lib/net/BasePIFuNet.py @@ -0,0 +1,84 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import torch.nn as nn +import pytorch_lightning as pl + +from .geometry import index, orthogonal, perspective + + +class BasePIFuNet(pl.LightningModule): + def __init__( + self, + projection_mode='orthogonal', + error_term=nn.MSELoss(), + ): + """ + :param projection_mode: + Either orthogonal or perspective. + It will call the corresponding function for projection. + :param error_term: + nn Loss between the predicted [B, Res, N] and the label [B, Res, N] + """ + super(BasePIFuNet, self).__init__() + self.name = 'base' + + self.error_term = error_term + + self.index = index + self.projection = orthogonal if projection_mode == 'orthogonal' else perspective + + def forward(self, points, images, calibs, transforms=None): + ''' + :param points: [B, 3, N] world space coordinates of points + :param images: [B, C, H, W] input images + :param calibs: [B, 3, 4] calibration matrices for each image + :param transforms: Optional [B, 2, 3] image space coordinate transforms + :return: [B, Res, N] predictions for each point + ''' + features = self.filter(images) + preds = self.query(features, points, calibs, transforms) + return preds + + def filter(self, images): + ''' + Filter the input images + store all intermediate features. + :param images: [B, C, H, W] input images + ''' + return None + + def query(self, features, points, calibs, transforms=None): + ''' + Given 3D points, query the network predictions for each point. + Image features should be pre-computed before this call. + store all intermediate features. + query() function may behave differently during training/testing. + :param points: [B, 3, N] world space coordinates of points + :param calibs: [B, 3, 4] calibration matrices for each image + :param transforms: Optional [B, 2, 3] image space coordinate transforms + :param labels: Optional [B, Res, N] gt labeling + :return: [B, Res, N] predictions for each point + ''' + return None + + def get_error(self, preds, labels): + ''' + Get the network loss from the last query + :return: loss term + ''' + return self.error_term(preds, labels) diff --git a/lib/net/FBNet.py b/lib/net/FBNet.py new file mode 100644 index 0000000000000000000000000000000000000000..23541cc6559b7ae68c8a1f606f0637bb712542f1 --- /dev/null +++ b/lib/net/FBNet.py @@ -0,0 +1,387 @@ +''' +Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu. +BSD License. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. +IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL +DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, +WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING +OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +''' +import torch +import torch.nn as nn +import functools +import numpy as np +import pytorch_lightning as pl + + +############################################################################### +# Functions +############################################################################### +def weights_init(m): + classname = m.__class__.__name__ + if classname.find('Conv') != -1: + m.weight.data.normal_(0.0, 0.02) + elif classname.find('BatchNorm2d') != -1: + m.weight.data.normal_(1.0, 0.02) + m.bias.data.fill_(0) + + +def get_norm_layer(norm_type='instance'): + if norm_type == 'batch': + norm_layer = functools.partial(nn.BatchNorm2d, affine=True) + elif norm_type == 'instance': + norm_layer = functools.partial(nn.InstanceNorm2d, affine=False) + else: + raise NotImplementedError('normalization layer [%s] is not found' % + norm_type) + return norm_layer + + +def define_G(input_nc, + output_nc, + ngf, + netG, + n_downsample_global=3, + n_blocks_global=9, + n_local_enhancers=1, + n_blocks_local=3, + norm='instance', + gpu_ids=[], + last_op=nn.Tanh()): + norm_layer = get_norm_layer(norm_type=norm) + if netG == 'global': + netG = GlobalGenerator(input_nc, + output_nc, + ngf, + n_downsample_global, + n_blocks_global, + norm_layer, + last_op=last_op) + elif netG == 'local': + netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, + n_blocks_global, n_local_enhancers, + n_blocks_local, norm_layer) + elif netG == 'encoder': + netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, + norm_layer) + else: + raise ('generator not implemented!') + # print(netG) + if len(gpu_ids) > 0: + assert (torch.cuda.is_available()) + netG.cuda(gpu_ids[0]) + netG.apply(weights_init) + return netG + + +def print_network(net): + if isinstance(net, list): + net = net[0] + num_params = 0 + for param in net.parameters(): + num_params += param.numel() + print(net) + print('Total number of parameters: %d' % num_params) + + +############################################################################## +# Generator +############################################################################## +class LocalEnhancer(pl.LightningModule): + def __init__(self, + input_nc, + output_nc, + ngf=32, + n_downsample_global=3, + n_blocks_global=9, + n_local_enhancers=1, + n_blocks_local=3, + norm_layer=nn.BatchNorm2d, + padding_type='reflect'): + super(LocalEnhancer, self).__init__() + self.n_local_enhancers = n_local_enhancers + + ###### global generator model ##### + ngf_global = ngf * (2**n_local_enhancers) + model_global = GlobalGenerator(input_nc, output_nc, ngf_global, + n_downsample_global, n_blocks_global, + norm_layer).model + model_global = [model_global[i] for i in range(len(model_global) - 3) + ] # get rid of final convolution layers + self.model = nn.Sequential(*model_global) + + ###### local enhancer layers ##### + for n in range(1, n_local_enhancers + 1): + # downsample + ngf_global = ngf * (2**(n_local_enhancers - n)) + model_downsample = [ + nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), + norm_layer(ngf_global), + nn.ReLU(True), + nn.Conv2d(ngf_global, + ngf_global * 2, + kernel_size=3, + stride=2, + padding=1), + norm_layer(ngf_global * 2), + nn.ReLU(True) + ] + # residual blocks + model_upsample = [] + for i in range(n_blocks_local): + model_upsample += [ + ResnetBlock(ngf_global * 2, + padding_type=padding_type, + norm_layer=norm_layer) + ] + + # upsample + model_upsample += [ + nn.ConvTranspose2d(ngf_global * 2, + ngf_global, + kernel_size=3, + stride=2, + padding=1, + output_padding=1), + norm_layer(ngf_global), + nn.ReLU(True) + ] + + # final convolution + if n == n_local_enhancers: + model_upsample += [ + nn.ReflectionPad2d(3), + nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), + nn.Tanh() + ] + + setattr(self, 'model' + str(n) + '_1', + nn.Sequential(*model_downsample)) + setattr(self, 'model' + str(n) + '_2', + nn.Sequential(*model_upsample)) + + self.downsample = nn.AvgPool2d(3, + stride=2, + padding=[1, 1], + count_include_pad=False) + + def forward(self, input): + # create input pyramid + input_downsampled = [input] + for i in range(self.n_local_enhancers): + input_downsampled.append(self.downsample(input_downsampled[-1])) + + # output at coarest level + output_prev = self.model(input_downsampled[-1]) + # build up one layer at a time + for n_local_enhancers in range(1, self.n_local_enhancers + 1): + model_downsample = getattr(self, + 'model' + str(n_local_enhancers) + '_1') + model_upsample = getattr(self, + 'model' + str(n_local_enhancers) + '_2') + input_i = input_downsampled[self.n_local_enhancers - + n_local_enhancers] + output_prev = model_upsample( + model_downsample(input_i) + output_prev) + return output_prev + + +class GlobalGenerator(pl.LightningModule): + def __init__(self, + input_nc, + output_nc, + ngf=64, + n_downsampling=3, + n_blocks=9, + norm_layer=nn.BatchNorm2d, + padding_type='reflect', + last_op=nn.Tanh()): + assert (n_blocks >= 0) + super(GlobalGenerator, self).__init__() + activation = nn.ReLU(True) + + model = [ + nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), + norm_layer(ngf), activation + ] + # downsample + for i in range(n_downsampling): + mult = 2**i + model += [ + nn.Conv2d(ngf * mult, + ngf * mult * 2, + kernel_size=3, + stride=2, + padding=1), + norm_layer(ngf * mult * 2), activation + ] + + # resnet blocks + mult = 2**n_downsampling + for i in range(n_blocks): + model += [ + ResnetBlock(ngf * mult, + padding_type=padding_type, + activation=activation, + norm_layer=norm_layer) + ] + + # upsample + for i in range(n_downsampling): + mult = 2**(n_downsampling - i) + model += [ + nn.ConvTranspose2d(ngf * mult, + int(ngf * mult / 2), + kernel_size=3, + stride=2, + padding=1, + output_padding=1), + norm_layer(int(ngf * mult / 2)), activation + ] + model += [ + nn.ReflectionPad2d(3), + nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0) + ] + if last_op is not None: + model += [last_op] + self.model = nn.Sequential(*model) + + def forward(self, input): + return self.model(input) + + +# Define a resnet block +class ResnetBlock(pl.LightningModule): + def __init__(self, + dim, + padding_type, + norm_layer, + activation=nn.ReLU(True), + use_dropout=False): + super(ResnetBlock, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, + activation, use_dropout) + + def build_conv_block(self, dim, padding_type, norm_layer, activation, + use_dropout): + conv_block = [] + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % + padding_type) + + conv_block += [ + nn.Conv2d(dim, dim, kernel_size=3, padding=p), + norm_layer(dim), activation + ] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % + padding_type) + conv_block += [ + nn.Conv2d(dim, dim, kernel_size=3, padding=p), + norm_layer(dim) + ] + + return nn.Sequential(*conv_block) + + def forward(self, x): + out = x + self.conv_block(x) + return out + + +class Encoder(pl.LightningModule): + def __init__(self, + input_nc, + output_nc, + ngf=32, + n_downsampling=4, + norm_layer=nn.BatchNorm2d): + super(Encoder, self).__init__() + self.output_nc = output_nc + + model = [ + nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), + norm_layer(ngf), + nn.ReLU(True) + ] + # downsample + for i in range(n_downsampling): + mult = 2**i + model += [ + nn.Conv2d(ngf * mult, + ngf * mult * 2, + kernel_size=3, + stride=2, + padding=1), + norm_layer(ngf * mult * 2), + nn.ReLU(True) + ] + + # upsample + for i in range(n_downsampling): + mult = 2**(n_downsampling - i) + model += [ + nn.ConvTranspose2d(ngf * mult, + int(ngf * mult / 2), + kernel_size=3, + stride=2, + padding=1, + output_padding=1), + norm_layer(int(ngf * mult / 2)), + nn.ReLU(True) + ] + + model += [ + nn.ReflectionPad2d(3), + nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), + nn.Tanh() + ] + self.model = nn.Sequential(*model) + + def forward(self, input, inst): + outputs = self.model(input) + + # instance-wise average pooling + outputs_mean = outputs.clone() + inst_list = np.unique(inst.cpu().numpy().astype(int)) + for i in inst_list: + for b in range(input.size()[0]): + indices = (inst[b:b + 1] == int(i)).nonzero() # n x 4 + for j in range(self.output_nc): + output_ins = outputs[indices[:, 0] + b, indices[:, 1] + j, + indices[:, 2], indices[:, 3]] + mean_feat = torch.mean(output_ins).expand_as(output_ins) + outputs_mean[indices[:, 0] + b, indices[:, 1] + j, + indices[:, 2], indices[:, 3]] = mean_feat + return outputs_mean diff --git a/lib/net/HGFilters.py b/lib/net/HGFilters.py new file mode 100755 index 0000000000000000000000000000000000000000..b8578cc42fb6c2630fea884ea86e5d53ab5f6d5d --- /dev/null +++ b/lib/net/HGFilters.py @@ -0,0 +1,197 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.net.net_util import * +import torch.nn as nn +import torch.nn.functional as F + + +class HourGlass(nn.Module): + def __init__(self, num_modules, depth, num_features, opt): + super(HourGlass, self).__init__() + self.num_modules = num_modules + self.depth = depth + self.features = num_features + self.opt = opt + + self._generate_network(self.depth) + + def _generate_network(self, level): + self.add_module('b1_' + str(level), + ConvBlock(self.features, self.features, self.opt)) + + self.add_module('b2_' + str(level), + ConvBlock(self.features, self.features, self.opt)) + + if level > 1: + self._generate_network(level - 1) + else: + self.add_module('b2_plus_' + str(level), + ConvBlock(self.features, self.features, self.opt)) + + self.add_module('b3_' + str(level), + ConvBlock(self.features, self.features, self.opt)) + + def _forward(self, level, inp): + # Upper branch + up1 = inp + up1 = self._modules['b1_' + str(level)](up1) + + # Lower branch + low1 = F.avg_pool2d(inp, 2, stride=2) + low1 = self._modules['b2_' + str(level)](low1) + + if level > 1: + low2 = self._forward(level - 1, low1) + else: + low2 = low1 + low2 = self._modules['b2_plus_' + str(level)](low2) + + low3 = low2 + low3 = self._modules['b3_' + str(level)](low3) + + # NOTE: for newer PyTorch (1.3~), it seems that training results are degraded due to implementation diff in F.grid_sample + # if the pretrained model behaves weirdly, switch with the commented line. + # NOTE: I also found that "bicubic" works better. + up2 = F.interpolate(low3, + scale_factor=2, + mode='bicubic', + align_corners=True) + # up2 = F.interpolate(low3, scale_factor=2, mode='nearest) + + return up1 + up2 + + def forward(self, x): + return self._forward(self.depth, x) + + +class HGFilter(nn.Module): + def __init__(self, opt, num_modules, in_dim): + super(HGFilter, self).__init__() + self.num_modules = num_modules + + self.opt = opt + [k, s, d, p] = self.opt.conv1 + + # self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3) + self.conv1 = nn.Conv2d(in_dim, + 64, + kernel_size=k, + stride=s, + dilation=d, + padding=p) + + if self.opt.norm == 'batch': + self.bn1 = nn.BatchNorm2d(64) + elif self.opt.norm == 'group': + self.bn1 = nn.GroupNorm(32, 64) + + if self.opt.hg_down == 'conv64': + self.conv2 = ConvBlock(64, 64, self.opt) + self.down_conv2 = nn.Conv2d(64, + 128, + kernel_size=3, + stride=2, + padding=1) + elif self.opt.hg_down == 'conv128': + self.conv2 = ConvBlock(64, 128, self.opt) + self.down_conv2 = nn.Conv2d(128, + 128, + kernel_size=3, + stride=2, + padding=1) + elif self.opt.hg_down == 'ave_pool': + self.conv2 = ConvBlock(64, 128, self.opt) + else: + raise NameError('Unknown Fan Filter setting!') + + self.conv3 = ConvBlock(128, 128, self.opt) + self.conv4 = ConvBlock(128, 256, self.opt) + + # Stacking part + for hg_module in range(self.num_modules): + self.add_module('m' + str(hg_module), + HourGlass(1, opt.num_hourglass, 256, self.opt)) + + self.add_module('top_m_' + str(hg_module), + ConvBlock(256, 256, self.opt)) + self.add_module( + 'conv_last' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + if self.opt.norm == 'batch': + self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) + elif self.opt.norm == 'group': + self.add_module('bn_end' + str(hg_module), + nn.GroupNorm(32, 256)) + + self.add_module( + 'l' + str(hg_module), + nn.Conv2d(256, + opt.hourglass_dim, + kernel_size=1, + stride=1, + padding=0)) + + if hg_module < self.num_modules - 1: + self.add_module( + 'bl' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module( + 'al' + str(hg_module), + nn.Conv2d(opt.hourglass_dim, + 256, + kernel_size=1, + stride=1, + padding=0)) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x)), True) + tmpx = x + if self.opt.hg_down == 'ave_pool': + x = F.avg_pool2d(self.conv2(x), 2, stride=2) + elif self.opt.hg_down in ['conv64', 'conv128']: + x = self.conv2(x) + x = self.down_conv2(x) + else: + raise NameError('Unknown Fan Filter setting!') + + x = self.conv3(x) + x = self.conv4(x) + + previous = x + + outputs = [] + for i in range(self.num_modules): + hg = self._modules['m' + str(i)](previous) + + ll = hg + ll = self._modules['top_m_' + str(i)](ll) + + ll = F.relu( + self._modules['bn_end' + str(i)]( + self._modules['conv_last' + str(i)](ll)), True) + + # Predict heatmaps + tmp_out = self._modules['l' + str(i)](ll) + outputs.append(tmp_out) + + if i < self.num_modules - 1: + ll = self._modules['bl' + str(i)](ll) + tmp_out_ = self._modules['al' + str(i)](tmp_out) + previous = previous + ll + tmp_out_ + + return outputs diff --git a/lib/net/HGPIFuNet.py b/lib/net/HGPIFuNet.py new file mode 100644 index 0000000000000000000000000000000000000000..e80f553768177d59caec8d569712fd0dcf461174 --- /dev/null +++ b/lib/net/HGPIFuNet.py @@ -0,0 +1,410 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.net.voxelize import Voxelization +from lib.dataset.mesh_util import cal_sdf_batch, feat_select, read_smpl_constants +from lib.net.NormalNet import NormalNet +from lib.net.MLP import MLP +from lib.dataset.mesh_util import SMPLX +from lib.net.VE import VolumeEncoder +from lib.net.HGFilters import * +from termcolor import colored +from lib.net.BasePIFuNet import BasePIFuNet +import torch.nn as nn +import torch +import os + + +maskout = False + + +class HGPIFuNet(BasePIFuNet): + ''' + HG PIFu network uses Hourglass stacks as the image filter. + It does the following: + 1. Compute image feature stacks and store it in self.im_feat_list + self.im_feat_list[-1] is the last stack (output stack) + 2. Calculate calibration + 3. If training, it index on every intermediate stacks, + If testing, it index on the last stack. + 4. Classification. + 5. During training, error is calculated on all stacks. + ''' + + def __init__(self, + cfg, + projection_mode='orthogonal', + error_term=nn.MSELoss()): + + super(HGPIFuNet, self).__init__(projection_mode=projection_mode, + error_term=error_term) + + self.l1_loss = nn.SmoothL1Loss() + self.opt = cfg.net + self.root = cfg.root + self.overfit = cfg.overfit + + channels_IF = self.opt.mlp_dim + + self.use_filter = self.opt.use_filter + self.prior_type = self.opt.prior_type + self.smpl_feats = self.opt.smpl_feats + + self.smpl_dim = self.opt.smpl_dim + self.voxel_dim = self.opt.voxel_dim + self.hourglass_dim = self.opt.hourglass_dim + self.sdf_clip = cfg.sdf_clip / 100.0 + + self.in_geo = [item[0] for item in self.opt.in_geo] + self.in_nml = [item[0] for item in self.opt.in_nml] + + self.in_geo_dim = sum([item[1] for item in self.opt.in_geo]) + self.in_nml_dim = sum([item[1] for item in self.opt.in_nml]) + + self.in_total = self.in_geo + self.in_nml + self.smpl_feat_dict = None + self.smplx_data = SMPLX() + + if self.prior_type == 'icon': + if 'image' in self.in_geo: + self.channels_filter = [[0, 1, 2, 3, 4, 5], [0, 1, 2, 6, 7, 8]] + else: + self.channels_filter = [[0, 1, 2], [3, 4, 5]] + + else: + if 'image' in self.in_geo: + self.channels_filter = [[0, 1, 2, 3, 4, 5, 6, 7, 8]] + else: + self.channels_filter = [[0, 1, 2, 3, 4, 5]] + + channels_IF[0] = self.hourglass_dim if self.use_filter else len( + self.channels_filter[0]) + + if self.prior_type == 'icon' and 'vis' not in self.smpl_feats: + if self.use_filter: + channels_IF[0] += self.hourglass_dim + else: + channels_IF[0] += len(self.channels_filter[0]) + + if self.prior_type == 'icon': + channels_IF[0] += self.smpl_dim + elif self.prior_type == 'pamir': + channels_IF[0] += self.voxel_dim + smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras = read_smpl_constants( + self.smplx_data.tedra_dir) + self.voxelization = Voxelization( + smpl_vertex_code, + smpl_face_code, + smpl_faces, + smpl_tetras, + volume_res=128, + sigma=0.05, + smooth_kernel_size=7, + batch_size=cfg.batch_size, + device=torch.device(f"cuda:{cfg.gpus[0]}")) + self.ve = VolumeEncoder(3, self.voxel_dim, self.opt.num_stack) + else: + channels_IF[0] += 1 + + self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"] + self.pamir_keys = [ + "voxel_verts", "voxel_faces", "pad_v_num", "pad_f_num" + ] + + self.if_regressor = MLP( + filter_channels=channels_IF, + name='if', + res_layers=self.opt.res_layers, + norm=self.opt.norm_mlp, + last_op=nn.Sigmoid() if not cfg.test_mode else None) + + # network + if self.use_filter: + if self.opt.gtype == "HGPIFuNet": + self.F_filter = HGFilter(self.opt, self.opt.num_stack, + len(self.channels_filter[0])) + else: + print( + colored(f"Backbone {self.opt.gtype} is unimplemented", + 'green')) + + summary_log = f"{self.prior_type.upper()}:\n" + \ + f"w/ Global Image Encoder: {self.use_filter}\n" + \ + f"Image Features used by MLP: {self.in_geo}\n" + + if self.prior_type == "icon": + summary_log += f"Geometry Features used by MLP: {self.smpl_feats}\n" + summary_log += f"Dim of Image Features (local): 6\n" + summary_log += f"Dim of Geometry Features (ICON): {self.smpl_dim}\n" + elif self.prior_type == "pamir": + summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n" + summary_log += f"Dim of Geometry Features (PaMIR): {self.voxel_dim}\n" + else: + summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n" + summary_log += f"Dim of Geometry Features (PIFu): 1 (z-value)\n" + + summary_log += f"Dim of MLP's first layer: {channels_IF[0]}\n" + + print(colored(summary_log, "yellow")) + + self.normal_filter = NormalNet(cfg) + init_net(self) + + def get_normal(self, in_tensor_dict): + + # insert normal features + if (not self.training) and (not self.overfit): + # print(colored("infer normal","blue")) + with torch.no_grad(): + feat_lst = [] + if "image" in self.in_geo: + feat_lst.append( + in_tensor_dict['image']) # [1, 3, 512, 512] + if 'normal_F' in self.in_geo and 'normal_B' in self.in_geo: + if 'normal_F' not in in_tensor_dict.keys( + ) or 'normal_B' not in in_tensor_dict.keys(): + (nmlF, nmlB) = self.normal_filter(in_tensor_dict) + else: + nmlF = in_tensor_dict['normal_F'] + nmlB = in_tensor_dict['normal_B'] + feat_lst.append(nmlF) # [1, 3, 512, 512] + feat_lst.append(nmlB) # [1, 3, 512, 512] + in_filter = torch.cat(feat_lst, dim=1) + + else: + in_filter = torch.cat([in_tensor_dict[key] for key in self.in_geo], + dim=1) + + return in_filter + + def get_mask(self, in_filter, size=128): + + mask = F.interpolate(in_filter[:, self.channels_filter[0]], + size=(size, size), + mode="bilinear", + align_corners=True).abs().sum(dim=1, + keepdim=True) != 0.0 + + return mask + + def filter(self, in_tensor_dict, return_inter=False): + ''' + Filter the input images + store all intermediate features. + :param images: [B, C, H, W] input images + ''' + + in_filter = self.get_normal(in_tensor_dict) + + features_G = [] + + if self.prior_type == 'icon': + if self.use_filter: + features_F = self.F_filter(in_filter[:, + self.channels_filter[0]] + ) # [(B,hg_dim,128,128) * 4] + features_B = self.F_filter(in_filter[:, + self.channels_filter[1]] + ) # [(B,hg_dim,128,128) * 4] + else: + features_F = [in_filter[:, self.channels_filter[0]]] + features_B = [in_filter[:, self.channels_filter[1]]] + for idx in range(len(features_F)): + features_G.append( + torch.cat([features_F[idx], features_B[idx]], dim=1)) + else: + if self.use_filter: + features_G = self.F_filter(in_filter[:, + self.channels_filter[0]]) + else: + features_G = [in_filter[:, self.channels_filter[0]]] + + if self.prior_type == 'icon': + self.smpl_feat_dict = { + k: in_tensor_dict[k] + for k in self.icon_keys + } + elif self.prior_type == "pamir": + self.smpl_feat_dict = { + k: in_tensor_dict[k] + for k in self.pamir_keys + } + else: + pass + # print(colored("use z rather than icon or pamir", "green")) + + # If it is not in training, only produce the last im_feat + if not self.training: + features_out = [features_G[-1]] + else: + features_out = features_G + + if maskout: + features_out_mask = [] + for feat in features_out: + features_out_mask.append( + feat * self.get_mask(in_filter, size=feat.shape[2])) + features_out = features_out_mask + + if return_inter: + return features_out, in_filter + else: + return features_out + + def query(self, features, points, calibs, transforms=None, regressor=None): + + xyz = self.projection(points, calibs, transforms) + + (xy, z) = xyz.split([2, 1], dim=1) + + in_cube = (xyz > -1.0) & (xyz < 1.0) + in_cube = in_cube.all(dim=1, keepdim=True).detach().float() + + preds_list = [] + + if self.prior_type == 'icon': + + # smpl_verts [B, N_vert, 3] + # smpl_faces [B, N_face, 3] + # points [B, 3, N] + + smpl_sdf, smpl_norm, smpl_cmap, smpl_vis = cal_sdf_batch( + self.smpl_feat_dict['smpl_verts'], + self.smpl_feat_dict['smpl_faces'], + self.smpl_feat_dict['smpl_cmap'], + self.smpl_feat_dict['smpl_vis'], + xyz.permute(0, 2, 1).contiguous()) + + # smpl_sdf [B, N, 1] + # smpl_norm [B, N, 3] + # smpl_cmap [B, N, 3] + # smpl_vis [B, N, 1] + + # set ourlier point features as uniform values + smpl_outlier = torch.abs(smpl_sdf).ge(self.sdf_clip) + smpl_sdf[smpl_outlier] = torch.sign(smpl_sdf[smpl_outlier]) + + feat_lst = [smpl_sdf] + if 'cmap' in self.smpl_feats: + smpl_cmap[smpl_outlier.repeat( + 1, 1, 3)] = smpl_sdf[smpl_outlier].repeat(1, 1, 3) + feat_lst.append(smpl_cmap) + if 'norm' in self.smpl_feats: + feat_lst.append(smpl_norm) + if 'vis' in self.smpl_feats: + feat_lst.append(smpl_vis) + + smpl_feat = torch.cat(feat_lst, dim=2).permute(0, 2, 1) + vol_feats = features + + elif self.prior_type == "pamir": + + voxel_verts = self.smpl_feat_dict[ + 'voxel_verts'][:, :-self.smpl_feat_dict['pad_v_num'][0], :] + voxel_faces = self.smpl_feat_dict[ + 'voxel_faces'][:, :-self.smpl_feat_dict['pad_f_num'][0], :] + + self.voxelization.update_param( + batch_size=voxel_faces.shape[0], + smpl_tetra=voxel_faces[0].detach().cpu().numpy()) + vol = self.voxelization(voxel_verts) # vol ~ [0,1] + vol_feats = self.ve(vol, intermediate_output=self.training) + else: + vol_feats = features + + for im_feat, vol_feat in zip(features, vol_feats): + + # [B, Feat_i + z, N] + # normal feature choice by smpl_vis + if self.prior_type == 'icon': + if 'vis' in self.smpl_feats: + point_local_feat = feat_select(self.index(im_feat, xy), + smpl_feat[:, [-1], :]) + if maskout: + normal_mask = torch.tile( + point_local_feat.sum(dim=1, keepdims=True) == 0.0, + (1, smpl_feat.shape[1], 1)) + normal_mask[:, 1:, :] = False + smpl_feat[normal_mask] = -1.0 + point_feat_list = [point_local_feat, smpl_feat[:, :-1, :]] + else: + point_local_feat = self.index(im_feat, xy) + point_feat_list = [point_local_feat, smpl_feat[:, :, :]] + + elif self.prior_type == 'pamir': + # im_feat [B, hg_dim, 128, 128] + # vol_feat [B, vol_dim, 32, 32, 32] + point_feat_list = [ + self.index(im_feat, xy), + self.index(vol_feat, xyz) + ] + + else: + point_feat_list = [self.index(im_feat, xy), z] + + point_feat = torch.cat(point_feat_list, 1) + + # out of image plane is always set to 0 + preds = regressor(point_feat) + preds = in_cube * preds + + preds_list.append(preds) + + return preds_list + + def get_error(self, preds_if_list, labels): + """calcaulate error + + Args: + preds_list (list): list of torch.tensor(B, 3, N) + labels (torch.tensor): (B, N_knn, N) + + Returns: + torch.tensor: error + """ + error_if = 0 + + for pred_id in range(len(preds_if_list)): + pred_if = preds_if_list[pred_id] + error_if += self.error_term(pred_if, labels) + + error_if /= len(preds_if_list) + + return error_if + + def forward(self, in_tensor_dict): + """ + sample_tensor [B, 3, N] + calib_tensor [B, 4, 4] + label_tensor [B, 1, N] + smpl_feat_tensor [B, 59, N] + """ + + sample_tensor = in_tensor_dict['sample'] + calib_tensor = in_tensor_dict['calib'] + label_tensor = in_tensor_dict['label'] + + in_feat = self.filter(in_tensor_dict) + + preds_if_list = self.query(in_feat, + sample_tensor, + calib_tensor, + regressor=self.if_regressor) + + error = self.get_error(preds_if_list, label_tensor) + + return preds_if_list[-1], error diff --git a/lib/net/MLP.py b/lib/net/MLP.py new file mode 100644 index 0000000000000000000000000000000000000000..b9abd1dcf42481bb0434a72e24be27bfeef92772 --- /dev/null +++ b/lib/net/MLP.py @@ -0,0 +1,72 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. + +import torch +import torch.nn as nn +import pytorch_lightning as pl + + +class MLP(pl.LightningModule): + def __init__(self, + filter_channels, + name=None, + res_layers=[], + norm='group', + last_op=None): + + super(MLP, self).__init__() + + self.filters = nn.ModuleList() + self.norms = nn.ModuleList() + self.res_layers = res_layers + self.norm = norm + self.last_op = last_op + self.name = name + self.activate = nn.LeakyReLU(inplace=True) + + for l in range(0, len(filter_channels) - 1): + if l in self.res_layers: + self.filters.append( + nn.Conv1d(filter_channels[l] + filter_channels[0], + filter_channels[l + 1], 1)) + else: + self.filters.append( + nn.Conv1d(filter_channels[l], filter_channels[l + 1], 1)) + + if l != len(filter_channels) - 2: + if norm == 'group': + self.norms.append(nn.GroupNorm(32, filter_channels[l + 1])) + elif norm == 'batch': + self.norms.append(nn.BatchNorm1d(filter_channels[l + 1])) + elif norm == 'instance': + self.norms.append(nn.InstanceNorm1d(filter_channels[l + + 1])) + elif norm == 'weight': + self.filters[l] = nn.utils.weight_norm(self.filters[l], + name='weight') + # print(self.filters[l].weight_g.size(), + # self.filters[l].weight_v.size()) + + def forward(self, feature): + ''' + feature may include multiple view inputs + args: + feature: [B, C_in, N] + return: + [B, C_out, N] prediction + ''' + y = feature + tmpy = feature + + for i, f in enumerate(self.filters): + + y = f(y if i not in self.res_layers else torch.cat([y, tmpy], 1)) + if i != len(self.filters) - 1: + if self.norm not in ['batch', 'group', 'instance']: + y = self.activate(y) + else: + y = self.activate(self.norms[i](y)) + + if self.last_op is not None: + y = self.last_op(y) + + return y diff --git a/lib/net/NormalNet.py b/lib/net/NormalNet.py new file mode 100644 index 0000000000000000000000000000000000000000..96b372d8a4c934006201135d7b8f0a88da263aba --- /dev/null +++ b/lib/net/NormalNet.py @@ -0,0 +1,122 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.net.FBNet import define_G +from lib.net.net_util import init_net, VGGLoss +from lib.net.HGFilters import * +from lib.net.BasePIFuNet import BasePIFuNet +import torch +import torch.nn as nn + + +class NormalNet(BasePIFuNet): + ''' + HG PIFu network uses Hourglass stacks as the image filter. + It does the following: + 1. Compute image feature stacks and store it in self.im_feat_list + self.im_feat_list[-1] is the last stack (output stack) + 2. Calculate calibration + 3. If training, it index on every intermediate stacks, + If testing, it index on the last stack. + 4. Classification. + 5. During training, error is calculated on all stacks. + ''' + + def __init__(self, cfg, error_term=nn.SmoothL1Loss()): + + super(NormalNet, self).__init__(error_term=error_term) + + self.l1_loss = nn.SmoothL1Loss() + + self.opt = cfg.net + + if self.training: + self.vgg_loss = [VGGLoss()] + + self.in_nmlF = [ + item[0] for item in self.opt.in_nml + if '_F' in item[0] or item[0] == 'image' + ] + self.in_nmlB = [ + item[0] for item in self.opt.in_nml + if '_B' in item[0] or item[0] == 'image' + ] + self.in_nmlF_dim = sum([ + item[1] for item in self.opt.in_nml + if '_F' in item[0] or item[0] == 'image' + ]) + self.in_nmlB_dim = sum([ + item[1] for item in self.opt.in_nml + if '_B' in item[0] or item[0] == 'image' + ]) + + self.netF = define_G(self.in_nmlF_dim, 3, 64, "global", 4, 9, 1, 3, + "instance") + self.netB = define_G(self.in_nmlB_dim, 3, 64, "global", 4, 9, 1, 3, + "instance") + + init_net(self) + + def forward(self, in_tensor): + + inF_list = [] + inB_list = [] + + for name in self.in_nmlF: + inF_list.append(in_tensor[name]) + for name in self.in_nmlB: + inB_list.append(in_tensor[name]) + + nmlF = self.netF(torch.cat(inF_list, dim=1)) + nmlB = self.netB(torch.cat(inB_list, dim=1)) + + # ||normal|| == 1 + nmlF /= torch.norm(nmlF, dim=1) + nmlB /= torch.norm(nmlB, dim=1) + + # output: float_arr [-1,1] with [B, C, H, W] + + mask = (in_tensor['image'].abs().sum(dim=1, keepdim=True) != + 0.0).detach().float() + + nmlF = nmlF * mask + nmlB = nmlB * mask + + return nmlF, nmlB + + def get_norm_error(self, prd_F, prd_B, tgt): + """calculate normal loss + + Args: + pred (torch.tensor): [B, 6, 512, 512] + tagt (torch.tensor): [B, 6, 512, 512] + """ + + tgt_F, tgt_B = tgt['normal_F'], tgt['normal_B'] + + l1_F_loss = self.l1_loss(prd_F, tgt_F) + l1_B_loss = self.l1_loss(prd_B, tgt_B) + + with torch.no_grad(): + vgg_F_loss = self.vgg_loss[0](prd_F, tgt_F) + vgg_B_loss = self.vgg_loss[0](prd_B, tgt_B) + + total_loss = [ + 5.0 * l1_F_loss + vgg_F_loss, 5.0 * l1_B_loss + vgg_B_loss + ] + + return total_loss diff --git a/lib/net/VE.py b/lib/net/VE.py new file mode 100644 index 0000000000000000000000000000000000000000..65c1dfab953ea2d5135a04050b2fb020aa03bf7d --- /dev/null +++ b/lib/net/VE.py @@ -0,0 +1,183 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + + +import torch.nn as nn +import pytorch_lightning as pl + + +class BaseNetwork(pl.LightningModule): + def __init__(self): + super(BaseNetwork, self).__init__() + + def init_weights(self, init_type='xavier', gain=0.02): + ''' + initializes network's weights + init_type: normal | xavier | kaiming | orthogonal + https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39 + ''' + def init_func(m): + classname = m.__class__.__name__ + if hasattr(m, 'weight') and (classname.find('Conv') != -1 + or classname.find('Linear') != -1): + if init_type == 'normal': + nn.init.normal_(m.weight.data, 0.0, gain) + elif init_type == 'xavier': + nn.init.xavier_normal_(m.weight.data, gain=gain) + elif init_type == 'kaiming': + nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif init_type == 'orthogonal': + nn.init.orthogonal_(m.weight.data, gain=gain) + + if hasattr(m, 'bias') and m.bias is not None: + nn.init.constant_(m.bias.data, 0.0) + + elif classname.find('BatchNorm2d') != -1: + nn.init.normal_(m.weight.data, 1.0, gain) + nn.init.constant_(m.bias.data, 0.0) + + self.apply(init_func) + + +class Residual3D(BaseNetwork): + def __init__(self, numIn, numOut): + super(Residual3D, self).__init__() + self.numIn = numIn + self.numOut = numOut + self.with_bias = True + # self.bn = nn.GroupNorm(4, self.numIn) + self.bn = nn.BatchNorm3d(self.numIn) + self.relu = nn.ReLU(inplace=True) + self.conv1 = nn.Conv3d(self.numIn, + self.numOut, + bias=self.with_bias, + kernel_size=3, + stride=1, + padding=2, + dilation=2) + # self.bn1 = nn.GroupNorm(4, self.numOut) + self.bn1 = nn.BatchNorm3d(self.numOut) + self.conv2 = nn.Conv3d(self.numOut, + self.numOut, + bias=self.with_bias, + kernel_size=3, + stride=1, + padding=1) + # self.bn2 = nn.GroupNorm(4, self.numOut) + self.bn2 = nn.BatchNorm3d(self.numOut) + self.conv3 = nn.Conv3d(self.numOut, + self.numOut, + bias=self.with_bias, + kernel_size=3, + stride=1, + padding=1) + + if self.numIn != self.numOut: + self.conv4 = nn.Conv3d(self.numIn, + self.numOut, + bias=self.with_bias, + kernel_size=1) + self.init_weights() + + def forward(self, x): + residual = x + # out = self.bn(x) + # out = self.relu(out) + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.bn2(out) + # out = self.conv3(out) + # out = self.relu(out) + + if self.numIn != self.numOut: + residual = self.conv4(x) + + return out + residual + + +class VolumeEncoder(BaseNetwork): + """CycleGan Encoder""" + + def __init__(self, num_in=3, num_out=32, num_stacks=2): + super(VolumeEncoder, self).__init__() + self.num_in = num_in + self.num_out = num_out + self.num_inter = 8 + self.num_stacks = num_stacks + self.with_bias = True + + self.relu = nn.ReLU(inplace=True) + self.conv1 = nn.Conv3d(self.num_in, + self.num_inter, + bias=self.with_bias, + kernel_size=5, + stride=2, + padding=4, + dilation=2) + # self.bn1 = nn.GroupNorm(4, self.num_inter) + self.bn1 = nn.BatchNorm3d(self.num_inter) + self.conv2 = nn.Conv3d(self.num_inter, + self.num_out, + bias=self.with_bias, + kernel_size=5, + stride=2, + padding=4, + dilation=2) + # self.bn2 = nn.GroupNorm(4, self.num_out) + self.bn2 = nn.BatchNorm3d(self.num_out) + + self.conv_out1 = nn.Conv3d(self.num_out, + self.num_out, + bias=self.with_bias, + kernel_size=3, + stride=1, + padding=1, + dilation=1) + self.conv_out2 = nn.Conv3d(self.num_out, + self.num_out, + bias=self.with_bias, + kernel_size=3, + stride=1, + padding=1, + dilation=1) + + for idx in range(self.num_stacks): + self.add_module("res" + str(idx), + Residual3D(self.num_out, self.num_out)) + + self.init_weights() + + def forward(self, x, intermediate_output=True): + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out_lst = [] + for idx in range(self.num_stacks): + out = self._modules["res" + str(idx)](out) + out_lst.append(out) + + if intermediate_output: + return out_lst + else: + return [out_lst[-1]] diff --git a/lib/net/__init__.py b/lib/net/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..3373108ec282a13b43ba61181bc6ff4e090ebfda --- /dev/null +++ b/lib/net/__init__.py @@ -0,0 +1,4 @@ +from .BasePIFuNet import BasePIFuNet +from .HGPIFuNet import HGPIFuNet +from .NormalNet import NormalNet +from .VE import VolumeEncoder diff --git a/lib/net/geometry.py b/lib/net/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..a3e6464abf8298d7d7b00c3f5a6bcf74a30fa991 --- /dev/null +++ b/lib/net/geometry.py @@ -0,0 +1,82 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import torch + + +def index(feat, uv): + ''' + :param feat: [B, C, H, W] image features + :param uv: [B, 2, N] uv coordinates in the image plane, range [0, 1] + :return: [B, C, N] image features at the uv coordinates + ''' + uv = uv.transpose(1, 2) # [B, N, 2] + + (B, N, _) = uv.shape + C = feat.shape[1] + + if uv.shape[-1] == 3: + # uv = uv[:,:,[2,1,0]] + # uv = uv * torch.tensor([1.0,-1.0,1.0]).type_as(uv)[None,None,...] + uv = uv.unsqueeze(2).unsqueeze(3) # [B, N, 1, 1, 3] + else: + uv = uv.unsqueeze(2) # [B, N, 1, 2] + + # NOTE: for newer PyTorch, it seems that training results are degraded due to implementation diff in F.grid_sample + # for old versions, simply remove the aligned_corners argument. + samples = torch.nn.functional.grid_sample( + feat, uv, align_corners=True) # [B, C, N, 1] + return samples.view(B, C, N) # [B, C, N] + + +def orthogonal(points, calibrations, transforms=None): + ''' + Compute the orthogonal projections of 3D points into the image plane by given projection matrix + :param points: [B, 3, N] Tensor of 3D points + :param calibrations: [B, 3, 4] Tensor of projection matrix + :param transforms: [B, 2, 3] Tensor of image transform matrix + :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane + ''' + rot = calibrations[:, :3, :3] + trans = calibrations[:, :3, 3:4] + pts = torch.baddbmm(trans, rot, points) # [B, 3, N] + if transforms is not None: + scale = transforms[:2, :2] + shift = transforms[:2, 2:3] + pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :]) + return pts + + +def perspective(points, calibrations, transforms=None): + ''' + Compute the perspective projections of 3D points into the image plane by given projection matrix + :param points: [Bx3xN] Tensor of 3D points + :param calibrations: [Bx3x4] Tensor of projection matrix + :param transforms: [Bx2x3] Tensor of image transform matrix + :return: xy: [Bx2xN] Tensor of xy coordinates in the image plane + ''' + rot = calibrations[:, :3, :3] + trans = calibrations[:, :3, 3:4] + homo = torch.baddbmm(trans, rot, points) # [B, 3, N] + xy = homo[:, :2, :] / homo[:, 2:3, :] + if transforms is not None: + scale = transforms[:2, :2] + shift = transforms[:2, 2:3] + xy = torch.baddbmm(shift, scale, xy) + + xyz = torch.cat([xy, homo[:, 2:3, :]], 1) + return xyz diff --git a/lib/net/local_affine.py b/lib/net/local_affine.py new file mode 100644 index 0000000000000000000000000000000000000000..0992d4f04a5739c9c43a9a71db26c5bd2441ffcd --- /dev/null +++ b/lib/net/local_affine.py @@ -0,0 +1,57 @@ +# Copyright 2021 by Haozhe Wu, Tsinghua University, Department of Computer Science and Technology. +# All rights reserved. +# This file is part of the pytorch-nicp, +# and is released under the "MIT License Agreement". Please see the LICENSE +# file that should have been included as part of this package. + +import torch +import torch.nn as nn +import torch.sparse as sp + +# reference: https://github.com/wuhaozhe/pytorch-nicp +class LocalAffine(nn.Module): + def __init__(self, num_points, batch_size=1, edges=None): + ''' + specify the number of points, the number of points should be constant across the batch + and the edges torch.Longtensor() with shape N * 2 + the local affine operator supports batch operation + batch size must be constant + add additional pooling on top of w matrix + ''' + super(LocalAffine, self).__init__() + self.A = nn.Parameter(torch.eye(3).unsqueeze( + 0).unsqueeze(0).repeat(batch_size, num_points, 1, 1)) + self.b = nn.Parameter(torch.zeros(3).unsqueeze(0).unsqueeze( + 0).unsqueeze(3).repeat(batch_size, num_points, 1, 1)) + self.edges = edges + self.num_points = num_points + + def stiffness(self): + ''' + calculate the stiffness of local affine transformation + f norm get infinity gradient when w is zero matrix, + ''' + if self.edges is None: + raise Exception("edges cannot be none when calculate stiff") + idx1 = self.edges[:, 0] + idx2 = self.edges[:, 1] + affine_weight = torch.cat((self.A, self.b), dim=3) + w1 = torch.index_select(affine_weight, dim=1, index=idx1) + w2 = torch.index_select(affine_weight, dim=1, index=idx2) + w_diff = (w1 - w2) ** 2 + w_rigid = (torch.linalg.det(self.A) - 1.0) ** 2 + return w_diff, w_rigid + + def forward(self, x, return_stiff=False): + ''' + x should have shape of B * N * 3 + ''' + x = x.unsqueeze(3) + out_x = torch.matmul(self.A, x) + out_x = out_x + self.b + out_x.squeeze_(3) + if return_stiff: + stiffness, rigid = self.stiffness() + return out_x, stiffness, rigid + else: + return out_x diff --git a/lib/net/net_util.py b/lib/net/net_util.py new file mode 100644 index 0000000000000000000000000000000000000000..5eea98f5c4ce0960aa123b7c529aeac31917d45f --- /dev/null +++ b/lib/net/net_util.py @@ -0,0 +1,329 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from torchvision import models +import torch +from torch.nn import init +import torch.nn as nn +import torch.nn.functional as F +import functools +from torch.autograd import grad + + +def gradient(inputs, outputs): + d_points = torch.ones_like(outputs, + requires_grad=False, + device=outputs.device) + points_grad = grad(outputs=outputs, + inputs=inputs, + grad_outputs=d_points, + create_graph=True, + retain_graph=True, + only_inputs=True, + allow_unused=True)[0] + return points_grad + + +# def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): +# "3x3 convolution with padding" +# return nn.Conv2d(in_planes, out_planes, kernel_size=3, +# stride=strd, padding=padding, bias=bias) + + +def conv3x3(in_planes, + out_planes, + kernel=3, + strd=1, + dilation=1, + padding=1, + bias=False): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=kernel, + dilation=dilation, + stride=strd, + padding=padding, + bias=bias) + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=1, + stride=stride, + bias=False) + + +def init_weights(net, init_type='normal', init_gain=0.02): + """Initialize network weights. + + Parameters: + net (network) -- network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + + We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might + work better for some applications. Feel free to try yourself. + """ + def init_func(m): # define the initialization function + classname = m.__class__.__name__ + if hasattr(m, 'weight') and (classname.find('Conv') != -1 + or classname.find('Linear') != -1): + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, init_gain) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=init_gain) + elif init_type == 'kaiming': + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif init_type == 'orthogonal': + init.orthogonal_(m.weight.data, gain=init_gain) + else: + raise NotImplementedError( + 'initialization method [%s] is not implemented' % + init_type) + if hasattr(m, 'bias') and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif classname.find( + 'BatchNorm2d' + ) != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. + init.normal_(m.weight.data, 1.0, init_gain) + init.constant_(m.bias.data, 0.0) + + # print('initialize network with %s' % init_type) + net.apply(init_func) # apply the initialization function + + +def init_net(net, init_type='xavier', init_gain=0.02, gpu_ids=[]): + """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights + Parameters: + net (network) -- the network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Return an initialized network. + """ + if len(gpu_ids) > 0: + assert (torch.cuda.is_available()) + net = torch.nn.DataParallel(net) # multi-GPUs + init_weights(net, init_type, init_gain=init_gain) + return net + + +def imageSpaceRotation(xy, rot): + ''' + args: + xy: (B, 2, N) input + rot: (B, 2) x,y axis rotation angles + + rotation center will be always image center (other rotation center can be represented by additional z translation) + ''' + disp = rot.unsqueeze(2).sin().expand_as(xy) + return (disp * xy).sum(dim=1) + + +def cal_gradient_penalty(netD, + real_data, + fake_data, + device, + type='mixed', + constant=1.0, + lambda_gp=10.0): + """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 + + Arguments: + netD (network) -- discriminator network + real_data (tensor array) -- real images + fake_data (tensor array) -- generated images from the generator + device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') + type (str) -- if we mix real and fake data or not [real | fake | mixed]. + constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2 + lambda_gp (float) -- weight for this loss + + Returns the gradient penalty loss + """ + if lambda_gp > 0.0: + # either use real images, fake images, or a linear interpolation of two. + if type == 'real': + interpolatesv = real_data + elif type == 'fake': + interpolatesv = fake_data + elif type == 'mixed': + alpha = torch.rand(real_data.shape[0], 1) + alpha = alpha.expand( + real_data.shape[0], + real_data.nelement() // + real_data.shape[0]).contiguous().view(*real_data.shape) + alpha = alpha.to(device) + interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) + else: + raise NotImplementedError('{} not implemented'.format(type)) + interpolatesv.requires_grad_(True) + disc_interpolates = netD(interpolatesv) + gradients = torch.autograd.grad( + outputs=disc_interpolates, + inputs=interpolatesv, + grad_outputs=torch.ones(disc_interpolates.size()).to(device), + create_graph=True, + retain_graph=True, + only_inputs=True) + gradients = gradients[0].view(real_data.size(0), -1) # flat the data + gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** + 2).mean() * lambda_gp # added eps + return gradient_penalty, gradients + else: + return 0.0, None + + +def get_norm_layer(norm_type='instance'): + """Return a normalization layer + Parameters: + norm_type (str) -- the name of the normalization layer: batch | instance | none + For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). + For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. + """ + if norm_type == 'batch': + norm_layer = functools.partial(nn.BatchNorm2d, + affine=True, + track_running_stats=True) + elif norm_type == 'instance': + norm_layer = functools.partial(nn.InstanceNorm2d, + affine=False, + track_running_stats=False) + elif norm_type == 'group': + norm_layer = functools.partial(nn.GroupNorm, 32) + elif norm_type == 'none': + norm_layer = None + else: + raise NotImplementedError('normalization layer [%s] is not found' % + norm_type) + return norm_layer + + +class Flatten(nn.Module): + def forward(self, input): + return input.view(input.size(0), -1) + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, out_planes, opt): + super(ConvBlock, self).__init__() + [k, s, d, p] = opt.conv3x3 + self.conv1 = conv3x3(in_planes, int(out_planes / 2), k, s, d, p) + self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4), k, s, d, + p) + self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4), k, s, d, + p) + + if opt.norm == 'batch': + self.bn1 = nn.BatchNorm2d(in_planes) + self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) + self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) + self.bn4 = nn.BatchNorm2d(in_planes) + elif opt.norm == 'group': + self.bn1 = nn.GroupNorm(32, in_planes) + self.bn2 = nn.GroupNorm(32, int(out_planes / 2)) + self.bn3 = nn.GroupNorm(32, int(out_planes / 4)) + self.bn4 = nn.GroupNorm(32, in_planes) + + if in_planes != out_planes: + self.downsample = nn.Sequential( + self.bn4, + nn.ReLU(True), + nn.Conv2d(in_planes, + out_planes, + kernel_size=1, + stride=1, + bias=False), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out1 = self.bn1(x) + out1 = F.relu(out1, True) + out1 = self.conv1(out1) + + out2 = self.bn2(out1) + out2 = F.relu(out2, True) + out2 = self.conv2(out2) + + out3 = self.bn3(out2) + out3 = F.relu(out3, True) + out3 = self.conv3(out3) + + out3 = torch.cat((out1, out2, out3), 1) + + if self.downsample is not None: + residual = self.downsample(residual) + + out3 += residual + + return out3 + + +class Vgg19(torch.nn.Module): + def __init__(self, requires_grad=False): + super(Vgg19, self).__init__() + vgg_pretrained_features = models.vgg19(pretrained=True).features + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + self.slice5 = torch.nn.Sequential() + for x in range(2): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(2, 7): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice5.add_module(str(x), vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + h_relu1 = self.slice1(X) + h_relu2 = self.slice2(h_relu1) + h_relu3 = self.slice3(h_relu2) + h_relu4 = self.slice4(h_relu3) + h_relu5 = self.slice5(h_relu4) + out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] + return out + + +class VGGLoss(nn.Module): + def __init__(self): + super(VGGLoss, self).__init__() + self.vgg = Vgg19().cuda() + self.criterion = nn.L1Loss() + self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0] + + def forward(self, x, y): + x_vgg, y_vgg = self.vgg(x), self.vgg(y) + loss = 0 + for i in range(len(x_vgg)): + loss += self.weights[i] * self.criterion(x_vgg[i], + y_vgg[i].detach()) + return loss diff --git a/lib/net/voxelize.py b/lib/net/voxelize.py new file mode 100644 index 0000000000000000000000000000000000000000..2354b755e8aef61dfb144b7217cf24afc68ab2c1 --- /dev/null +++ b/lib/net/voxelize.py @@ -0,0 +1,184 @@ +from __future__ import division, print_function +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from torch.autograd import Function + +import voxelize_cuda + + +class VoxelizationFunction(Function): + """ + Definition of differentiable voxelization function + Currently implemented only for cuda Tensors + """ + @staticmethod + def forward(ctx, smpl_vertices, smpl_face_center, smpl_face_normal, + smpl_vertex_code, smpl_face_code, smpl_tetrahedrons, + volume_res, sigma, smooth_kernel_size): + """ + forward pass + Output format: (batch_size, z_dims, y_dims, x_dims, channel_num) + """ + assert (smpl_vertices.size()[1] == smpl_vertex_code.size()[1]) + assert (smpl_face_center.size()[1] == smpl_face_normal.size()[1]) + assert (smpl_face_center.size()[1] == smpl_face_code.size()[1]) + ctx.batch_size = smpl_vertices.size()[0] + ctx.volume_res = volume_res + ctx.sigma = sigma + ctx.smooth_kernel_size = smooth_kernel_size + ctx.smpl_vertex_num = smpl_vertices.size()[1] + ctx.device = smpl_vertices.device + + smpl_vertices = smpl_vertices.contiguous() + smpl_face_center = smpl_face_center.contiguous() + smpl_face_normal = smpl_face_normal.contiguous() + smpl_vertex_code = smpl_vertex_code.contiguous() + smpl_face_code = smpl_face_code.contiguous() + smpl_tetrahedrons = smpl_tetrahedrons.contiguous() + + occ_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res, + ctx.volume_res, + ctx.volume_res).fill_(0.0) + semantic_volume = torch.cuda.FloatTensor(ctx.batch_size, + ctx.volume_res, + ctx.volume_res, + ctx.volume_res, 3).fill_(0.0) + weight_sum_volume = torch.cuda.FloatTensor(ctx.batch_size, + ctx.volume_res, + ctx.volume_res, + ctx.volume_res).fill_(1e-3) + + # occ_volume [B, volume_res, volume_res, volume_res] + # semantic_volume [B, volume_res, volume_res, volume_res, 3] + # weight_sum_volume [B, volume_res, volume_res, volume_res] + + occ_volume, semantic_volume, weight_sum_volume = voxelize_cuda.forward_semantic_voxelization( + smpl_vertices, smpl_vertex_code, smpl_tetrahedrons, occ_volume, + semantic_volume, weight_sum_volume, sigma) + + return semantic_volume + + +class Voxelization(nn.Module): + """ + Wrapper around the autograd function VoxelizationFunction + """ + + def __init__(self, smpl_vertex_code, smpl_face_code, smpl_face_indices, + smpl_tetraderon_indices, volume_res, sigma, + smooth_kernel_size, batch_size, device): + super(Voxelization, self).__init__() + assert (len(smpl_face_indices.shape) == 2) + assert (len(smpl_tetraderon_indices.shape) == 2) + assert (smpl_face_indices.shape[1] == 3) + assert (smpl_tetraderon_indices.shape[1] == 4) + + self.volume_res = volume_res + self.sigma = sigma + self.smooth_kernel_size = smooth_kernel_size + self.batch_size = batch_size + self.device = device + + self.smpl_vertex_code = smpl_vertex_code + self.smpl_face_code = smpl_face_code + self.smpl_face_indices = smpl_face_indices + self.smpl_tetraderon_indices = smpl_tetraderon_indices + + def update_param(self, batch_size, smpl_tetra): + + self.batch_size = batch_size + self.smpl_tetraderon_indices = smpl_tetra + + smpl_vertex_code_batch = np.tile(self.smpl_vertex_code, + (self.batch_size, 1, 1)) + smpl_face_code_batch = np.tile(self.smpl_face_code, + (self.batch_size, 1, 1)) + smpl_face_indices_batch = np.tile(self.smpl_face_indices, + (self.batch_size, 1, 1)) + smpl_tetraderon_indices_batch = np.tile(self.smpl_tetraderon_indices, + (self.batch_size, 1, 1)) + + smpl_vertex_code_batch = torch.from_numpy( + smpl_vertex_code_batch).contiguous().to(self.device) + smpl_face_code_batch = torch.from_numpy( + smpl_face_code_batch).contiguous().to(self.device) + smpl_face_indices_batch = torch.from_numpy( + smpl_face_indices_batch).contiguous().to(self.device) + smpl_tetraderon_indices_batch = torch.from_numpy( + smpl_tetraderon_indices_batch).contiguous().to(self.device) + + self.register_buffer('smpl_vertex_code_batch', smpl_vertex_code_batch) + self.register_buffer('smpl_face_code_batch', smpl_face_code_batch) + self.register_buffer('smpl_face_indices_batch', + smpl_face_indices_batch) + self.register_buffer('smpl_tetraderon_indices_batch', + smpl_tetraderon_indices_batch) + + def forward(self, smpl_vertices): + """ + Generate semantic volumes from SMPL vertices + """ + assert (smpl_vertices.size()[0] == self.batch_size) + self.check_input(smpl_vertices) + smpl_faces = self.vertices_to_faces(smpl_vertices) + smpl_tetrahedrons = self.vertices_to_tetrahedrons(smpl_vertices) + smpl_face_center = self.calc_face_centers(smpl_faces) + smpl_face_normal = self.calc_face_normals(smpl_faces) + smpl_surface_vertex_num = self.smpl_vertex_code_batch.size()[1] + smpl_vertices_surface = smpl_vertices[:, :smpl_surface_vertex_num, :] + vol = VoxelizationFunction.apply(smpl_vertices_surface, + smpl_face_center, smpl_face_normal, + self.smpl_vertex_code_batch, + self.smpl_face_code_batch, + smpl_tetrahedrons, self.volume_res, + self.sigma, self.smooth_kernel_size) + return vol.permute((0, 4, 1, 2, 3)) # (bzyxc --> bcdhw) + + def vertices_to_faces(self, vertices): + assert (vertices.ndimension() == 3) + bs, nv = vertices.shape[:2] + device = vertices.device + face = self.smpl_face_indices_batch + ( + torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] + vertices_ = vertices.reshape((bs * nv, 3)) + return vertices_[face.long()] + + def vertices_to_tetrahedrons(self, vertices): + assert (vertices.ndimension() == 3) + bs, nv = vertices.shape[:2] + device = vertices.device + tets = self.smpl_tetraderon_indices_batch + ( + torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] + vertices_ = vertices.reshape((bs * nv, 3)) + return vertices_[tets.long()] + + def calc_face_centers(self, face_verts): + assert len(face_verts.shape) == 4 + assert face_verts.shape[2] == 3 + assert face_verts.shape[3] == 3 + bs, nf = face_verts.shape[:2] + face_centers = (face_verts[:, :, 0, :] + face_verts[:, :, 1, :] + + face_verts[:, :, 2, :]) / 3.0 + face_centers = face_centers.reshape((bs, nf, 3)) + return face_centers + + def calc_face_normals(self, face_verts): + assert len(face_verts.shape) == 4 + assert face_verts.shape[2] == 3 + assert face_verts.shape[3] == 3 + bs, nf = face_verts.shape[:2] + face_verts = face_verts.reshape((bs * nf, 3, 3)) + v10 = face_verts[:, 0] - face_verts[:, 1] + v12 = face_verts[:, 2] - face_verts[:, 1] + normals = F.normalize(torch.cross(v10, v12), eps=1e-5) + normals = normals.reshape((bs, nf, 3)) + return normals + + def check_input(self, x): + if x.device == 'cpu': + raise TypeError('Voxelization module supports only cuda tensors') + if x.type() != 'torch.cuda.FloatTensor': + raise TypeError( + 'Voxelization module supports only float32 tensors') diff --git a/lib/pymaf/configs/pymaf_config.yaml b/lib/pymaf/configs/pymaf_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5f9b57256adbc29f5611b88db31472c2857e3955 --- /dev/null +++ b/lib/pymaf/configs/pymaf_config.yaml @@ -0,0 +1,47 @@ +SOLVER: + MAX_ITER: 500000 + TYPE: Adam + BASE_LR: 0.00005 + GAMMA: 0.1 + STEPS: [0] + EPOCHS: [0] +DEBUG: False +LOGDIR: '' +DEVICE: cuda +NUM_WORKERS: 8 +SEED_VALUE: -1 +LOSS: + KP_2D_W: 300.0 + KP_3D_W: 300.0 + SHAPE_W: 0.06 + POSE_W: 60.0 + VERT_W: 0.0 + INDEX_WEIGHTS: 2.0 + # Loss weights for surface parts. (24 Parts) + PART_WEIGHTS: 0.3 + # Loss weights for UV regression. + POINT_REGRESSION_WEIGHTS: 0.5 +TRAIN: + NUM_WORKERS: 8 + BATCH_SIZE: 64 + PIN_MEMORY: True +TEST: + BATCH_SIZE: 32 +MODEL: + PyMAF: + BACKBONE: 'res50' + MLP_DIM: [256, 128, 64, 5] + N_ITER: 3 + AUX_SUPV_ON: True + DP_HEATMAP_SIZE: 56 +RES_MODEL: + DECONV_WITH_BIAS: False + NUM_DECONV_LAYERS: 3 + NUM_DECONV_FILTERS: + - 256 + - 256 + - 256 + NUM_DECONV_KERNELS: + - 4 + - 4 + - 4 diff --git a/lib/pymaf/core/__init__.py b/lib/pymaf/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/pymaf/core/base_trainer.py b/lib/pymaf/core/base_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..764491bd99f89445decaf2217837a0f084ed0271 --- /dev/null +++ b/lib/pymaf/core/base_trainer.py @@ -0,0 +1,107 @@ +# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/utils/base_trainer.py +from __future__ import division +import logging +from utils import CheckpointSaver +from tensorboardX import SummaryWriter + +import torch +from tqdm import tqdm + +tqdm.monitor_interval = 0 + + +logger = logging.getLogger(__name__) + + +class BaseTrainer(object): + """Base class for Trainer objects. + Takes care of checkpointing/logging/resuming training. + """ + + def __init__(self, options): + self.options = options + if options.multiprocessing_distributed: + self.device = torch.device('cuda', options.gpu) + else: + self.device = torch.device( + 'cuda' if torch.cuda.is_available() else 'cpu') + # override this function to define your model, optimizers etc. + self.saver = CheckpointSaver(save_dir=options.checkpoint_dir, + overwrite=options.overwrite) + if options.rank == 0: + self.summary_writer = SummaryWriter(self.options.summary_dir) + self.init_fn() + + self.checkpoint = None + if options.resume and self.saver.exists_checkpoint(): + self.checkpoint = self.saver.load_checkpoint( + self.models_dict, self.optimizers_dict) + + if self.checkpoint is None: + self.epoch_count = 0 + self.step_count = 0 + else: + self.epoch_count = self.checkpoint['epoch'] + self.step_count = self.checkpoint['total_step_count'] + + if self.checkpoint is not None: + self.checkpoint_batch_idx = self.checkpoint['batch_idx'] + else: + self.checkpoint_batch_idx = 0 + + self.best_performance = float('inf') + + def load_pretrained(self, checkpoint_file=None): + """Load a pretrained checkpoint. + This is different from resuming training using --resume. + """ + if checkpoint_file is not None: + checkpoint = torch.load(checkpoint_file) + for model in self.models_dict: + if model in checkpoint: + self.models_dict[model].load_state_dict(checkpoint[model], + strict=True) + print(f'Checkpoint {model} loaded') + + def move_dict_to_device(self, dict, device, tensor2float=False): + for k, v in dict.items(): + if isinstance(v, torch.Tensor): + if tensor2float: + dict[k] = v.float().to(device) + else: + dict[k] = v.to(device) + + # The following methods (with the possible exception of test) have to be implemented in the derived classes + def train(self, epoch): + raise NotImplementedError('You need to provide an train method') + + def init_fn(self): + raise NotImplementedError('You need to provide an _init_fn method') + + def train_step(self, input_batch): + raise NotImplementedError('You need to provide a _train_step method') + + def train_summaries(self, input_batch): + raise NotImplementedError( + 'You need to provide a _train_summaries method') + + def visualize(self, input_batch): + raise NotImplementedError('You need to provide a visualize method') + + def validate(self): + pass + + def test(self): + pass + + def evaluate(self): + pass + + def fit(self): + # Run training for num_epochs epochs + for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), + total=self.options.num_epochs, + initial=self.epoch_count): + self.epoch_count = epoch + self.train(epoch) + return diff --git a/lib/pymaf/core/cfgs.py b/lib/pymaf/core/cfgs.py new file mode 100644 index 0000000000000000000000000000000000000000..09ac4fa48483aa9e595b7e4b27dfa7426cb11d33 --- /dev/null +++ b/lib/pymaf/core/cfgs.py @@ -0,0 +1,100 @@ +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import os +import json +from yacs.config import CfgNode as CN + +# Configuration variables +cfg = CN(new_allowed=True) + +cfg.OUTPUT_DIR = 'results' +cfg.DEVICE = 'cuda' +cfg.DEBUG = False +cfg.LOGDIR = '' +cfg.VAL_VIS_BATCH_FREQ = 200 +cfg.TRAIN_VIS_ITER_FERQ = 1000 +cfg.SEED_VALUE = -1 + +cfg.TRAIN = CN(new_allowed=True) + +cfg.LOSS = CN(new_allowed=True) +cfg.LOSS.KP_2D_W = 300.0 +cfg.LOSS.KP_3D_W = 300.0 +cfg.LOSS.SHAPE_W = 0.06 +cfg.LOSS.POSE_W = 60.0 +cfg.LOSS.VERT_W = 0.0 + +# Loss weights for dense correspondences +cfg.LOSS.INDEX_WEIGHTS = 2.0 +# Loss weights for surface parts. (24 Parts) +cfg.LOSS.PART_WEIGHTS = 0.3 +# Loss weights for UV regression. +cfg.LOSS.POINT_REGRESSION_WEIGHTS = 0.5 + +cfg.MODEL = CN(new_allowed=True) + +cfg.MODEL.PyMAF = CN(new_allowed=True) + +# switch +cfg.TRAIN.VAL_LOOP = True + +cfg.TEST = CN(new_allowed=True) + + +def get_cfg_defaults(): + """Get a yacs CfgNode object with default values for my_project.""" + # Return a clone so that the defaults will not be altered + # This is for the "local variable" use pattern + # return cfg.clone() + return cfg + + +def update_cfg(cfg_file): + # cfg = get_cfg_defaults() + cfg.merge_from_file(cfg_file) + # return cfg.clone() + return cfg + + +def parse_args(args): + cfg_file = args.cfg_file + if args.cfg_file is not None: + cfg = update_cfg(args.cfg_file) + else: + cfg = get_cfg_defaults() + + # if args.misc is not None: + # cfg.merge_from_list(args.misc) + + return cfg + + +def parse_args_extend(args): + if args.resume: + if not os.path.exists(args.log_dir): + raise ValueError( + 'Experiment are set to resume mode, but log directory does not exist.' + ) + + # load log's cfg + cfg_file = os.path.join(args.log_dir, 'cfg.yaml') + cfg = update_cfg(cfg_file) + + if args.misc is not None: + cfg.merge_from_list(args.misc) + else: + parse_args(args) diff --git a/lib/pymaf/core/constants.py b/lib/pymaf/core/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..f2e832437d732fb5c32fc50f975bd9f1c7a750ec --- /dev/null +++ b/lib/pymaf/core/constants.py @@ -0,0 +1,153 @@ +# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/constants.py +FOCAL_LENGTH = 5000. +IMG_RES = 224 + +# Mean and standard deviation for normalizing input image +IMG_NORM_MEAN = [0.485, 0.456, 0.406] +IMG_NORM_STD = [0.229, 0.224, 0.225] +""" +We create a superset of joints containing the OpenPose joints together with the ones that each dataset provides. +We keep a superset of 24 joints such that we include all joints from every dataset. +If a dataset doesn't provide annotations for a specific joint, we simply ignore it. +The joints used here are the following: +""" +JOINT_NAMES = [ + # 25 OpenPose joints (in the order provided by OpenPose) + 'OP Nose', + 'OP Neck', + 'OP RShoulder', + 'OP RElbow', + 'OP RWrist', + 'OP LShoulder', + 'OP LElbow', + 'OP LWrist', + 'OP MidHip', + 'OP RHip', + 'OP RKnee', + 'OP RAnkle', + 'OP LHip', + 'OP LKnee', + 'OP LAnkle', + 'OP REye', + 'OP LEye', + 'OP REar', + 'OP LEar', + 'OP LBigToe', + 'OP LSmallToe', + 'OP LHeel', + 'OP RBigToe', + 'OP RSmallToe', + 'OP RHeel', + # 24 Ground Truth joints (superset of joints from different datasets) + 'Right Ankle', + 'Right Knee', + 'Right Hip', # 2 + 'Left Hip', + 'Left Knee', # 4 + 'Left Ankle', + 'Right Wrist', # 6 + 'Right Elbow', + 'Right Shoulder', # 8 + 'Left Shoulder', + 'Left Elbow', # 10 + 'Left Wrist', + 'Neck (LSP)', # 12 + 'Top of Head (LSP)', + 'Pelvis (MPII)', # 14 + 'Thorax (MPII)', + 'Spine (H36M)', # 16 + 'Jaw (H36M)', + 'Head (H36M)', # 18 + 'Nose', + 'Left Eye', + 'Right Eye', + 'Left Ear', + 'Right Ear' +] + +# Dict containing the joints in numerical order +JOINT_IDS = {JOINT_NAMES[i]: i for i in range(len(JOINT_NAMES))} + +# Map joints to SMPL joints +JOINT_MAP = { + 'OP Nose': 24, + 'OP Neck': 12, + 'OP RShoulder': 17, + 'OP RElbow': 19, + 'OP RWrist': 21, + 'OP LShoulder': 16, + 'OP LElbow': 18, + 'OP LWrist': 20, + 'OP MidHip': 0, + 'OP RHip': 2, + 'OP RKnee': 5, + 'OP RAnkle': 8, + 'OP LHip': 1, + 'OP LKnee': 4, + 'OP LAnkle': 7, + 'OP REye': 25, + 'OP LEye': 26, + 'OP REar': 27, + 'OP LEar': 28, + 'OP LBigToe': 29, + 'OP LSmallToe': 30, + 'OP LHeel': 31, + 'OP RBigToe': 32, + 'OP RSmallToe': 33, + 'OP RHeel': 34, + 'Right Ankle': 8, + 'Right Knee': 5, + 'Right Hip': 45, + 'Left Hip': 46, + 'Left Knee': 4, + 'Left Ankle': 7, + 'Right Wrist': 21, + 'Right Elbow': 19, + 'Right Shoulder': 17, + 'Left Shoulder': 16, + 'Left Elbow': 18, + 'Left Wrist': 20, + 'Neck (LSP)': 47, + 'Top of Head (LSP)': 48, + 'Pelvis (MPII)': 49, + 'Thorax (MPII)': 50, + 'Spine (H36M)': 51, + 'Jaw (H36M)': 52, + 'Head (H36M)': 53, + 'Nose': 24, + 'Left Eye': 26, + 'Right Eye': 25, + 'Left Ear': 28, + 'Right Ear': 27 +} + +# Joint selectors +# Indices to get the 14 LSP joints from the 17 H36M joints +H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9] +H36M_TO_J14 = H36M_TO_J17[:14] +# Indices to get the 14 LSP joints from the ground truth joints +J24_TO_J17 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18, 14, 16, 17] +J24_TO_J14 = J24_TO_J17[:14] +J24_TO_J19 = J24_TO_J17[:14] + [19, 20, 21, 22, 23] +J24_TO_JCOCO = [19, 20, 21, 22, 23, 9, 8, 10, 7, 11, 6, 3, 2, 4, 1, 5, 0] + +# Permutation of SMPL pose parameters when flipping the shape +SMPL_JOINTS_FLIP_PERM = [ + 0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, + 20, 23, 22 +] +SMPL_POSE_FLIP_PERM = [] +for i in SMPL_JOINTS_FLIP_PERM: + SMPL_POSE_FLIP_PERM.append(3 * i) + SMPL_POSE_FLIP_PERM.append(3 * i + 1) + SMPL_POSE_FLIP_PERM.append(3 * i + 2) +# Permutation indices for the 24 ground truth joints +J24_FLIP_PERM = [ + 5, 4, 3, 2, 1, 0, 11, 10, 9, 8, 7, 6, 12, 13, 14, 15, 16, 17, 18, 19, 21, + 20, 23, 22 +] +# Permutation indices for the full set of 49 joints +J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\ + + [25+i for i in J24_FLIP_PERM] +SMPL_J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\ + + [25+i for i in SMPL_JOINTS_FLIP_PERM] diff --git a/lib/pymaf/core/fits_dict.py b/lib/pymaf/core/fits_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..ab78cdd237586365c430824f60b937d35708d97c --- /dev/null +++ b/lib/pymaf/core/fits_dict.py @@ -0,0 +1,133 @@ +''' +This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/train/fits_dict.py +''' +import os +import cv2 +import torch +import numpy as np +from torchgeometry import angle_axis_to_rotation_matrix, rotation_matrix_to_angle_axis + +from core import path_config, constants + +import logging + +logger = logging.getLogger(__name__) + + +class FitsDict(): + """ Dictionary keeping track of the best fit per image in the training set """ + + def __init__(self, options, train_dataset): + self.options = options + self.train_dataset = train_dataset + self.fits_dict = {} + self.valid_fit_state = {} + # array used to flip SMPL pose parameters + self.flipped_parts = torch.tensor(constants.SMPL_POSE_FLIP_PERM, + dtype=torch.int64) + # Load dictionary state + for ds_name, ds in train_dataset.dataset_dict.items(): + if ds_name in ['h36m']: + dict_file = os.path.join(path_config.FINAL_FITS_DIR, + ds_name + '.npy') + self.fits_dict[ds_name] = torch.from_numpy(np.load(dict_file)) + self.valid_fit_state[ds_name] = torch.ones(len( + self.fits_dict[ds_name]), + dtype=torch.uint8) + else: + dict_file = os.path.join(path_config.FINAL_FITS_DIR, + ds_name + '.npz') + fits_dict = np.load(dict_file) + opt_pose = torch.from_numpy(fits_dict['pose']) + opt_betas = torch.from_numpy(fits_dict['betas']) + opt_valid_fit = torch.from_numpy(fits_dict['valid_fit']).to( + torch.uint8) + self.fits_dict[ds_name] = torch.cat([opt_pose, opt_betas], + dim=1) + self.valid_fit_state[ds_name] = opt_valid_fit + + if not options.single_dataset: + for ds in train_dataset.datasets: + if ds.dataset not in ['h36m']: + ds.pose = self.fits_dict[ds.dataset][:, :72].numpy() + ds.betas = self.fits_dict[ds.dataset][:, 72:].numpy() + ds.has_smpl = self.valid_fit_state[ds.dataset].numpy() + + def save(self): + """ Save dictionary state to disk """ + for ds_name in self.train_dataset.dataset_dict.keys(): + dict_file = os.path.join(self.options.checkpoint_dir, + ds_name + '_fits.npy') + np.save(dict_file, self.fits_dict[ds_name].cpu().numpy()) + + def __getitem__(self, x): + """ Retrieve dictionary entries """ + dataset_name, ind, rot, is_flipped = x + batch_size = len(dataset_name) + pose = torch.zeros((batch_size, 72)) + betas = torch.zeros((batch_size, 10)) + for ds, i, n in zip(dataset_name, ind, range(batch_size)): + params = self.fits_dict[ds][i] + pose[n, :] = params[:72] + betas[n, :] = params[72:] + pose = pose.clone() + # Apply flipping and rotation + pose = self.flip_pose(self.rotate_pose(pose, rot), is_flipped) + betas = betas.clone() + return pose, betas + + def get_vaild_state(self, dataset_name, ind): + batch_size = len(dataset_name) + valid_fit = torch.zeros(batch_size, dtype=torch.uint8) + for ds, i, n in zip(dataset_name, ind, range(batch_size)): + valid_fit[n] = self.valid_fit_state[ds][i] + valid_fit = valid_fit.clone() + return valid_fit + + def __setitem__(self, x, val): + """ Update dictionary entries """ + dataset_name, ind, rot, is_flipped, update = x + pose, betas = val + batch_size = len(dataset_name) + # Undo flipping and rotation + pose = self.rotate_pose(self.flip_pose(pose, is_flipped), -rot) + params = torch.cat((pose, betas), dim=-1).cpu() + for ds, i, n in zip(dataset_name, ind, range(batch_size)): + if update[n]: + self.fits_dict[ds][i] = params[n] + + def flip_pose(self, pose, is_flipped): + """flip SMPL pose parameters""" + is_flipped = is_flipped.byte() + pose_f = pose.clone() + pose_f[is_flipped, :] = pose[is_flipped][:, self.flipped_parts] + # we also negate the second and the third dimension of the axis-angle representation + pose_f[is_flipped, 1::3] *= -1 + pose_f[is_flipped, 2::3] *= -1 + return pose_f + + def rotate_pose(self, pose, rot): + """Rotate SMPL pose parameters by rot degrees""" + pose = pose.clone() + cos = torch.cos(-np.pi * rot / 180.) + sin = torch.sin(-np.pi * rot / 180.) + zeros = torch.zeros_like(cos) + r3 = torch.zeros(cos.shape[0], 1, 3, device=cos.device) + r3[:, 0, -1] = 1 + R = torch.cat([ + torch.stack([cos, -sin, zeros], dim=-1).unsqueeze(1), + torch.stack([sin, cos, zeros], dim=-1).unsqueeze(1), r3 + ], + dim=1) + global_pose = pose[:, :3] + global_pose_rotmat = angle_axis_to_rotation_matrix(global_pose) + global_pose_rotmat_3b3 = global_pose_rotmat[:, :3, :3] + global_pose_rotmat_3b3 = torch.matmul(R, global_pose_rotmat_3b3) + global_pose_rotmat[:, :3, :3] = global_pose_rotmat_3b3 + global_pose_rotmat = global_pose_rotmat[:, :-1, :-1].cpu().numpy() + global_pose_np = np.zeros((global_pose.shape[0], 3)) + for i in range(global_pose.shape[0]): + aa, _ = cv2.Rodrigues(global_pose_rotmat[i]) + global_pose_np[i, :] = aa.squeeze() + pose[:, :3] = torch.from_numpy(global_pose_np).to(pose.device) + return pose diff --git a/lib/pymaf/core/path_config.py b/lib/pymaf/core/path_config.py new file mode 100644 index 0000000000000000000000000000000000000000..0623abf2552c75688aa4e7c0ad2e03ae74c5f698 --- /dev/null +++ b/lib/pymaf/core/path_config.py @@ -0,0 +1,38 @@ +""" +This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/path_config.py +path configuration +This file contains definitions of useful data stuctures and the paths +for the datasets and data files necessary to run the code. +Things you need to change: *_ROOT that indicate the path to each dataset +""" +import os + +# pymaf +pymaf_data_dir = os.path.join(os.path.dirname(__file__), + "../../../data/pymaf_data") + +SMPL_MEAN_PARAMS = os.path.join(pymaf_data_dir, 'smpl_mean_params.npz') +SMPL_MODEL_DIR = os.path.join(pymaf_data_dir, '../smpl_related/models/smpl') + +CUBE_PARTS_FILE = os.path.join(pymaf_data_dir, 'cube_parts.npy') +JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(pymaf_data_dir, + 'J_regressor_extra.npy') +JOINT_REGRESSOR_H36M = os.path.join(pymaf_data_dir, 'J_regressor_h36m.npy') +VERTEX_TEXTURE_FILE = os.path.join(pymaf_data_dir, 'vertex_texture.npy') +SMPL_MEAN_PARAMS = os.path.join(pymaf_data_dir, 'smpl_mean_params.npz') +SMPL_MODEL_DIR = os.path.join(pymaf_data_dir, '../smpl_related/models/smpl') +CHECKPOINT_FILE = os.path.join(pymaf_data_dir, + 'pretrained_model/PyMAF_model_checkpoint.pt') + +# pare +pare_data_dir = os.path.join(os.path.dirname(__file__), + "../../../data/pare_data") +CFG = os.path.join(pare_data_dir, 'pare/checkpoints/pare_w_3dpw_config.yaml') +CKPT = os.path.join(pare_data_dir, + 'pare/checkpoints/pare_w_3dpw_checkpoint.ckpt') + +# hybrik +hybrik_data_dir = os.path.join(os.path.dirname(__file__), + "../../../data/hybrik_data") +HYBRIK_CFG = os.path.join(hybrik_data_dir, 'hybrik_config.yaml') +HYBRIK_CKPT = os.path.join(hybrik_data_dir, 'pretrained_w_cam.pth') diff --git a/lib/pymaf/core/train_options.py b/lib/pymaf/core/train_options.py new file mode 100644 index 0000000000000000000000000000000000000000..43daeb4486e8751b903b3852b066e5b3e13bd9de --- /dev/null +++ b/lib/pymaf/core/train_options.py @@ -0,0 +1,135 @@ +import argparse + + +class TrainOptions(): + def __init__(self): + self.parser = argparse.ArgumentParser() + + gen = self.parser.add_argument_group('General') + gen.add_argument( + '--resume', + dest='resume', + default=False, + action='store_true', + help='Resume from checkpoint (Use latest checkpoint by default') + + io = self.parser.add_argument_group('io') + io.add_argument('--log_dir', + default='logs', + help='Directory to store logs') + io.add_argument( + '--pretrained_checkpoint', + default=None, + help='Load a pretrained checkpoint at the beginning training') + + train = self.parser.add_argument_group('Training Options') + train.add_argument('--num_epochs', + type=int, + default=200, + help='Total number of training epochs') + train.add_argument('--regressor', + type=str, + choices=['hmr', 'pymaf_net'], + default='pymaf_net', + help='Name of the SMPL regressor.') + train.add_argument('--cfg_file', + type=str, + default='./configs/pymaf_config.yaml', + help='config file path for PyMAF.') + train.add_argument( + '--img_res', + type=int, + default=224, + help='Rescale bounding boxes to size [img_res, img_res] before feeding them in the network' + ) + train.add_argument( + '--rot_factor', + type=float, + default=30, + help='Random rotation in the range [-rot_factor, rot_factor]') + train.add_argument( + '--noise_factor', + type=float, + default=0.4, + help='Randomly multiply pixel values with factor in the range [1-noise_factor, 1+noise_factor]' + ) + train.add_argument( + '--scale_factor', + type=float, + default=0.25, + help='Rescale bounding boxes by a factor of [1-scale_factor,1+scale_factor]' + ) + train.add_argument( + '--openpose_train_weight', + default=0., + help='Weight for OpenPose keypoints during training') + train.add_argument('--gt_train_weight', + default=1., + help='Weight for GT keypoints during training') + train.add_argument('--eval_dataset', + type=str, + default='h36m-p2-mosh', + help='Name of the evaluation dataset.') + train.add_argument('--single_dataset', + default=False, + action='store_true', + help='Use a single dataset') + train.add_argument('--single_dataname', + type=str, + default='h36m', + help='Name of the single dataset.') + train.add_argument('--eval_pve', + default=False, + action='store_true', + help='evaluate PVE') + train.add_argument('--overwrite', + default=False, + action='store_true', + help='overwrite the latest checkpoint') + + train.add_argument('--distributed', + action='store_true', + help='Use distributed training') + train.add_argument('--dist_backend', + default='nccl', + type=str, + help='distributed backend') + train.add_argument('--dist_url', + default='tcp://127.0.0.1:10356', + type=str, + help='url used to set up distributed training') + train.add_argument('--world_size', + default=1, + type=int, + help='number of nodes for distributed training') + train.add_argument("--local_rank", default=0, type=int) + train.add_argument('--rank', + default=0, + type=int, + help='node rank for distributed training') + train.add_argument( + '--multiprocessing_distributed', + action='store_true', + help='Use multi-processing distributed training to launch ' + 'N processes per node, which has N GPUs. This is the ' + 'fastest way to use PyTorch for either single node or ' + 'multi node data parallel training') + + misc = self.parser.add_argument_group('Misc Options') + misc.add_argument('--misc', + help="Modify config options using the command-line", + default=None, + nargs=argparse.REMAINDER) + return + + def parse_args(self): + """Parse input arguments.""" + self.args = self.parser.parse_args() + self.save_dump() + return self.args + + def save_dump(self): + """Store all argument values to a json file. + The default location is logs/expname/args.json. + """ + pass diff --git a/lib/pymaf/models/__init__.py b/lib/pymaf/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c85ca9042485c0af8b1f29a4d4eaa1547935a40f --- /dev/null +++ b/lib/pymaf/models/__init__.py @@ -0,0 +1,3 @@ +from .hmr import hmr +from .pymaf_net import pymaf_net +from .smpl import SMPL diff --git a/lib/pymaf/models/hmr.py b/lib/pymaf/models/hmr.py new file mode 100644 index 0000000000000000000000000000000000000000..9fb1cd6b7e2e4581f2c5d9cb5e952049b5a075e1 --- /dev/null +++ b/lib/pymaf/models/hmr.py @@ -0,0 +1,303 @@ +# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/hmr.py + +import torch +import torch.nn as nn +import torchvision.models.resnet as resnet +import numpy as np +import math +from lib.pymaf.utils.geometry import rot6d_to_rotmat + +import logging + +logger = logging.getLogger(__name__) + +BN_MOMENTUM = 0.1 + + +class Bottleneck(nn.Module): + """ Redefinition of Bottleneck residual block + Adapted from the official PyTorch implementation + """ + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super().__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, + planes, + kernel_size=3, + stride=stride, + padding=1, + bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet_Backbone(nn.Module): + """ Feature Extrator with ResNet backbone + """ + + def __init__(self, model='res50', pretrained=True): + if model == 'res50': + block, layers = Bottleneck, [3, 4, 6, 3] + else: + pass # TODO + + self.inplanes = 64 + super().__init__() + npose = 24 * 6 + self.conv1 = nn.Conv2d(3, + 64, + kernel_size=7, + stride=2, + padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + self.avgpool = nn.AvgPool2d(7, stride=1) + + if pretrained: + resnet_imagenet = resnet.resnet50(pretrained=True) + self.load_state_dict(resnet_imagenet.state_dict(), strict=False) + logger.info('loaded resnet50 imagenet pretrained model') + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _make_deconv_layer(self, num_layers, num_filters, num_kernels): + assert num_layers == len(num_filters), \ + 'ERROR: num_deconv_layers is different len(num_deconv_filters)' + assert num_layers == len(num_kernels), \ + 'ERROR: num_deconv_layers is different len(num_deconv_filters)' + + def _get_deconv_cfg(deconv_kernel, index): + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel == 2: + padding = 0 + output_padding = 0 + + return deconv_kernel, padding, output_padding + + layers = [] + for i in range(num_layers): + kernel, padding, output_padding = _get_deconv_cfg( + num_kernels[i], i) + + planes = num_filters[i] + layers.append( + nn.ConvTranspose2d(in_channels=self.inplanes, + out_channels=planes, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias)) + layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)) + layers.append(nn.ReLU(inplace=True)) + self.inplanes = planes + + return nn.Sequential(*layers) + + def forward(self, x): + + batch_size = x.shape[0] + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x1 = self.layer1(x) + x2 = self.layer2(x1) + x3 = self.layer3(x2) + x4 = self.layer4(x3) + + xf = self.avgpool(x4) + xf = xf.view(xf.size(0), -1) + + x_featmap = x4 + + return x_featmap, xf + + +class HMR(nn.Module): + """ SMPL Iterative Regressor with ResNet50 backbone + """ + + def __init__(self, block, layers, smpl_mean_params): + self.inplanes = 64 + super().__init__() + npose = 24 * 6 + self.conv1 = nn.Conv2d(3, + 64, + kernel_size=7, + stride=2, + padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + self.avgpool = nn.AvgPool2d(7, stride=1) + self.fc1 = nn.Linear(512 * block.expansion + npose + 13, 1024) + self.drop1 = nn.Dropout() + self.fc2 = nn.Linear(1024, 1024) + self.drop2 = nn.Dropout() + self.decpose = nn.Linear(1024, npose) + self.decshape = nn.Linear(1024, 10) + self.deccam = nn.Linear(1024, 3) + nn.init.xavier_uniform_(self.decpose.weight, gain=0.01) + nn.init.xavier_uniform_(self.decshape.weight, gain=0.01) + nn.init.xavier_uniform_(self.deccam.weight, gain=0.01) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + mean_params = np.load(smpl_mean_params) + init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0) + init_shape = torch.from_numpy( + mean_params['shape'][:].astype('float32')).unsqueeze(0) + init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0) + self.register_buffer('init_pose', init_pose) + self.register_buffer('init_shape', init_shape) + self.register_buffer('init_cam', init_cam) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, + x, + init_pose=None, + init_shape=None, + init_cam=None, + n_iter=3): + + batch_size = x.shape[0] + + if init_pose is None: + init_pose = self.init_pose.expand(batch_size, -1) + if init_shape is None: + init_shape = self.init_shape.expand(batch_size, -1) + if init_cam is None: + init_cam = self.init_cam.expand(batch_size, -1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x1 = self.layer1(x) + x2 = self.layer2(x1) + x3 = self.layer3(x2) + x4 = self.layer4(x3) + + xf = self.avgpool(x4) + xf = xf.view(xf.size(0), -1) + + pred_pose = init_pose + pred_shape = init_shape + pred_cam = init_cam + for i in range(n_iter): + xc = torch.cat([xf, pred_pose, pred_shape, pred_cam], 1) + xc = self.fc1(xc) + xc = self.drop1(xc) + xc = self.fc2(xc) + xc = self.drop2(xc) + pred_pose = self.decpose(xc) + pred_pose + pred_shape = self.decshape(xc) + pred_shape + pred_cam = self.deccam(xc) + pred_cam + + pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) + + return pred_rotmat, pred_shape, pred_cam + + +def hmr(smpl_mean_params, pretrained=True, **kwargs): + """ Constructs an HMR model with ResNet50 backbone. + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = HMR(Bottleneck, [3, 4, 6, 3], smpl_mean_params, **kwargs) + if pretrained: + resnet_imagenet = resnet.resnet50(pretrained=True) + model.load_state_dict(resnet_imagenet.state_dict(), strict=False) + return model diff --git a/lib/pymaf/models/maf_extractor.py b/lib/pymaf/models/maf_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..8a4a1b2ef69ee16f0fd8a5cf15642ff32766ca97 --- /dev/null +++ b/lib/pymaf/models/maf_extractor.py @@ -0,0 +1,140 @@ +# This script is borrowed and extended from https://github.com/shunsukesaito/PIFu/blob/master/lib/model/SurfaceClassifier.py + +from packaging import version +import torch +import scipy +import os +import numpy as np +import torch.nn as nn +import torch.nn.functional as F + +from lib.common.config import cfg +from lib.pymaf.utils.geometry import projection + +import logging + +logger = logging.getLogger(__name__) + + +class MAF_Extractor(nn.Module): + ''' Mesh-aligned Feature Extrator + + As discussed in the paper, we extract mesh-aligned features based on 2D projection of the mesh vertices. + The features extrated from spatial feature maps will go through a MLP for dimension reduction. + ''' + + def __init__(self, device=torch.device('cuda')): + super().__init__() + + self.device = device + self.filters = [] + self.num_views = 1 + filter_channels = cfg.MODEL.PyMAF.MLP_DIM + self.last_op = nn.ReLU(True) + + for l in range(0, len(filter_channels) - 1): + if 0 != l: + self.filters.append( + nn.Conv1d(filter_channels[l] + filter_channels[0], + filter_channels[l + 1], 1)) + else: + self.filters.append( + nn.Conv1d(filter_channels[l], filter_channels[l + 1], 1)) + + self.add_module("conv%d" % l, self.filters[l]) + + self.im_feat = None + self.cam = None + + # downsample SMPL mesh and assign part labels + # from https://github.com/nkolot/GraphCMR/blob/master/data/mesh_downsampling.npz + mesh_downsampling_path = os.path.join( + os.path.dirname(__file__), + "../../../data/pymaf_data/mesh_downsampling.npz") + smpl_mesh_graph = np.load(mesh_downsampling_path, + allow_pickle=True, + encoding='latin1') + + A = smpl_mesh_graph['A'] + U = smpl_mesh_graph['U'] + D = smpl_mesh_graph['D'] # shape: (2,) + + # downsampling + ptD = [] + for i in range(len(D)): + d = scipy.sparse.coo_matrix(D[i]) + i = torch.LongTensor(np.array([d.row, d.col])) + v = torch.FloatTensor(d.data) + ptD.append(torch.sparse.FloatTensor(i, v, d.shape)) + + # downsampling mapping from 6890 points to 431 points + # ptD[0].to_dense() - Size: [1723, 6890] + # ptD[1].to_dense() - Size: [431. 1723] + Dmap = torch.matmul(ptD[1].to_dense(), + ptD[0].to_dense()) # 6890 -> 431 + self.register_buffer('Dmap', Dmap) + + def reduce_dim(self, feature): + ''' + Dimension reduction by multi-layer perceptrons + :param feature: list of [B, C_s, N] point-wise features before dimension reduction + :return: [B, C_p x N] concatantion of point-wise features after dimension reduction + ''' + y = feature + tmpy = feature + for i, f in enumerate(self.filters): + y = self._modules['conv' + + str(i)](y if i == 0 else torch.cat([y, tmpy], 1)) + if i != len(self.filters) - 1: + y = F.leaky_relu(y) + if self.num_views > 1 and i == len(self.filters) // 2: + y = y.view(-1, self.num_views, y.shape[1], + y.shape[2]).mean(dim=1) + tmpy = feature.view(-1, self.num_views, feature.shape[1], + feature.shape[2]).mean(dim=1) + + y = self.last_op(y) + + y = y.view(y.shape[0], -1) + return y + + def sampling(self, points, im_feat=None, z_feat=None): + ''' + Given 2D points, sample the point-wise features for each point, + the dimension of point-wise features will be reduced from C_s to C_p by MLP. + Image features should be pre-computed before this call. + :param points: [B, N, 2] image coordinates of points + :im_feat: [B, C_s, H_s, W_s] spatial feature maps + :return: [B, C_p x N] concatantion of point-wise features after dimension reduction + ''' + if im_feat is None: + im_feat = self.im_feat + + batch_size = im_feat.shape[0] + + if version.parse(torch.__version__) >= version.parse('1.3.0'): + # Default grid_sample behavior has changed to align_corners=False since 1.3.0. + point_feat = torch.nn.functional.grid_sample( + im_feat, points.unsqueeze(2), align_corners=True)[..., 0] + else: + point_feat = torch.nn.functional.grid_sample( + im_feat, points.unsqueeze(2))[..., 0] + + mesh_align_feat = self.reduce_dim(point_feat) + return mesh_align_feat + + def forward(self, p, s_feat=None, cam=None, **kwargs): + ''' Returns mesh-aligned features for the 3D mesh points. + + Args: + p (tensor): [B, N_m, 3] mesh vertices + s_feat (tensor): [B, C_s, H_s, W_s] spatial feature maps + cam (tensor): [B, 3] camera + Return: + mesh_align_feat (tensor): [B, C_p x N_m] mesh-aligned features + ''' + if cam is None: + cam = self.cam + p_proj_2d = projection(p, cam, retain_z=False) + mesh_align_feat = self.sampling(p_proj_2d, s_feat) + return mesh_align_feat diff --git a/lib/pymaf/models/pymaf_net.py b/lib/pymaf/models/pymaf_net.py new file mode 100644 index 0000000000000000000000000000000000000000..2807abaa3c7da0be6913d2fd68cb0ad1721e2bf1 --- /dev/null +++ b/lib/pymaf/models/pymaf_net.py @@ -0,0 +1,362 @@ +import torch +import torch.nn as nn +import numpy as np + +from lib.pymaf.utils.geometry import rot6d_to_rotmat, projection, rotation_matrix_to_angle_axis +from .maf_extractor import MAF_Extractor +from .smpl import SMPL, SMPL_MODEL_DIR, SMPL_MEAN_PARAMS, H36M_TO_J14 +from .hmr import ResNet_Backbone +from .res_module import IUV_predict_layer +from lib.common.config import cfg +import logging + +logger = logging.getLogger(__name__) + +BN_MOMENTUM = 0.1 + + +class Regressor(nn.Module): + def __init__(self, feat_dim, smpl_mean_params): + super().__init__() + + npose = 24 * 6 + + self.fc1 = nn.Linear(feat_dim + npose + 13, 1024) + self.drop1 = nn.Dropout() + self.fc2 = nn.Linear(1024, 1024) + self.drop2 = nn.Dropout() + self.decpose = nn.Linear(1024, npose) + self.decshape = nn.Linear(1024, 10) + self.deccam = nn.Linear(1024, 3) + nn.init.xavier_uniform_(self.decpose.weight, gain=0.01) + nn.init.xavier_uniform_(self.decshape.weight, gain=0.01) + nn.init.xavier_uniform_(self.deccam.weight, gain=0.01) + + self.smpl = SMPL(SMPL_MODEL_DIR, batch_size=64, create_transl=False) + + mean_params = np.load(smpl_mean_params) + init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0) + init_shape = torch.from_numpy( + mean_params['shape'][:].astype('float32')).unsqueeze(0) + init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0) + self.register_buffer('init_pose', init_pose) + self.register_buffer('init_shape', init_shape) + self.register_buffer('init_cam', init_cam) + + def forward(self, + x, + init_pose=None, + init_shape=None, + init_cam=None, + n_iter=1, + J_regressor=None): + batch_size = x.shape[0] + + if init_pose is None: + init_pose = self.init_pose.expand(batch_size, -1) + if init_shape is None: + init_shape = self.init_shape.expand(batch_size, -1) + if init_cam is None: + init_cam = self.init_cam.expand(batch_size, -1) + + pred_pose = init_pose + pred_shape = init_shape + pred_cam = init_cam + for i in range(n_iter): + xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1) + xc = self.fc1(xc) + xc = self.drop1(xc) + xc = self.fc2(xc) + xc = self.drop2(xc) + pred_pose = self.decpose(xc) + pred_pose + pred_shape = self.decshape(xc) + pred_shape + pred_cam = self.deccam(xc) + pred_cam + + pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) + + pred_output = self.smpl(betas=pred_shape, + body_pose=pred_rotmat[:, 1:], + global_orient=pred_rotmat[:, 0].unsqueeze(1), + pose2rot=False) + + pred_vertices = pred_output.vertices + pred_joints = pred_output.joints + pred_smpl_joints = pred_output.smpl_joints + pred_keypoints_2d = projection(pred_joints, pred_cam) + pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, + 3)).reshape( + -1, 72) + + if J_regressor is not None: + pred_joints = torch.matmul(J_regressor, pred_vertices) + pred_pelvis = pred_joints[:, [0], :].clone() + pred_joints = pred_joints[:, H36M_TO_J14, :] + pred_joints = pred_joints - pred_pelvis + + output = { + 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), + 'verts': pred_vertices, + 'kp_2d': pred_keypoints_2d, + 'kp_3d': pred_joints, + 'smpl_kp_3d': pred_smpl_joints, + 'rotmat': pred_rotmat, + 'pred_cam': pred_cam, + 'pred_shape': pred_shape, + 'pred_pose': pred_pose, + } + return output + + def forward_init(self, + x, + init_pose=None, + init_shape=None, + init_cam=None, + n_iter=1, + J_regressor=None): + batch_size = x.shape[0] + + if init_pose is None: + init_pose = self.init_pose.expand(batch_size, -1) + if init_shape is None: + init_shape = self.init_shape.expand(batch_size, -1) + if init_cam is None: + init_cam = self.init_cam.expand(batch_size, -1) + + pred_pose = init_pose + pred_shape = init_shape + pred_cam = init_cam + + pred_rotmat = rot6d_to_rotmat(pred_pose.contiguous()).view( + batch_size, 24, 3, 3) + + pred_output = self.smpl(betas=pred_shape, + body_pose=pred_rotmat[:, 1:], + global_orient=pred_rotmat[:, 0].unsqueeze(1), + pose2rot=False) + + pred_vertices = pred_output.vertices + pred_joints = pred_output.joints + pred_smpl_joints = pred_output.smpl_joints + pred_keypoints_2d = projection(pred_joints, pred_cam) + pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, + 3)).reshape( + -1, 72) + + if J_regressor is not None: + pred_joints = torch.matmul(J_regressor, pred_vertices) + pred_pelvis = pred_joints[:, [0], :].clone() + pred_joints = pred_joints[:, H36M_TO_J14, :] + pred_joints = pred_joints - pred_pelvis + + output = { + 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), + 'verts': pred_vertices, + 'kp_2d': pred_keypoints_2d, + 'kp_3d': pred_joints, + 'smpl_kp_3d': pred_smpl_joints, + 'rotmat': pred_rotmat, + 'pred_cam': pred_cam, + 'pred_shape': pred_shape, + 'pred_pose': pred_pose, + } + return output + + +class PyMAF(nn.Module): + """ PyMAF based Deep Regressor for Human Mesh Recovery + PyMAF: 3D Human Pose and Shape Regression with Pyramidal Mesh Alignment Feedback Loop, in ICCV, 2021 + """ + + def __init__(self, smpl_mean_params=SMPL_MEAN_PARAMS, pretrained=True): + super().__init__() + self.feature_extractor = ResNet_Backbone( + model=cfg.MODEL.PyMAF.BACKBONE, pretrained=pretrained) + + # deconv layers + self.inplanes = self.feature_extractor.inplanes + self.deconv_with_bias = cfg.RES_MODEL.DECONV_WITH_BIAS + self.deconv_layers = self._make_deconv_layer( + cfg.RES_MODEL.NUM_DECONV_LAYERS, + cfg.RES_MODEL.NUM_DECONV_FILTERS, + cfg.RES_MODEL.NUM_DECONV_KERNELS, + ) + + self.maf_extractor = nn.ModuleList() + for _ in range(cfg.MODEL.PyMAF.N_ITER): + self.maf_extractor.append(MAF_Extractor()) + ma_feat_len = self.maf_extractor[-1].Dmap.shape[ + 0] * cfg.MODEL.PyMAF.MLP_DIM[-1] + + grid_size = 21 + xv, yv = torch.meshgrid([ + torch.linspace(-1, 1, grid_size), + torch.linspace(-1, 1, grid_size) + ]) + points_grid = torch.stack([xv.reshape(-1), + yv.reshape(-1)]).unsqueeze(0) + self.register_buffer('points_grid', points_grid) + grid_feat_len = grid_size * grid_size * cfg.MODEL.PyMAF.MLP_DIM[-1] + + self.regressor = nn.ModuleList() + for i in range(cfg.MODEL.PyMAF.N_ITER): + if i == 0: + ref_infeat_dim = grid_feat_len + else: + ref_infeat_dim = ma_feat_len + self.regressor.append( + Regressor(feat_dim=ref_infeat_dim, + smpl_mean_params=smpl_mean_params)) + + dp_feat_dim = 256 + self.with_uv = cfg.LOSS.POINT_REGRESSION_WEIGHTS > 0 + if cfg.MODEL.PyMAF.AUX_SUPV_ON: + self.dp_head = IUV_predict_layer(feat_dim=dp_feat_dim) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _make_deconv_layer(self, num_layers, num_filters, num_kernels): + """ + Deconv_layer used in Simple Baselines: + Xiao et al. Simple Baselines for Human Pose Estimation and Tracking + https://github.com/microsoft/human-pose-estimation.pytorch + """ + assert num_layers == len(num_filters), \ + 'ERROR: num_deconv_layers is different len(num_deconv_filters)' + assert num_layers == len(num_kernels), \ + 'ERROR: num_deconv_layers is different len(num_deconv_filters)' + + def _get_deconv_cfg(deconv_kernel, index): + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel == 2: + padding = 0 + output_padding = 0 + + return deconv_kernel, padding, output_padding + + layers = [] + for i in range(num_layers): + kernel, padding, output_padding = _get_deconv_cfg( + num_kernels[i], i) + + planes = num_filters[i] + layers.append( + nn.ConvTranspose2d(in_channels=self.inplanes, + out_channels=planes, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias)) + layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)) + layers.append(nn.ReLU(inplace=True)) + self.inplanes = planes + + return nn.Sequential(*layers) + + def forward(self, x, J_regressor=None): + + batch_size = x.shape[0] + + # spatial features and global features + s_feat, g_feat = self.feature_extractor(x) + + assert cfg.MODEL.PyMAF.N_ITER >= 0 and cfg.MODEL.PyMAF.N_ITER <= 3 + if cfg.MODEL.PyMAF.N_ITER == 1: + deconv_blocks = [self.deconv_layers] + elif cfg.MODEL.PyMAF.N_ITER == 2: + deconv_blocks = [self.deconv_layers[0:6], self.deconv_layers[6:9]] + elif cfg.MODEL.PyMAF.N_ITER == 3: + deconv_blocks = [ + self.deconv_layers[0:3], self.deconv_layers[3:6], + self.deconv_layers[6:9] + ] + + out_list = {} + + # initial parameters + # TODO: remove the initial mesh generation during forward to reduce runtime + # by generating initial mesh the beforehand: smpl_output = self.init_smpl + smpl_output = self.regressor[0].forward_init(g_feat, + J_regressor=J_regressor) + + out_list['smpl_out'] = [smpl_output] + out_list['dp_out'] = [] + + # for visulization + vis_feat_list = [s_feat.detach()] + + # parameter predictions + for rf_i in range(cfg.MODEL.PyMAF.N_ITER): + pred_cam = smpl_output['pred_cam'] + pred_shape = smpl_output['pred_shape'] + pred_pose = smpl_output['pred_pose'] + + pred_cam = pred_cam.detach() + pred_shape = pred_shape.detach() + pred_pose = pred_pose.detach() + + s_feat_i = deconv_blocks[rf_i](s_feat) + s_feat = s_feat_i + vis_feat_list.append(s_feat_i.detach()) + + self.maf_extractor[rf_i].im_feat = s_feat_i + self.maf_extractor[rf_i].cam = pred_cam + + if rf_i == 0: + sample_points = torch.transpose( + self.points_grid.expand(batch_size, -1, -1), 1, 2) + ref_feature = self.maf_extractor[rf_i].sampling(sample_points) + else: + pred_smpl_verts = smpl_output['verts'].detach() + # TODO: use a more sparse SMPL implementation (with 431 vertices) for acceleration + pred_smpl_verts_ds = torch.matmul( + self.maf_extractor[rf_i].Dmap.unsqueeze(0), + pred_smpl_verts) # [B, 431, 3] + ref_feature = self.maf_extractor[rf_i]( + pred_smpl_verts_ds) # [B, 431 * n_feat] + + smpl_output = self.regressor[rf_i](ref_feature, + pred_pose, + pred_shape, + pred_cam, + n_iter=1, + J_regressor=J_regressor) + out_list['smpl_out'].append(smpl_output) + + if self.training and cfg.MODEL.PyMAF.AUX_SUPV_ON: + iuv_out_dict = self.dp_head(s_feat) + out_list['dp_out'].append(iuv_out_dict) + + return out_list + + +def pymaf_net(smpl_mean_params, pretrained=True): + """ Constructs an PyMAF model with ResNet50 backbone. + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = PyMAF(smpl_mean_params, pretrained) + return model diff --git a/lib/pymaf/models/res_module.py b/lib/pymaf/models/res_module.py new file mode 100644 index 0000000000000000000000000000000000000000..35d23f4388be7eab70203ff3b6002eba561525b7 --- /dev/null +++ b/lib/pymaf/models/res_module.py @@ -0,0 +1,385 @@ +# code brought in part from https://github.com/microsoft/human-pose-estimation.pytorch/blob/master/lib/models/pose_resnet.py + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict +from lib.pymaf.core.cfgs import cfg + +import logging + +logger = logging.getLogger(__name__) + +BN_MOMENTUM = 0.1 + + +def conv3x3(in_planes, out_planes, stride=1, bias=False, groups=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes * groups, + out_planes * groups, + kernel_size=3, + stride=stride, + padding=1, + bias=bias, + groups=groups) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1): + super().__init__() + self.conv1 = conv3x3(inplanes, planes, stride, groups=groups) + self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes, groups=groups) + self.bn2 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1): + super().__init__() + self.conv1 = nn.Conv2d(inplanes * groups, + planes * groups, + kernel_size=1, + bias=False, + groups=groups) + self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(planes * groups, + planes * groups, + kernel_size=3, + stride=stride, + padding=1, + bias=False, + groups=groups) + self.bn2 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) + self.conv3 = nn.Conv2d(planes * groups, + planes * self.expansion * groups, + kernel_size=1, + bias=False, + groups=groups) + self.bn3 = nn.BatchNorm2d(planes * self.expansion * groups, + momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +resnet_spec = { + 18: (BasicBlock, [2, 2, 2, 2]), + 34: (BasicBlock, [3, 4, 6, 3]), + 50: (Bottleneck, [3, 4, 6, 3]), + 101: (Bottleneck, [3, 4, 23, 3]), + 152: (Bottleneck, [3, 8, 36, 3]) +} + + +class IUV_predict_layer(nn.Module): + def __init__(self, + feat_dim=256, + final_cov_k=3, + part_out_dim=25, + with_uv=True): + super().__init__() + + self.with_uv = with_uv + if self.with_uv: + self.predict_u = nn.Conv2d(in_channels=feat_dim, + out_channels=25, + kernel_size=final_cov_k, + stride=1, + padding=1 if final_cov_k == 3 else 0) + + self.predict_v = nn.Conv2d(in_channels=feat_dim, + out_channels=25, + kernel_size=final_cov_k, + stride=1, + padding=1 if final_cov_k == 3 else 0) + + self.predict_ann_index = nn.Conv2d( + in_channels=feat_dim, + out_channels=15, + kernel_size=final_cov_k, + stride=1, + padding=1 if final_cov_k == 3 else 0) + + self.predict_uv_index = nn.Conv2d(in_channels=feat_dim, + out_channels=25, + kernel_size=final_cov_k, + stride=1, + padding=1 if final_cov_k == 3 else 0) + + self.inplanes = feat_dim + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + return_dict = {} + + predict_uv_index = self.predict_uv_index(x) + predict_ann_index = self.predict_ann_index(x) + + return_dict['predict_uv_index'] = predict_uv_index + return_dict['predict_ann_index'] = predict_ann_index + + if self.with_uv: + predict_u = self.predict_u(x) + predict_v = self.predict_v(x) + return_dict['predict_u'] = predict_u + return_dict['predict_v'] = predict_v + else: + return_dict['predict_u'] = None + return_dict['predict_v'] = None + # return_dict['predict_u'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device) + # return_dict['predict_v'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device) + + return return_dict + + +class SmplResNet(nn.Module): + def __init__(self, + resnet_nums, + in_channels=3, + num_classes=229, + last_stride=2, + n_extra_feat=0, + truncate=0, + **kwargs): + super().__init__() + + self.inplanes = 64 + self.truncate = truncate + # extra = cfg.MODEL.EXTRA + # self.deconv_with_bias = extra.DECONV_WITH_BIAS + block, layers = resnet_spec[resnet_nums] + + self.conv1 = nn.Conv2d(in_channels, + 64, + kernel_size=7, + stride=2, + padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], + stride=2) if truncate < 2 else None + self.layer4 = self._make_layer( + block, 512, layers[3], + stride=last_stride) if truncate < 1 else None + + self.avg_pooling = nn.AdaptiveAvgPool2d(1) + + self.num_classes = num_classes + if num_classes > 0: + self.final_layer = nn.Linear(512 * block.expansion, num_classes) + nn.init.xavier_uniform_(self.final_layer.weight, gain=0.01) + + self.n_extra_feat = n_extra_feat + if n_extra_feat > 0: + self.trans_conv = nn.Sequential( + nn.Conv2d(n_extra_feat + 512 * block.expansion, + 512 * block.expansion, + kernel_size=1, + bias=False), + nn.BatchNorm2d(512 * block.expansion, momentum=BN_MOMENTUM), + nn.ReLU(True)) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x, infeat=None): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x1 = self.layer1(x) + x2 = self.layer2(x1) + x3 = self.layer3(x2) if self.truncate < 2 else x2 + x4 = self.layer4(x3) if self.truncate < 1 else x3 + + if infeat is not None: + x4 = self.trans_conv(torch.cat([infeat, x4], 1)) + + if self.num_classes > 0: + xp = self.avg_pooling(x4) + cls = self.final_layer(xp.view(xp.size(0), -1)) + if not cfg.DANET.USE_MEAN_PARA: + # for non-negative scale + scale = F.relu(cls[:, 0]).unsqueeze(1) + cls = torch.cat((scale, cls[:, 1:]), dim=1) + else: + cls = None + + return cls, {'x4': x4} + + def init_weights(self, pretrained=''): + if os.path.isfile(pretrained): + logger.info('=> loading pretrained model {}'.format(pretrained)) + # self.load_state_dict(pretrained_state_dict, strict=False) + checkpoint = torch.load(pretrained) + if isinstance(checkpoint, OrderedDict): + # state_dict = checkpoint + state_dict_old = self.state_dict() + for key in state_dict_old.keys(): + if key in checkpoint.keys(): + if state_dict_old[key].shape != checkpoint[key].shape: + del checkpoint[key] + state_dict = checkpoint + elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint: + state_dict_old = checkpoint['state_dict'] + state_dict = OrderedDict() + # delete 'module.' because it is saved from DataParallel module + for key in state_dict_old.keys(): + if key.startswith('module.'): + # state_dict[key[7:]] = state_dict[key] + # state_dict.pop(key) + state_dict[key[7:]] = state_dict_old[key] + else: + state_dict[key] = state_dict_old[key] + else: + raise RuntimeError( + 'No state_dict found in checkpoint file {}'.format( + pretrained)) + self.load_state_dict(state_dict, strict=False) + else: + logger.error('=> imagenet pretrained model dose not exist') + logger.error('=> please download it first') + raise ValueError('imagenet pretrained model does not exist') + + +class LimbResLayers(nn.Module): + def __init__(self, + resnet_nums, + inplanes, + outplanes=None, + groups=1, + **kwargs): + super().__init__() + + self.inplanes = inplanes + block, layers = resnet_spec[resnet_nums] + self.outplanes = 512 if outplanes == None else outplanes + self.layer4 = self._make_layer(block, + self.outplanes, + layers[3], + stride=2, + groups=groups) + + self.avg_pooling = nn.AdaptiveAvgPool2d(1) + + def _make_layer(self, block, planes, blocks, stride=1, groups=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes * groups, + planes * block.expansion * groups, + kernel_size=1, + stride=stride, + bias=False, + groups=groups), + nn.BatchNorm2d(planes * block.expansion * groups, + momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, groups=groups)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, groups=groups)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.layer4(x) + x = self.avg_pooling(x) + + return x diff --git a/lib/pymaf/models/smpl.py b/lib/pymaf/models/smpl.py new file mode 100644 index 0000000000000000000000000000000000000000..1037f2e048ce9a102a2281e3e865b5d24fe60fe6 --- /dev/null +++ b/lib/pymaf/models/smpl.py @@ -0,0 +1,92 @@ +# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/smpl.py + +import torch +import numpy as np +from smplx import SMPL as _SMPL +from smplx.body_models import ModelOutput +from smplx.lbs import vertices2joints +from collections import namedtuple + +from lib.pymaf.core import path_config, constants + +SMPL_MEAN_PARAMS = path_config.SMPL_MEAN_PARAMS +SMPL_MODEL_DIR = path_config.SMPL_MODEL_DIR + +# Indices to get the 14 LSP joints from the 17 H36M joints +H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9] +H36M_TO_J14 = H36M_TO_J17[:14] + + +class SMPL(_SMPL): + """ Extension of the official SMPL implementation to support more joints """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES] + J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA) + self.register_buffer( + 'J_regressor_extra', + torch.tensor(J_regressor_extra, dtype=torch.float32)) + self.joint_map = torch.tensor(joints, dtype=torch.long) + self.ModelOutput = namedtuple( + 'ModelOutput_', ModelOutput._fields + ( + 'smpl_joints', + 'joints_J19', + )) + self.ModelOutput.__new__.__defaults__ = (None, ) * len( + self.ModelOutput._fields) + + def forward(self, *args, **kwargs): + kwargs['get_skin'] = True + smpl_output = super().forward(*args, **kwargs) + extra_joints = vertices2joints(self.J_regressor_extra, + smpl_output.vertices) + # smpl_output.joints: [B, 45, 3] extra_joints: [B, 9, 3] + vertices = smpl_output.vertices + joints = torch.cat([smpl_output.joints, extra_joints], dim=1) + smpl_joints = smpl_output.joints[:, :24] + joints = joints[:, self.joint_map, :] # [B, 49, 3] + joints_J24 = joints[:, -24:, :] + joints_J19 = joints_J24[:, constants.J24_TO_J19, :] + output = self.ModelOutput(vertices=vertices, + global_orient=smpl_output.global_orient, + body_pose=smpl_output.body_pose, + joints=joints, + joints_J19=joints_J19, + smpl_joints=smpl_joints, + betas=smpl_output.betas, + full_pose=smpl_output.full_pose) + return output + + +def get_smpl_faces(): + smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False) + return smpl.faces + + +def get_part_joints(smpl_joints): + batch_size = smpl_joints.shape[0] + + # part_joints = torch.zeros().to(smpl_joints.device) + + one_seg_pairs = [(0, 1), (0, 2), (0, 3), (3, 6), (9, 12), (9, 13), (9, 14), + (12, 15), (13, 16), (14, 17)] + two_seg_pairs = [(1, 4), (2, 5), (4, 7), (5, 8), (16, 18), (17, 19), + (18, 20), (19, 21)] + + one_seg_pairs.extend(two_seg_pairs) + + single_joints = [(10), (11), (15), (22), (23)] + + part_joints = [] + + for j_p in one_seg_pairs: + new_joint = torch.mean(smpl_joints[:, j_p], dim=1, keepdim=True) + part_joints.append(new_joint) + + for j_p in single_joints: + part_joints.append(smpl_joints[:, j_p:j_p + 1]) + + part_joints = torch.cat(part_joints, dim=1) + + return part_joints diff --git a/lib/pymaf/utils/__init__.py b/lib/pymaf/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/pymaf/utils/geometry.py b/lib/pymaf/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..51eb9bc0b958d220b02ec2ec643a6480d636f639 --- /dev/null +++ b/lib/pymaf/utils/geometry.py @@ -0,0 +1,452 @@ +import torch +import numpy as np +from torch.nn import functional as F +""" +Useful geometric operations, e.g. Perspective projection and a differentiable Rodrigues formula +Parts of the code are taken from https://github.com/MandyMo/pytorch_HMR +""" + + +def batch_rodrigues(theta): + """Convert axis-angle representation to rotation matrix. + Args: + theta: size = [B, 3] + Returns: + Rotation matrix corresponding to the quaternion -- size = [B, 3, 3] + """ + l1norm = torch.norm(theta + 1e-8, p=2, dim=1) + angle = torch.unsqueeze(l1norm, -1) + normalized = torch.div(theta, angle) + angle = angle * 0.5 + v_cos = torch.cos(angle) + v_sin = torch.sin(angle) + quat = torch.cat([v_cos, v_sin * normalized], dim=1) + return quat_to_rotmat(quat) + + +def quat_to_rotmat(quat): + """Convert quaternion coefficients to rotation matrix. + Args: + quat: size = [B, 4] 4 <===>(w, x, y, z) + Returns: + Rotation matrix corresponding to the quaternion -- size = [B, 3, 3] + """ + norm_quat = quat + norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True) + w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, + 2], norm_quat[:, + 3] + + B = quat.size(0) + + w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) + wx, wy, wz = w * x, w * y, w * z + xy, xz, yz = x * y, x * z, y * z + + rotMat = torch.stack([ + w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy, + w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz, + w2 - x2 - y2 + z2 + ], + dim=1).view(B, 3, 3) + return rotMat + + +def rotation_matrix_to_angle_axis(rotation_matrix): + """ + This function is borrowed from https://github.com/kornia/kornia + + Convert 3x4 rotation matrix to Rodrigues vector + + Args: + rotation_matrix (Tensor): rotation matrix. + + Returns: + Tensor: Rodrigues vector transformation. + + Shape: + - Input: :math:`(N, 3, 4)` + - Output: :math:`(N, 3)` + + Example: + >>> input = torch.rand(2, 3, 4) # Nx4x4 + >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3 + """ + if rotation_matrix.shape[1:] == (3, 3): + rot_mat = rotation_matrix.reshape(-1, 3, 3) + hom = torch.tensor([0, 0, 1], + dtype=torch.float32, + device=rotation_matrix.device).reshape( + 1, 3, 1).expand(rot_mat.shape[0], -1, -1) + rotation_matrix = torch.cat([rot_mat, hom], dim=-1) + + quaternion = rotation_matrix_to_quaternion(rotation_matrix) + aa = quaternion_to_angle_axis(quaternion) + aa[torch.isnan(aa)] = 0.0 + return aa + + +def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor: + """ + This function is borrowed from https://github.com/kornia/kornia + + Convert quaternion vector to angle axis of rotation. + + Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h + + Args: + quaternion (torch.Tensor): tensor with quaternions. + + Return: + torch.Tensor: tensor with angle axis of rotation. + + Shape: + - Input: :math:`(*, 4)` where `*` means, any number of dimensions + - Output: :math:`(*, 3)` + + Example: + >>> quaternion = torch.rand(2, 4) # Nx4 + >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3 + """ + if not torch.is_tensor(quaternion): + raise TypeError("Input type is not a torch.Tensor. Got {}".format( + type(quaternion))) + + if not quaternion.shape[-1] == 4: + raise ValueError( + "Input must be a tensor of shape Nx4 or 4. Got {}".format( + quaternion.shape)) + # unpack input and compute conversion + q1: torch.Tensor = quaternion[..., 1] + q2: torch.Tensor = quaternion[..., 2] + q3: torch.Tensor = quaternion[..., 3] + sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3 + + sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta) + cos_theta: torch.Tensor = quaternion[..., 0] + two_theta: torch.Tensor = 2.0 * torch.where( + cos_theta < 0.0, torch.atan2(-sin_theta, -cos_theta), + torch.atan2(sin_theta, cos_theta)) + + k_pos: torch.Tensor = two_theta / sin_theta + k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta) + k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg) + + angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3] + angle_axis[..., 0] += q1 * k + angle_axis[..., 1] += q2 * k + angle_axis[..., 2] += q3 * k + return angle_axis + + +def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6): + """ + This function is borrowed from https://github.com/kornia/kornia + + Convert 3x4 rotation matrix to 4d quaternion vector + + This algorithm is based on algorithm described in + https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201 + + Args: + rotation_matrix (Tensor): the rotation matrix to convert. + + Return: + Tensor: the rotation in quaternion + + Shape: + - Input: :math:`(N, 3, 4)` + - Output: :math:`(N, 4)` + + Example: + >>> input = torch.rand(4, 3, 4) # Nx3x4 + >>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4 + """ + if not torch.is_tensor(rotation_matrix): + raise TypeError("Input type is not a torch.Tensor. Got {}".format( + type(rotation_matrix))) + + if len(rotation_matrix.shape) > 3: + raise ValueError( + "Input size must be a three dimensional tensor. Got {}".format( + rotation_matrix.shape)) + if not rotation_matrix.shape[-2:] == (3, 4): + raise ValueError( + "Input size must be a N x 3 x 4 tensor. Got {}".format( + rotation_matrix.shape)) + + rmat_t = torch.transpose(rotation_matrix, 1, 2) + + mask_d2 = rmat_t[:, 2, 2] < eps + + mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1] + mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1] + + t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2] + q0 = torch.stack([ + rmat_t[:, 1, 2] - rmat_t[:, 2, 1], t0, + rmat_t[:, 0, 1] + rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2] + ], -1) + t0_rep = t0.repeat(4, 1).t() + + t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2] + q1 = torch.stack([ + rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] + rmat_t[:, 1, 0], + t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1] + ], -1) + t1_rep = t1.repeat(4, 1).t() + + t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2] + q2 = torch.stack([ + rmat_t[:, 0, 1] - rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2], + rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2 + ], -1) + t2_rep = t2.repeat(4, 1).t() + + t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2] + q3 = torch.stack([ + t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1], + rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] - rmat_t[:, 1, 0] + ], -1) + t3_rep = t3.repeat(4, 1).t() + + mask_c0 = mask_d2 * mask_d0_d1 + mask_c1 = mask_d2 * ~mask_d0_d1 + mask_c2 = ~mask_d2 * mask_d0_nd1 + mask_c3 = ~mask_d2 * ~mask_d0_nd1 + mask_c0 = mask_c0.view(-1, 1).type_as(q0) + mask_c1 = mask_c1.view(-1, 1).type_as(q1) + mask_c2 = mask_c2.view(-1, 1).type_as(q2) + mask_c3 = mask_c3.view(-1, 1).type_as(q3) + + q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3 + q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa + t2_rep * mask_c2 + t3_rep * mask_c3) # noqa + q *= 0.5 + return q + + +def rot6d_to_rotmat(x): + """Convert 6D rotation representation to 3x3 rotation matrix. + Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019 + Input: + (B,6) Batch of 6-D rotation representations + Output: + (B,3,3) Batch of corresponding rotation matrices + """ + x = x.view(-1, 3, 2) + a1 = x[:, :, 0] + a2 = x[:, :, 1] + b1 = F.normalize(a1) + b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) + b3 = torch.cross(b1, b2) + return torch.stack((b1, b2, b3), dim=-1) + + +def projection(pred_joints, pred_camera, retain_z=False): + pred_cam_t = torch.stack([ + pred_camera[:, 1], pred_camera[:, 2], 2 * 5000. / + (224. * pred_camera[:, 0] + 1e-9) + ], + dim=-1) + batch_size = pred_joints.shape[0] + camera_center = torch.zeros(batch_size, 2) + pred_keypoints_2d = perspective_projection( + pred_joints, + rotation=torch.eye(3).unsqueeze(0).expand(batch_size, -1, + -1).to(pred_joints.device), + translation=pred_cam_t, + focal_length=5000., + camera_center=camera_center, + retain_z=retain_z) + # Normalize keypoints to [-1,1] + pred_keypoints_2d = pred_keypoints_2d / (224. / 2.) + return pred_keypoints_2d + + +def perspective_projection(points, + rotation, + translation, + focal_length, + camera_center, + retain_z=False): + """ + This function computes the perspective projection of a set of points. + Input: + points (bs, N, 3): 3D points + rotation (bs, 3, 3): Camera rotation + translation (bs, 3): Camera translation + focal_length (bs,) or scalar: Focal length + camera_center (bs, 2): Camera center + """ + batch_size = points.shape[0] + K = torch.zeros([batch_size, 3, 3], device=points.device) + K[:, 0, 0] = focal_length + K[:, 1, 1] = focal_length + K[:, 2, 2] = 1. + K[:, :-1, -1] = camera_center + + # Transform points + points = torch.einsum('bij,bkj->bki', rotation, points) + points = points + translation.unsqueeze(1) + + # Apply perspective distortion + projected_points = points / points[:, :, -1].unsqueeze(-1) + + # Apply camera intrinsics + projected_points = torch.einsum('bij,bkj->bki', K, projected_points) + + if retain_z: + return projected_points + else: + return projected_points[:, :, :-1] + + +def estimate_translation_np(S, + joints_2d, + joints_conf, + focal_length=5000, + img_size=224): + """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. + Input: + S: (25, 3) 3D joint locations + joints: (25, 3) 2D joint locations and confidence + Returns: + (3,) camera translation vector + """ + + num_joints = S.shape[0] + # focal length + f = np.array([focal_length, focal_length]) + # optical center + center = np.array([img_size / 2., img_size / 2.]) + + # transformations + Z = np.reshape(np.tile(S[:, 2], (2, 1)).T, -1) + XY = np.reshape(S[:, 0:2], -1) + O = np.tile(center, num_joints) + F = np.tile(f, num_joints) + weight2 = np.reshape(np.tile(np.sqrt(joints_conf), (2, 1)).T, -1) + + # least squares + Q = np.array([ + F * np.tile(np.array([1, 0]), num_joints), + F * np.tile(np.array([0, 1]), num_joints), + O - np.reshape(joints_2d, -1) + ]).T + c = (np.reshape(joints_2d, -1) - O) * Z - F * XY + + # weighted least squares + W = np.diagflat(weight2) + Q = np.dot(W, Q) + c = np.dot(W, c) + + # square matrix + A = np.dot(Q.T, Q) + b = np.dot(Q.T, c) + + # solution + trans = np.linalg.solve(A, b) + + return trans + + +def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.): + """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. + Input: + S: (B, 49, 3) 3D joint locations + joints: (B, 49, 3) 2D joint locations and confidence + Returns: + (B, 3) camera translation vectors + """ + + device = S.device + # Use only joints 25:49 (GT joints) + S = S[:, 25:, :].cpu().numpy() + joints_2d = joints_2d[:, 25:, :].cpu().numpy() + joints_conf = joints_2d[:, :, -1] + joints_2d = joints_2d[:, :, :-1] + trans = np.zeros((S.shape[0], 3), dtype=np.float32) + # Find the translation for each example in the batch + for i in range(S.shape[0]): + S_i = S[i] + joints_i = joints_2d[i] + conf_i = joints_conf[i] + trans[i] = estimate_translation_np(S_i, + joints_i, + conf_i, + focal_length=focal_length, + img_size=img_size) + return torch.from_numpy(trans).to(device) + + +def Rot_y(angle, category='torch', prepend_dim=True, device=None): + '''Rotate around y-axis by angle + Args: + category: 'torch' or 'numpy' + prepend_dim: prepend an extra dimension + Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) + ''' + m = np.array([[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.], + [-np.sin(angle), 0., np.cos(angle)]]) + if category == 'torch': + if prepend_dim: + return torch.tensor(m, dtype=torch.float, + device=device).unsqueeze(0) + else: + return torch.tensor(m, dtype=torch.float, device=device) + elif category == 'numpy': + if prepend_dim: + return np.expand_dims(m, 0) + else: + return m + else: + raise ValueError("category must be 'torch' or 'numpy'") + + +def Rot_x(angle, category='torch', prepend_dim=True, device=None): + '''Rotate around x-axis by angle + Args: + category: 'torch' or 'numpy' + prepend_dim: prepend an extra dimension + Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) + ''' + m = np.array([[1., 0., 0.], [0., np.cos(angle), -np.sin(angle)], + [0., np.sin(angle), np.cos(angle)]]) + if category == 'torch': + if prepend_dim: + return torch.tensor(m, dtype=torch.float, + device=device).unsqueeze(0) + else: + return torch.tensor(m, dtype=torch.float, device=device) + elif category == 'numpy': + if prepend_dim: + return np.expand_dims(m, 0) + else: + return m + else: + raise ValueError("category must be 'torch' or 'numpy'") + + +def Rot_z(angle, category='torch', prepend_dim=True, device=None): + '''Rotate around z-axis by angle + Args: + category: 'torch' or 'numpy' + prepend_dim: prepend an extra dimension + Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) + ''' + m = np.array([[np.cos(angle), -np.sin(angle), 0.], + [np.sin(angle), np.cos(angle), 0.], [0., 0., 1.]]) + if category == 'torch': + if prepend_dim: + return torch.tensor(m, dtype=torch.float, + device=device).unsqueeze(0) + else: + return torch.tensor(m, dtype=torch.float, device=device) + elif category == 'numpy': + if prepend_dim: + return np.expand_dims(m, 0) + else: + return m + else: + raise ValueError("category must be 'torch' or 'numpy'") diff --git a/lib/pymaf/utils/imutils.py b/lib/pymaf/utils/imutils.py new file mode 100644 index 0000000000000000000000000000000000000000..268822521a49dd3d4913f92b9f14667499354932 --- /dev/null +++ b/lib/pymaf/utils/imutils.py @@ -0,0 +1,497 @@ +""" +This file contains functions that are used to perform data augmentation. +""" +from turtle import reset +import cv2 +import io +import torch +import numpy as np +import scipy.misc +from PIL import Image +from rembg.bg import remove +import human_det + +from lib.pymaf.core import constants +from lib.pymaf.utils.streamer import aug_matrix +from lib.common.cloth_extraction import load_segmentation +from torchvision import transforms + + +def load_img(img_file): + + img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED) + if len(img.shape) == 2: + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + if not img_file.endswith("png"): + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + else: + img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) + + return img + + +def get_bbox(img, det): + + input = np.float32(img) + input = (input / 255.0 - + (0.5, 0.5, 0.5)) / (0.5, 0.5, 0.5) # TO [-1.0, 1.0] + input = input.transpose(2, 0, 1) # TO [3 x H x W] + bboxes, probs = det(torch.from_numpy(input).float().unsqueeze(0)) + + probs = probs.unsqueeze(3) + bboxes = (bboxes * probs).sum(dim=1, keepdim=True) / probs.sum( + dim=1, keepdim=True) + bbox = bboxes[0, 0, 0].cpu().numpy() + + return bbox + + +def get_transformer(input_res): + + image_to_tensor = transforms.Compose([ + transforms.Resize(input_res), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + + mask_to_tensor = transforms.Compose([ + transforms.Resize(input_res), + transforms.ToTensor(), + transforms.Normalize((0.0, ), (1.0, )) + ]) + + image_to_pymaf_tensor = transforms.Compose([ + transforms.Resize(size=224), + transforms.Normalize(mean=constants.IMG_NORM_MEAN, + std=constants.IMG_NORM_STD) + ]) + + image_to_pixie_tensor = transforms.Compose([ + transforms.Resize(224) + ]) + + def image_to_hybrik_tensor(img): + # mean + img[0].add_(-0.406) + img[1].add_(-0.457) + img[2].add_(-0.480) + + # std + img[0].div_(0.225) + img[1].div_(0.224) + img[2].div_(0.229) + return img + + return [image_to_tensor, mask_to_tensor, image_to_pymaf_tensor, image_to_pixie_tensor, image_to_hybrik_tensor] + + +def process_image(img_file, det, hps_type, input_res=512, device=None, seg_path=None): + """Read image, do preprocessing and possibly crop it according to the bounding box. + If there are bounding box annotations, use them to crop the image. + If no bounding box is specified but openpose detections are available, use them to get the bounding box. + """ + + [image_to_tensor, mask_to_tensor, image_to_pymaf_tensor, + image_to_pixie_tensor, image_to_hybrik_tensor] = get_transformer(input_res) + + img_ori = load_img(img_file) + + in_height, in_width, _ = img_ori.shape + M = aug_matrix(in_width, in_height, input_res*2, input_res*2) + + # from rectangle to square + img_for_crop = cv2.warpAffine(img_ori, M[0:2, :], + (input_res*2, input_res*2), flags=cv2.INTER_CUBIC) + + if det is not None: + + # detection for bbox + bbox = get_bbox(img_for_crop, det) + + width = bbox[2] - bbox[0] + height = bbox[3] - bbox[1] + center = np.array([(bbox[0] + bbox[2]) / 2.0, + (bbox[1] + bbox[3]) / 2.0]) + + else: + # Assume that the person is centerered in the image + height = img_for_crop.shape[0] + width = img_for_crop.shape[1] + center = np.array([width // 2, height // 2]) + + scale = max(height, width) / 180 + + if hps_type == 'hybrik': + img_np = crop_for_hybrik(img_for_crop, center, + np.array([scale * 180, scale * 180])) + else: + img_np, cropping_parameters = crop( + img_for_crop, center, scale, (input_res, input_res)) + + with torch.no_grad(): + buf = io.BytesIO() + Image.fromarray(img_np).save(buf, format='png') + img_pil = Image.open( + io.BytesIO(remove(buf.getvalue()))).convert("RGBA") + + # for icon + img_rgb = image_to_tensor(img_pil.convert("RGB")) + img_mask = torch.tensor(1.0) - (mask_to_tensor(img_pil.split()[-1]) < + torch.tensor(0.5)).float() + img_tensor = img_rgb * img_mask + + # for hps + img_hps = img_np.astype(np.float32) / 255. + img_hps = torch.from_numpy(img_hps).permute(2, 0, 1) + + if hps_type == 'bev': + img_hps = img_np[:, :, [2, 1, 0]] + elif hps_type == 'hybrik': + img_hps = image_to_hybrik_tensor(img_hps).unsqueeze(0).to(device) + elif hps_type != 'pixie': + img_hps = image_to_pymaf_tensor(img_hps).unsqueeze(0).to(device) + else: + img_hps = image_to_pixie_tensor(img_hps).unsqueeze(0).to(device) + + # uncrop params + uncrop_param = {'center': center, + 'scale': scale, + 'ori_shape': img_ori.shape, + 'box_shape': img_np.shape, + 'crop_shape': img_for_crop.shape, + 'M': M} + + if not (seg_path is None): + segmentations = load_segmentation(seg_path, (in_height, in_width)) + seg_coord_normalized = [] + for seg in segmentations: + coord_normalized = [] + for xy in seg['coordinates']: + xy_h = np.vstack((xy[:, 0], xy[:, 1], np.ones(len(xy)))).T + warped_indeces = M[0:2, :] @ xy_h[:, :, None] + warped_indeces = np.array(warped_indeces).astype(int) + warped_indeces.resize((warped_indeces.shape[:2])) + + # cropped_indeces = crop_segmentation(warped_indeces, center, scale, (input_res, input_res), img_np.shape) + cropped_indeces = crop_segmentation( + warped_indeces, (input_res, input_res), cropping_parameters) + + indices = np.vstack( + (cropped_indeces[:, 0], cropped_indeces[:, 1])).T + + # Convert to NDC coordinates + seg_cropped_normalized = 2*(indices / input_res) - 1 + # Don't know why we need to divide by 50 but it works ¯\_(ツ)_/¯ (probably some scaling factor somewhere) + # Divide only by 45 on the horizontal axis to take the curve of the human body into account + seg_cropped_normalized[:, 0] = ( + 1/40) * seg_cropped_normalized[:, 0] + seg_cropped_normalized[:, 1] = ( + 1/50) * seg_cropped_normalized[:, 1] + coord_normalized.append(seg_cropped_normalized) + + seg['coord_normalized'] = coord_normalized + seg_coord_normalized.append(seg) + + return img_tensor, img_hps, img_ori, img_mask, uncrop_param, seg_coord_normalized + + return img_tensor, img_hps, img_ori, img_mask, uncrop_param + + +def get_transform(center, scale, res): + """Generate transformation matrix.""" + h = 200 * scale + t = np.zeros((3, 3)) + t[0, 0] = float(res[1]) / h + t[1, 1] = float(res[0]) / h + t[0, 2] = res[1] * (-float(center[0]) / h + .5) + t[1, 2] = res[0] * (-float(center[1]) / h + .5) + t[2, 2] = 1 + + return t + + +def transform(pt, center, scale, res, invert=0): + """Transform pixel location to different reference.""" + t = get_transform(center, scale, res) + if invert: + t = np.linalg.inv(t) + new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.]).T + new_pt = np.dot(t, new_pt) + return np.around(new_pt[:2]).astype(np.int16) + + +def crop(img, center, scale, res): + """Crop image according to the supplied bounding box.""" + + # Upper left point + ul = np.array(transform([0, 0], center, scale, res, invert=1)) + + # Bottom right point + br = np.array(transform(res, center, scale, res, invert=1)) + + new_shape = [br[1] - ul[1], br[0] - ul[0]] + if len(img.shape) > 2: + new_shape += [img.shape[2]] + new_img = np.zeros(new_shape) + + # Range to fill new array + new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0] + new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1] + + # Range to sample from original image + old_x = max(0, ul[0]), min(len(img[0]), br[0]) + old_y = max(0, ul[1]), min(len(img), br[1]) + + new_img[new_y[0]:new_y[1], new_x[0]:new_x[1] + ] = img[old_y[0]:old_y[1], old_x[0]:old_x[1]] + if len(img.shape) == 2: + new_img = np.array(Image.fromarray(new_img).resize(res)) + else: + new_img = np.array(Image.fromarray( + new_img.astype(np.uint8)).resize(res)) + + return new_img, (old_x, new_x, old_y, new_y, new_shape) + + +def crop_segmentation(org_coord, res, cropping_parameters): + old_x, new_x, old_y, new_y, new_shape = cropping_parameters + + new_coord = np.zeros((org_coord.shape)) + new_coord[:, 0] = new_x[0] + (org_coord[:, 0] - old_x[0]) + new_coord[:, 1] = new_y[0] + (org_coord[:, 1] - old_y[0]) + + new_coord[:, 0] = res[0] * (new_coord[:, 0] / new_shape[1]) + new_coord[:, 1] = res[1] * (new_coord[:, 1] / new_shape[0]) + + return new_coord + + +def crop_for_hybrik(img, center, scale): + inp_h, inp_w = (256, 256) + trans = get_affine_transform(center, scale, 0, [inp_w, inp_h]) + new_img = cv2.warpAffine( + img, trans, (int(inp_w), int(inp_h)), flags=cv2.INTER_LINEAR) + return new_img + + +def get_affine_transform(center, + scale, + rot, + output_size, + shift=np.array([0, 0], dtype=np.float32), + inv=0): + + def get_dir(src_point, rot_rad): + """Rotate the point by `rot_rad` degree.""" + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + + src_result = [0, 0] + src_result[0] = src_point[0] * cs - src_point[1] * sn + src_result[1] = src_point[0] * sn + src_point[1] * cs + + return src_result + + def get_3rd_point(a, b): + """Return vector c that perpendicular to (a - b).""" + direct = a - b + return b + np.array([-direct[1], direct[0]], dtype=np.float32) + + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): + scale = np.array([scale, scale]) + + scale_tmp = scale + src_w = scale_tmp[0] + dst_w = output_size[0] + dst_h = output_size[1] + + rot_rad = np.pi * rot / 180 + src_dir = get_dir([0, src_w * -0.5], rot_rad) + dst_dir = np.array([0, dst_w * -0.5], np.float32) + + src = np.zeros((3, 2), dtype=np.float32) + dst = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale_tmp * shift + src[1, :] = center + src_dir + scale_tmp * shift + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir + + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + +def corner_align(ul, br): + + if ul[1]-ul[0] != br[1]-br[0]: + ul[1] = ul[0]+br[1]-br[0] + + return ul, br + + +def uncrop(img, center, scale, orig_shape): + """'Undo' the image cropping/resizing. + This function is used when evaluating mask/part segmentation. + """ + + res = img.shape[:2] + + # Upper left point + ul = np.array(transform([0, 0], center, scale, res, invert=1)) + # Bottom right point + br = np.array(transform(res, center, scale, res, invert=1)) + + # quick fix + ul, br = corner_align(ul, br) + + # size of cropped image + crop_shape = [br[1] - ul[1], br[0] - ul[0]] + new_img = np.zeros(orig_shape, dtype=np.uint8) + + # Range to fill new array + new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0] + new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1] + + # Range to sample from original image + old_x = max(0, ul[0]), min(orig_shape[1], br[0]) + old_y = max(0, ul[1]), min(orig_shape[0], br[1]) + + img = np.array(Image.fromarray(img.astype(np.uint8)).resize(crop_shape)) + + new_img[old_y[0]:old_y[1], old_x[0]:old_x[1] + ] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]] + + return new_img + + +def rot_aa(aa, rot): + """Rotate axis angle parameters.""" + # pose parameters + R = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], + [np.sin(np.deg2rad(-rot)), + np.cos(np.deg2rad(-rot)), 0], [0, 0, 1]]) + # find the rotation of the body in camera frame + per_rdg, _ = cv2.Rodrigues(aa) + # apply the global rotation to the global orientation + resrot, _ = cv2.Rodrigues(np.dot(R, per_rdg)) + aa = (resrot.T)[0] + return aa + + +def flip_img(img): + """Flip rgb images or masks. + channels come last, e.g. (256,256,3). + """ + img = np.fliplr(img) + return img + + +def flip_kp(kp, is_smpl=False): + """Flip keypoints.""" + if len(kp) == 24: + if is_smpl: + flipped_parts = constants.SMPL_JOINTS_FLIP_PERM + else: + flipped_parts = constants.J24_FLIP_PERM + elif len(kp) == 49: + if is_smpl: + flipped_parts = constants.SMPL_J49_FLIP_PERM + else: + flipped_parts = constants.J49_FLIP_PERM + kp = kp[flipped_parts] + kp[:, 0] = -kp[:, 0] + return kp + + +def flip_pose(pose): + """Flip pose. + The flipping is based on SMPL parameters. + """ + flipped_parts = constants.SMPL_POSE_FLIP_PERM + pose = pose[flipped_parts] + # we also negate the second and the third dimension of the axis-angle + pose[1::3] = -pose[1::3] + pose[2::3] = -pose[2::3] + return pose + + +def normalize_2d_kp(kp_2d, crop_size=224, inv=False): + # Normalize keypoints between -1, 1 + if not inv: + ratio = 1.0 / crop_size + kp_2d = 2.0 * kp_2d * ratio - 1.0 + else: + ratio = 1.0 / crop_size + kp_2d = (kp_2d + 1.0) / (2 * ratio) + + return kp_2d + + +def generate_heatmap(joints, heatmap_size, sigma=1, joints_vis=None): + ''' + param joints: [num_joints, 3] + param joints_vis: [num_joints, 3] + return: target, target_weight(1: visible, 0: invisible) + ''' + num_joints = joints.shape[0] + device = joints.device + cur_device = torch.device(device.type, device.index) + if not hasattr(heatmap_size, '__len__'): + # width height + heatmap_size = [heatmap_size, heatmap_size] + assert len(heatmap_size) == 2 + target_weight = np.ones((num_joints, 1), dtype=np.float32) + if joints_vis is not None: + target_weight[:, 0] = joints_vis[:, 0] + target = torch.zeros((num_joints, heatmap_size[1], heatmap_size[0]), + dtype=torch.float32, + device=cur_device) + + tmp_size = sigma * 3 + + for joint_id in range(num_joints): + mu_x = int(joints[joint_id][0] * heatmap_size[0] + 0.5) + mu_y = int(joints[joint_id][1] * heatmap_size[1] + 0.5) + # Check that any part of the gaussian is in-bounds + ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)] + br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)] + if ul[0] >= heatmap_size[0] or ul[1] >= heatmap_size[1] \ + or br[0] < 0 or br[1] < 0: + # If not, just return the image as is + target_weight[joint_id] = 0 + continue + + # # Generate gaussian + size = 2 * tmp_size + 1 + # x = np.arange(0, size, 1, np.float32) + # y = x[:, np.newaxis] + # x0 = y0 = size // 2 + # # The gaussian is not normalized, we want the center value to equal 1 + # g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2)) + # g = torch.from_numpy(g.astype(np.float32)) + + x = torch.arange(0, size, dtype=torch.float32, device=cur_device) + y = x.unsqueeze(-1) + x0 = y0 = size // 2 + # The gaussian is not normalized, we want the center value to equal 1 + g = torch.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2)) + + # Usable gaussian range + g_x = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0] + g_y = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1] + # Image range + img_x = max(0, ul[0]), min(br[0], heatmap_size[0]) + img_y = max(0, ul[1]), min(br[1], heatmap_size[1]) + + v = target_weight[joint_id] + if v > 0.5: + target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \ + g[g_y[0]:g_y[1], g_x[0]:g_x[1]] + + return target, target_weight diff --git a/lib/pymaf/utils/streamer.py b/lib/pymaf/utils/streamer.py new file mode 100644 index 0000000000000000000000000000000000000000..1753677159f9550dc26c8b40c04b3713f90b959b --- /dev/null +++ b/lib/pymaf/utils/streamer.py @@ -0,0 +1,142 @@ +import cv2 +import torch +import numpy as np +import imageio + + +def aug_matrix(w1, h1, w2, h2): + dx = (w2 - w1) / 2.0 + dy = (h2 - h1) / 2.0 + + matrix_trans = np.array([[1.0, 0, dx], + [0, 1.0, dy], + [0, 0, 1.0]]) + + scale = np.min([float(w2)/w1, float(h2)/h1]) + + M = get_affine_matrix( + center=(w2 / 2.0, h2 / 2.0), + translate=(0, 0), + scale=scale) + + M = np.array(M + [0., 0., 1.]).reshape(3, 3) + M = M.dot(matrix_trans) + + return M + + +def get_affine_matrix(center, translate, scale): + cx, cy = center + tx, ty = translate + + M = [1, 0, 0, + 0, 1, 0] + M = [x * scale for x in M] + + # Apply translation and of center translation: RSS * C^-1 + M[2] += M[0] * (-cx) + M[1] * (-cy) + M[5] += M[3] * (-cx) + M[4] * (-cy) + + # Apply center translation: T * C * RSS * C^-1 + M[2] += cx + tx + M[5] += cy + ty + return M + + +class BaseStreamer(): + """This streamer will return images at 512x512 size. + """ + + def __init__(self, + width=512, height=512, pad=True, + mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), + **kwargs): + self.width = width + self.height = height + self.pad = pad + self.mean = np.array(mean) + self.std = np.array(std) + + self.loader = self.create_loader() + + def create_loader(self): + raise NotImplementedError + yield np.zeros((600, 400, 3)) # in RGB (0, 255) + + def __getitem__(self, index): + image = next(self.loader) + in_height, in_width, _ = image.shape + M = aug_matrix(in_width, in_height, self.width, self.height, self.pad) + image = cv2.warpAffine( + image, M[0:2, :], (self.width, self.height), flags=cv2.INTER_CUBIC) + + input = np.float32(image) + input = (input / 255.0 - self.mean) / self.std # TO [-1.0, 1.0] + input = input.transpose(2, 0, 1) # TO [3 x H x W] + return torch.from_numpy(input).float() + + def __len__(self): + raise NotImplementedError + + +class CaptureStreamer(BaseStreamer): + """This streamer takes webcam as input. + """ + + def __init__(self, id=0, width=512, height=512, pad=True, **kwargs): + super().__init__(width, height, pad, **kwargs) + self.capture = cv2.VideoCapture(id) + + def create_loader(self): + while True: + _, image = self.capture.read() + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # RGB + yield image + + def __len__(self): + return 100_000_000 + + def __del__(self): + self.capture.release() + + +class VideoListStreamer(BaseStreamer): + """This streamer takes a list of video files as input. + """ + + def __init__(self, files, width=512, height=512, pad=True, **kwargs): + super().__init__(width, height, pad, **kwargs) + self.files = files + self.captures = [imageio.get_reader(f) for f in files] + self.nframes = sum([int(cap._meta["fps"] * cap._meta["duration"]) + for cap in self.captures]) + + def create_loader(self): + for capture in self.captures: + for image in capture: # RGB + yield image + + def __len__(self): + return self.nframes + + def __del__(self): + for capture in self.captures: + capture.close() + + +class ImageListStreamer(BaseStreamer): + """This streamer takes a list of image files as input. + """ + + def __init__(self, files, width=512, height=512, pad=True, **kwargs): + super().__init__(width, height, pad, **kwargs) + self.files = files + + def create_loader(self): + for f in self.files: + image = cv2.imread(f, cv2.IMREAD_UNCHANGED)[:, :, 0:3] + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # RGB + yield image + + def __len__(self): + return len(self.files) diff --git a/lib/pymaf/utils/transforms.py b/lib/pymaf/utils/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..996742aa6defb02afcd59f7f2ae0eb47576bf314 --- /dev/null +++ b/lib/pymaf/utils/transforms.py @@ -0,0 +1,78 @@ +# ------------------------------------------------------------------------------ +# Copyright (c) Microsoft +# Licensed under the MIT License. +# Written by Bin Xiao (Bin.Xiao@microsoft.com) +# ------------------------------------------------------------------------------ + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import cv2 +import numpy as np + + +def transform_preds(coords, center, scale, output_size): + target_coords = np.zeros(coords.shape) + trans = get_affine_transform(center, scale, 0, output_size, inv=1) + for p in range(coords.shape[0]): + target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans) + return target_coords + + +def get_affine_transform(center, + scale, + rot, + output_size, + shift=np.array([0, 0], dtype=np.float32), + inv=0): + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): + # print(scale) + scale = np.array([scale, scale]) + + scale_tmp = scale * 200.0 + src_w = scale_tmp[0] + dst_w = output_size[0] + dst_h = output_size[1] + + rot_rad = np.pi * rot / 180 + src_dir = get_dir([0, src_w * -0.5], rot_rad) + dst_dir = np.array([0, dst_w * -0.5], np.float32) + + src = np.zeros((3, 2), dtype=np.float32) + dst = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale_tmp * shift + src[1, :] = center + src_dir + scale_tmp * shift + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir + + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + +def affine_transform(pt, t): + new_pt = np.array([pt[0], pt[1], 1.]).T + new_pt = np.dot(t, new_pt) + return new_pt[:2] + + +def get_3rd_point(a, b): + direct = a - b + return b + np.array([-direct[1], direct[0]], dtype=np.float32) + + +def get_dir(src_point, rot_rad): + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + + src_result = [0, 0] + src_result[0] = src_point[0] * cs - src_point[1] * sn + src_result[1] = src_point[0] * sn + src_point[1] * cs + + return src_result diff --git a/lib/renderer/__init__.py b/lib/renderer/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/renderer/camera.py b/lib/renderer/camera.py new file mode 100755 index 0000000000000000000000000000000000000000..e6f4ce8c4c08becbff935fabcf5a1ffb65e4443c --- /dev/null +++ b/lib/renderer/camera.py @@ -0,0 +1,226 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import cv2 +import numpy as np + +from .glm import ortho + + +class Camera: + def __init__(self, width=1600, height=1200): + # Focal Length + # equivalent 50mm + focal = np.sqrt(width * width + height * height) + self.focal_x = focal + self.focal_y = focal + # Principal Point Offset + self.principal_x = width / 2 + self.principal_y = height / 2 + # Axis Skew + self.skew = 0 + # Image Size + self.width = width + self.height = height + + self.near = 1 + self.far = 10 + + # Camera Center + self.center = np.array([0, 0, 1.6]) + self.direction = np.array([0, 0, -1]) + self.right = np.array([1, 0, 0]) + self.up = np.array([0, 1, 0]) + + self.ortho_ratio = None + + def sanity_check(self): + self.center = self.center.reshape([-1]) + self.direction = self.direction.reshape([-1]) + self.right = self.right.reshape([-1]) + self.up = self.up.reshape([-1]) + + assert len(self.center) == 3 + assert len(self.direction) == 3 + assert len(self.right) == 3 + assert len(self.up) == 3 + + @staticmethod + def normalize_vector(v): + v_norm = np.linalg.norm(v) + return v if v_norm == 0 else v / v_norm + + def get_real_z_value(self, z): + z_near = self.near + z_far = self.far + z_n = 2.0 * z - 1.0 + z_e = 2.0 * z_near * z_far / (z_far + z_near - z_n * (z_far - z_near)) + return z_e + + def get_rotation_matrix(self): + rot_mat = np.eye(3) + s = self.right + s = self.normalize_vector(s) + rot_mat[0, :] = s + u = self.up + u = self.normalize_vector(u) + rot_mat[1, :] = -u + rot_mat[2, :] = self.normalize_vector(self.direction) + + return rot_mat + + def get_translation_vector(self): + rot_mat = self.get_rotation_matrix() + trans = -np.dot(rot_mat, self.center) + return trans + + def get_intrinsic_matrix(self): + int_mat = np.eye(3) + + int_mat[0, 0] = self.focal_x + int_mat[1, 1] = self.focal_y + int_mat[0, 1] = self.skew + int_mat[0, 2] = self.principal_x + int_mat[1, 2] = self.principal_y + + return int_mat + + def get_projection_matrix(self): + ext_mat = self.get_extrinsic_matrix() + int_mat = self.get_intrinsic_matrix() + + return np.matmul(int_mat, ext_mat) + + def get_extrinsic_matrix(self): + rot_mat = self.get_rotation_matrix() + int_mat = self.get_intrinsic_matrix() + trans = self.get_translation_vector() + + extrinsic = np.eye(4) + extrinsic[:3, :3] = rot_mat + extrinsic[:3, 3] = trans + + return extrinsic[:3, :] + + def set_rotation_matrix(self, rot_mat): + self.direction = rot_mat[2, :] + self.up = -rot_mat[1, :] + self.right = rot_mat[0, :] + + def set_intrinsic_matrix(self, int_mat): + self.focal_x = int_mat[0, 0] + self.focal_y = int_mat[1, 1] + self.skew = int_mat[0, 1] + self.principal_x = int_mat[0, 2] + self.principal_y = int_mat[1, 2] + + def set_projection_matrix(self, proj_mat): + res = cv2.decomposeProjectionMatrix(proj_mat) + int_mat, rot_mat, camera_center_homo = res[0], res[1], res[2] + camera_center = camera_center_homo[0:3] / camera_center_homo[3] + camera_center = camera_center.reshape(-1) + int_mat = int_mat / int_mat[2][2] + + self.set_intrinsic_matrix(int_mat) + self.set_rotation_matrix(rot_mat) + self.center = camera_center + + self.sanity_check() + + def get_gl_matrix(self): + z_near = self.near + z_far = self.far + rot_mat = self.get_rotation_matrix() + int_mat = self.get_intrinsic_matrix() + trans = self.get_translation_vector() + + extrinsic = np.eye(4) + extrinsic[:3, :3] = rot_mat + extrinsic[:3, 3] = trans + axis_adj = np.eye(4) + axis_adj[2, 2] = -1 + axis_adj[1, 1] = -1 + model_view = np.matmul(axis_adj, extrinsic) + + projective = np.zeros([4, 4]) + projective[:2, :2] = int_mat[:2, :2] + projective[:2, 2:3] = -int_mat[:2, 2:3] + projective[3, 2] = -1 + projective[2, 2] = (z_near + z_far) + projective[2, 3] = (z_near * z_far) + + if self.ortho_ratio is None: + ndc = ortho(0, self.width, 0, self.height, z_near, z_far) + perspective = np.matmul(ndc, projective) + else: + perspective = ortho(-self.width * self.ortho_ratio / 2, + self.width * self.ortho_ratio / 2, + -self.height * self.ortho_ratio / 2, + self.height * self.ortho_ratio / 2, z_near, + z_far) + + return perspective, model_view + + +def KRT_from_P(proj_mat, normalize_K=True): + res = cv2.decomposeProjectionMatrix(proj_mat) + K, Rot, camera_center_homog = res[0], res[1], res[2] + camera_center = camera_center_homog[0:3] / camera_center_homog[3] + trans = -Rot.dot(camera_center) + if normalize_K: + K = K / K[2][2] + return K, Rot, trans + + +def MVP_from_P(proj_mat, width, height, near=0.1, far=10000): + ''' + Convert OpenCV camera calibration matrix to OpenGL projection and model view matrix + :param proj_mat: OpenCV camera projeciton matrix + :param width: Image width + :param height: Image height + :param near: Z near value + :param far: Z far value + :return: OpenGL projection matrix and model view matrix + ''' + res = cv2.decomposeProjectionMatrix(proj_mat) + K, Rot, camera_center_homog = res[0], res[1], res[2] + camera_center = camera_center_homog[0:3] / camera_center_homog[3] + trans = -Rot.dot(camera_center) + K = K / K[2][2] + + extrinsic = np.eye(4) + extrinsic[:3, :3] = Rot + extrinsic[:3, 3:4] = trans + axis_adj = np.eye(4) + axis_adj[2, 2] = -1 + axis_adj[1, 1] = -1 + model_view = np.matmul(axis_adj, extrinsic) + + zFar = far + zNear = near + projective = np.zeros([4, 4]) + projective[:2, :2] = K[:2, :2] + projective[:2, 2:3] = -K[:2, 2:3] + projective[3, 2] = -1 + projective[2, 2] = (zNear + zFar) + projective[2, 3] = (zNear * zFar) + + ndc = ortho(0, width, 0, height, zNear, zFar) + + perspective = np.matmul(ndc, projective) + + return perspective, model_view diff --git a/lib/renderer/gl/__init__.py b/lib/renderer/gl/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/renderer/gl/cam_render.py b/lib/renderer/gl/cam_render.py new file mode 100755 index 0000000000000000000000000000000000000000..23488a946733f5cee953bf79032b1cd24c5473e1 --- /dev/null +++ b/lib/renderer/gl/cam_render.py @@ -0,0 +1,80 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from .render import Render + +GLUT = None + + +class CamRender(Render): + def __init__(self, + width=1600, + height=1200, + name='Cam Renderer', + program_files=['simple.fs', 'simple.vs'], + color_size=1, + ms_rate=1, + egl=False): + Render.__init__(self, + width, + height, + name, + program_files, + color_size, + ms_rate=ms_rate, + egl=egl) + self.camera = None + + if not egl: + global GLUT + import OpenGL.GLUT as GLUT + GLUT.glutDisplayFunc(self.display) + GLUT.glutKeyboardFunc(self.keyboard) + + def set_camera(self, camera): + self.camera = camera + self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() + + def keyboard(self, key, x, y): + # up + eps = 1 + # print(key) + if key == b'w': + self.camera.center += eps * self.camera.direction + elif key == b's': + self.camera.center -= eps * self.camera.direction + if key == b'a': + self.camera.center -= eps * self.camera.right + elif key == b'd': + self.camera.center += eps * self.camera.right + if key == b' ': + self.camera.center += eps * self.camera.up + elif key == b'x': + self.camera.center -= eps * self.camera.up + elif key == b'i': + self.camera.near += 0.1 * eps + self.camera.far += 0.1 * eps + elif key == b'o': + self.camera.near -= 0.1 * eps + self.camera.far -= 0.1 * eps + + self.projection_matrix, self.model_view_matrix = self.camera.get_gl_matrix( + ) + + def show(self): + if GLUT is not None: + GLUT.glutMainLoop() diff --git a/lib/renderer/gl/color_render.py b/lib/renderer/gl/color_render.py new file mode 100644 index 0000000000000000000000000000000000000000..4b6a7456bdb3d9ac8b6f5b9d746c3b70ad56d8cc --- /dev/null +++ b/lib/renderer/gl/color_render.py @@ -0,0 +1,162 @@ +''' +MIT License + +Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +''' +import numpy as np +import random + +from .framework import * +from .cam_render import CamRender + + +class ColorRender(CamRender): + def __init__(self, width=1600, height=1200, name='Color Renderer', egl=False): + program_files = ['color.vs', 'color.fs'] + CamRender.__init__(self, + width, + height, + name, + program_files=program_files, + color_size=3, egl=egl) + + # WARNING: this differs from vertex_buffer and vertex_data in Render + self.vert_buffer = {} + self.vert_data = {} + + # normal + self.norm_buffer = {} + self.norm_data = {} + + self.color_buffer = {} + self.color_data = {} + + self.vertex_dim = {} + self.n_vertices = {} + + self.rot_mat_unif = glGetUniformLocation(self.program, 'RotMat') + self.rot_matrix = np.eye(3) + + self.norm_mat_unif = glGetUniformLocation(self.program, 'NormMat') + self.normalize_matrix = np.eye(4) + + def set_norm_mat(self, scale, center): + N = np.eye(4) + N[:3, :3] = scale * np.eye(3) + N[:3, 3] = -scale * center + + self.normalize_matrix = N + + def set_mesh(self, vertices, faces, color, normals, mat_name='all'): + + self.vert_data[mat_name] = vertices[faces.reshape([-1])] + self.n_vertices[mat_name] = self.vert_data[mat_name].shape[0] + self.vertex_dim[mat_name] = self.vert_data[mat_name].shape[1] + self.color_data[mat_name] = color[faces.reshape([-1])] + self.norm_data[mat_name] = normals[faces.reshape([-1])] + + if mat_name not in self.vert_buffer.keys(): + self.vert_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.vert_data[mat_name], GL_STATIC_DRAW) + + if mat_name not in self.color_buffer.keys(): + self.color_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.color_data[mat_name], + GL_STATIC_DRAW) + + if mat_name not in self.norm_buffer.keys(): + self.norm_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.norm_data[mat_name], GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def cleanup(self): + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + for key in self.vert_data: + glDeleteBuffers(1, [self.vert_buffer[key]]) + glDeleteBuffers(1, [self.color_buffer[key]]) + glDeleteBuffers(1, [self.norm_buffer[key]]) + + self.norm_buffer = {} + self.norm_data = {} + + self.vert_buffer = {} + self.vert_data = {} + + self.color_buffer = {} + self.color_data = {} + + self.render_texture_mat = {} + + self.vertex_dim = {} + self.n_vertices = {} + + def draw(self): + self.draw_init() + + glEnable(GL_MULTISAMPLE) + + glUseProgram(self.program) + glUniformMatrix4fv(self.norm_mat_unif, 1, GL_FALSE, + self.normalize_matrix.transpose()) + glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, + self.model_view_matrix.transpose()) + glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, + self.projection_matrix.transpose()) + glUniformMatrix3fv(self.rot_mat_unif, 1, GL_FALSE, + self.rot_matrix.transpose()) + + for mat in self.vert_buffer: + + # Handle vertex buffer + glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat]) + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, self.vertex_dim[mat], GL_DOUBLE, GL_FALSE, + 0, None) + + # Handle color buffer + glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer[mat]) + glEnableVertexAttribArray(1) + glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle normal buffer + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat]) + glEnableVertexAttribArray(2) + glVertexAttribPointer(2, 3, GL_DOUBLE, GL_FALSE, 0, None) + + glDrawArrays(GL_TRIANGLES, 0, self.n_vertices[mat]) + + glDisableVertexAttribArray(2) + glDisableVertexAttribArray(1) + glDisableVertexAttribArray(0) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + glUseProgram(0) + + glDisable(GL_MULTISAMPLE) + + self.draw_end() diff --git a/lib/renderer/gl/data/color.fs b/lib/renderer/gl/data/color.fs new file mode 100644 index 0000000000000000000000000000000000000000..96e904bdca2e813e7e7ae75e1cc5dda7b1d0b5be --- /dev/null +++ b/lib/renderer/gl/data/color.fs @@ -0,0 +1,20 @@ +#version 330 core + +layout (location = 0) out vec4 FragColor; +layout (location = 1) out vec4 FragNormal; +layout (location = 2) out vec4 FragDepth; + +in vec3 Color; +in vec3 CamNormal; +in vec3 depth; + + +void main() +{ + FragColor = vec4(Color,1.0); + + vec3 cam_norm_normalized = normalize(CamNormal); + vec3 rgb = (cam_norm_normalized + 1.0) / 2.0; + FragNormal = vec4(rgb, 1.0); + FragDepth = vec4(depth.xyz, 1.0); +} diff --git a/lib/renderer/gl/data/color.vs b/lib/renderer/gl/data/color.vs new file mode 100644 index 0000000000000000000000000000000000000000..1256f7eb3f3605b8f848d452fcf601d5a18b95e2 --- /dev/null +++ b/lib/renderer/gl/data/color.vs @@ -0,0 +1,29 @@ +#version 330 core + +layout (location = 0) in vec3 a_Position; +layout (location = 1) in vec3 a_Color; +layout (location = 2) in vec3 a_Normal; + +out vec3 CamNormal; +out vec3 CamPos; +out vec3 Color; +out vec3 depth; + + +uniform mat3 RotMat; +uniform mat4 NormMat; +uniform mat4 ModelMat; +uniform mat4 PerspMat; + +void main() +{ + vec3 a_Position = (NormMat * vec4(a_Position,1.0)).xyz; + gl_Position = PerspMat * ModelMat * vec4(RotMat * a_Position, 1.0); + Color = a_Color; + + mat3 R = mat3(ModelMat) * RotMat; + CamNormal = (R * a_Normal); + + depth = vec3(gl_Position.z / gl_Position.w); + +} \ No newline at end of file diff --git a/lib/renderer/gl/data/normal.fs b/lib/renderer/gl/data/normal.fs new file mode 100644 index 0000000000000000000000000000000000000000..9e2770952e27d9265ccb100833245beed3ebebe5 --- /dev/null +++ b/lib/renderer/gl/data/normal.fs @@ -0,0 +1,12 @@ +#version 330 + +out vec4 FragColor; + +in vec3 CamNormal; + +void main() +{ + vec3 cam_norm_normalized = normalize(CamNormal); + vec3 rgb = (cam_norm_normalized + 1.0) / 2.0; + FragColor = vec4(rgb, 1.0); +} \ No newline at end of file diff --git a/lib/renderer/gl/data/normal.vs b/lib/renderer/gl/data/normal.vs new file mode 100644 index 0000000000000000000000000000000000000000..a0f7f50b1cedfd677843b2a60cf9051b2134c347 --- /dev/null +++ b/lib/renderer/gl/data/normal.vs @@ -0,0 +1,15 @@ +#version 330 + +layout (location = 0) in vec3 Position; +layout (location = 1) in vec3 Normal; + +out vec3 CamNormal; + +uniform mat4 ModelMat; +uniform mat4 PerspMat; + +void main() +{ + gl_Position = PerspMat * ModelMat * vec4(Position, 1.0); + CamNormal = (ModelMat * vec4(Normal, 0.0)).xyz; +} \ No newline at end of file diff --git a/lib/renderer/gl/data/prt.fs b/lib/renderer/gl/data/prt.fs new file mode 100755 index 0000000000000000000000000000000000000000..3737e2a5b51d5b001cdb2f3d8157793caa98fd54 --- /dev/null +++ b/lib/renderer/gl/data/prt.fs @@ -0,0 +1,157 @@ +#version 330 + +uniform vec3 SHCoeffs[9]; +uniform uint analytic; + +uniform uint hasNormalMap; +uniform uint hasAlbedoMap; + +uniform sampler2D AlbedoMap; +uniform sampler2D NormalMap; + +in VertexData { + vec3 Position; + vec3 Depth; + vec3 ModelNormal; + vec2 Texcoord; + vec3 Tangent; + vec3 Bitangent; + vec3 PRT1; + vec3 PRT2; + vec3 PRT3; + vec3 Label; +} VertexIn; + +layout (location = 0) out vec4 FragColor; +layout (location = 1) out vec4 FragNormal; +layout (location = 2) out vec4 FragPosition; +layout (location = 3) out vec4 FragAlbedo; +layout (location = 4) out vec4 FragShading; +layout (location = 5) out vec4 FragPRT1; +layout (location = 6) out vec4 FragPRT2; +// layout (location = 7) out vec4 FragPRT3; +layout (location = 7) out vec4 FragLabel; + + +vec4 gammaCorrection(vec4 vec, float g) +{ + return vec4(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g), vec.w); +} + +vec3 gammaCorrection(vec3 vec, float g) +{ + return vec3(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g)); +} + +void evaluateH(vec3 n, out float H[9]) +{ + float c1 = 0.429043, c2 = 0.511664, + c3 = 0.743125, c4 = 0.886227, c5 = 0.247708; + + H[0] = c4; + H[1] = 2.0 * c2 * n[1]; + H[2] = 2.0 * c2 * n[2]; + H[3] = 2.0 * c2 * n[0]; + H[4] = 2.0 * c1 * n[0] * n[1]; + H[5] = 2.0 * c1 * n[1] * n[2]; + H[6] = c3 * n[2] * n[2] - c5; + H[7] = 2.0 * c1 * n[2] * n[0]; + H[8] = c1 * (n[0] * n[0] - n[1] * n[1]); +} + +vec3 evaluateLightingModel(vec3 normal) +{ + float H[9]; + evaluateH(normal, H); + vec3 res = vec3(0.0); + for (int i = 0; i < 9; i++) { + res += H[i] * SHCoeffs[i]; + } + return res; +} + +// nC: coarse geometry normal, nH: fine normal from normal map +vec3 evaluateLightingModelHybrid(vec3 nC, vec3 nH, mat3 prt) +{ + float HC[9], HH[9]; + evaluateH(nC, HC); + evaluateH(nH, HH); + + vec3 res = vec3(0.0); + vec3 shadow = vec3(0.0); + vec3 unshadow = vec3(0.0); + for(int i = 0; i < 3; ++i){ + for(int j = 0; j < 3; ++j){ + int id = i*3+j; + res += HH[id]* SHCoeffs[id]; + shadow += prt[i][j] * SHCoeffs[id]; + unshadow += HC[id] * SHCoeffs[id]; + } + } + vec3 ratio = clamp(shadow/unshadow,0.0,1.0); + res = ratio * res; + + return res; +} + +vec3 evaluateLightingModelPRT(mat3 prt) +{ + vec3 res = vec3(0.0); + for(int i = 0; i < 3; ++i){ + for(int j = 0; j < 3; ++j){ + res += prt[i][j] * SHCoeffs[i*3+j]; + } + } + + return res; +} + +void main() +{ + vec2 uv = VertexIn.Texcoord; + vec3 nC = normalize(VertexIn.ModelNormal); + vec3 nml = nC; + mat3 prt = mat3(VertexIn.PRT1, VertexIn.PRT2, VertexIn.PRT3); + + if(hasAlbedoMap == uint(0)) + FragAlbedo = vec4(1.0); + else + FragAlbedo = texture(AlbedoMap, uv);//gammaCorrection(texture(AlbedoMap, uv), 1.0/2.2); + + if(hasNormalMap == uint(0)) + { + if(analytic == uint(0)) + FragShading = vec4(evaluateLightingModelPRT(prt), 1.0f); + else + FragShading = vec4(evaluateLightingModel(nC), 1.0f); + } + else + { + vec3 n_tan = normalize(texture(NormalMap, uv).rgb*2.0-vec3(1.0)); + + mat3 TBN = mat3(normalize(VertexIn.Tangent),normalize(VertexIn.Bitangent),nC); + vec3 nH = normalize(TBN * n_tan); + + if(analytic == uint(0)) + FragShading = vec4(evaluateLightingModelHybrid(nC,nH,prt),1.0f); + else + FragShading = vec4(evaluateLightingModel(nH), 1.0f); + + nml = nH; + } + + FragShading = gammaCorrection(FragShading, 2.2); + FragColor = clamp(FragAlbedo * FragShading, 0.0, 1.0); + FragNormal = vec4(0.5*(nml+vec3(1.0)), 1.0); + FragPosition = vec4(VertexIn.Depth.xyz, 1.0); + FragShading = vec4(clamp(0.5*FragShading.xyz, 0.0, 1.0),1.0); + // FragColor = gammaCorrection(clamp(FragAlbedo * FragShading, 0.0, 1.0),2.2); + // FragNormal = vec4(0.5*(nml+vec3(1.0)), 1.0); + // FragPosition = vec4(VertexIn.Position,VertexIn.Depth.x); + // FragShading = vec4(gammaCorrection(clamp(0.5*FragShading.xyz, 0.0, 1.0),2.2),1.0); + // FragAlbedo = gammaCorrection(FragAlbedo,2.2); + FragPRT1 = vec4(VertexIn.PRT1,1.0); + FragPRT2 = vec4(VertexIn.PRT2,1.0); + // FragPRT3 = vec4(VertexIn.PRT3,1.0); + FragLabel = vec4(VertexIn.Label,1.0); +} \ No newline at end of file diff --git a/lib/renderer/gl/data/prt.vs b/lib/renderer/gl/data/prt.vs new file mode 100755 index 0000000000000000000000000000000000000000..4bb55cc6fa74539fa48aed054cac120ac5cf7ec2 --- /dev/null +++ b/lib/renderer/gl/data/prt.vs @@ -0,0 +1,171 @@ +#version 330 + +layout (location = 0) in vec3 a_Position; +layout (location = 1) in vec3 a_Normal; +layout (location = 2) in vec2 a_TextureCoord; +layout (location = 3) in vec3 a_Tangent; +layout (location = 4) in vec3 a_Bitangent; +layout (location = 5) in vec3 a_PRT1; +layout (location = 6) in vec3 a_PRT2; +layout (location = 7) in vec3 a_PRT3; +layout (location = 8) in vec3 a_Label; + +out VertexData { + vec3 Position; + vec3 Depth; + vec3 ModelNormal; + vec2 Texcoord; + vec3 Tangent; + vec3 Bitangent; + vec3 PRT1; + vec3 PRT2; + vec3 PRT3; + vec3 Label; +} VertexOut; + +uniform mat3 RotMat; +uniform mat4 NormMat; +uniform mat4 ModelMat; +uniform mat4 PerspMat; + +float s_c3 = 0.94617469575; // (3*sqrt(5))/(4*sqrt(pi)) +float s_c4 = -0.31539156525;// (-sqrt(5))/(4*sqrt(pi)) +float s_c5 = 0.54627421529; // (sqrt(15))/(4*sqrt(pi)) + +float s_c_scale = 1.0/0.91529123286551084; +float s_c_scale_inv = 0.91529123286551084; + +float s_rc2 = 1.5853309190550713*s_c_scale; +float s_c4_div_c3 = s_c4/s_c3; +float s_c4_div_c3_x2 = (s_c4/s_c3)*2.0; + +float s_scale_dst2 = s_c3 * s_c_scale_inv; +float s_scale_dst4 = s_c5 * s_c_scale_inv; + +void OptRotateBand0(float x[1], mat3 R, out float dst[1]) +{ + dst[0] = x[0]; +} + +// 9 multiplies +void OptRotateBand1(float x[3], mat3 R, out float dst[3]) +{ + // derived from SlowRotateBand1 + dst[0] = ( R[1][1])*x[0] + (-R[1][2])*x[1] + ( R[1][0])*x[2]; + dst[1] = (-R[2][1])*x[0] + ( R[2][2])*x[1] + (-R[2][0])*x[2]; + dst[2] = ( R[0][1])*x[0] + (-R[0][2])*x[1] + ( R[0][0])*x[2]; +} + +// 48 multiplies +void OptRotateBand2(float x[5], mat3 R, out float dst[5]) +{ + // Sparse matrix multiply + float sh0 = x[3] + x[4] + x[4] - x[1]; + float sh1 = x[0] + s_rc2*x[2] + x[3] + x[4]; + float sh2 = x[0]; + float sh3 = -x[3]; + float sh4 = -x[1]; + + // Rotations. R0 and R1 just use the raw matrix columns + float r2x = R[0][0] + R[0][1]; + float r2y = R[1][0] + R[1][1]; + float r2z = R[2][0] + R[2][1]; + + float r3x = R[0][0] + R[0][2]; + float r3y = R[1][0] + R[1][2]; + float r3z = R[2][0] + R[2][2]; + + float r4x = R[0][1] + R[0][2]; + float r4y = R[1][1] + R[1][2]; + float r4z = R[2][1] + R[2][2]; + + // dense matrix multiplication one column at a time + + // column 0 + float sh0_x = sh0 * R[0][0]; + float sh0_y = sh0 * R[1][0]; + float d0 = sh0_x * R[1][0]; + float d1 = sh0_y * R[2][0]; + float d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3); + float d3 = sh0_x * R[2][0]; + float d4 = sh0_x * R[0][0] - sh0_y * R[1][0]; + + // column 1 + float sh1_x = sh1 * R[0][2]; + float sh1_y = sh1 * R[1][2]; + d0 += sh1_x * R[1][2]; + d1 += sh1_y * R[2][2]; + d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3); + d3 += sh1_x * R[2][2]; + d4 += sh1_x * R[0][2] - sh1_y * R[1][2]; + + // column 2 + float sh2_x = sh2 * r2x; + float sh2_y = sh2 * r2y; + d0 += sh2_x * r2y; + d1 += sh2_y * r2z; + d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2); + d3 += sh2_x * r2z; + d4 += sh2_x * r2x - sh2_y * r2y; + + // column 3 + float sh3_x = sh3 * r3x; + float sh3_y = sh3 * r3y; + d0 += sh3_x * r3y; + d1 += sh3_y * r3z; + d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2); + d3 += sh3_x * r3z; + d4 += sh3_x * r3x - sh3_y * r3y; + + // column 4 + float sh4_x = sh4 * r4x; + float sh4_y = sh4 * r4y; + d0 += sh4_x * r4y; + d1 += sh4_y * r4z; + d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2); + d3 += sh4_x * r4z; + d4 += sh4_x * r4x - sh4_y * r4y; + + // extra multipliers + dst[0] = d0; + dst[1] = -d1; + dst[2] = d2 * s_scale_dst2; + dst[3] = -d3; + dst[4] = d4 * s_scale_dst4; +} + +void main() +{ + // normalization + vec3 pos = (NormMat * vec4(a_Position,1.0)).xyz; + + mat3 R = mat3(ModelMat) * RotMat; + VertexOut.ModelNormal = (R * a_Normal); + VertexOut.Position = R * pos; + VertexOut.Texcoord = a_TextureCoord; + VertexOut.Tangent = (R * a_Tangent); + VertexOut.Bitangent = (R * a_Bitangent); + VertexOut.Label = a_Label; + + float PRT0, PRT1[3], PRT2[5]; + PRT0 = a_PRT1[0]; + PRT1[0] = a_PRT1[1]; + PRT1[1] = a_PRT1[2]; + PRT1[2] = a_PRT2[0]; + PRT2[0] = a_PRT2[1]; + PRT2[1] = a_PRT2[2]; + PRT2[2] = a_PRT3[0]; + PRT2[3] = a_PRT3[1]; + PRT2[4] = a_PRT3[2]; + + OptRotateBand1(PRT1, R, PRT1); + OptRotateBand2(PRT2, R, PRT2); + + VertexOut.PRT1 = vec3(PRT0,PRT1[0],PRT1[1]); + VertexOut.PRT2 = vec3(PRT1[2],PRT2[0],PRT2[1]); + VertexOut.PRT3 = vec3(PRT2[2],PRT2[3],PRT2[4]); + + gl_Position = PerspMat * ModelMat * vec4(RotMat * pos, 1.0); + + VertexOut.Depth = vec3(gl_Position.z / gl_Position.w); +} diff --git a/lib/renderer/gl/data/prt_uv.fs b/lib/renderer/gl/data/prt_uv.fs new file mode 100755 index 0000000000000000000000000000000000000000..6e90b25c62b41c8cf61afd29333372193047d5f1 --- /dev/null +++ b/lib/renderer/gl/data/prt_uv.fs @@ -0,0 +1,141 @@ +#version 330 + +uniform vec3 SHCoeffs[9]; +uniform uint analytic; + +uniform uint hasNormalMap; +uniform uint hasAlbedoMap; + +uniform sampler2D AlbedoMap; +uniform sampler2D NormalMap; + +in VertexData { + vec3 Position; + vec3 ModelNormal; + vec3 CameraNormal; + vec2 Texcoord; + vec3 Tangent; + vec3 Bitangent; + vec3 PRT1; + vec3 PRT2; + vec3 PRT3; +} VertexIn; + +layout (location = 0) out vec4 FragColor; +layout (location = 1) out vec4 FragPosition; +layout (location = 2) out vec4 FragNormal; + +vec4 gammaCorrection(vec4 vec, float g) +{ + return vec4(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g), vec.w); +} + +vec3 gammaCorrection(vec3 vec, float g) +{ + return vec3(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g)); +} + +void evaluateH(vec3 n, out float H[9]) +{ + float c1 = 0.429043, c2 = 0.511664, + c3 = 0.743125, c4 = 0.886227, c5 = 0.247708; + + H[0] = c4; + H[1] = 2.0 * c2 * n[1]; + H[2] = 2.0 * c2 * n[2]; + H[3] = 2.0 * c2 * n[0]; + H[4] = 2.0 * c1 * n[0] * n[1]; + H[5] = 2.0 * c1 * n[1] * n[2]; + H[6] = c3 * n[2] * n[2] - c5; + H[7] = 2.0 * c1 * n[2] * n[0]; + H[8] = c1 * (n[0] * n[0] - n[1] * n[1]); +} + +vec3 evaluateLightingModel(vec3 normal) +{ + float H[9]; + evaluateH(normal, H); + vec3 res = vec3(0.0); + for (int i = 0; i < 9; i++) { + res += H[i] * SHCoeffs[i]; + } + return res; +} + +// nC: coarse geometry normal, nH: fine normal from normal map +vec3 evaluateLightingModelHybrid(vec3 nC, vec3 nH, mat3 prt) +{ + float HC[9], HH[9]; + evaluateH(nC, HC); + evaluateH(nH, HH); + + vec3 res = vec3(0.0); + vec3 shadow = vec3(0.0); + vec3 unshadow = vec3(0.0); + for(int i = 0; i < 3; ++i){ + for(int j = 0; j < 3; ++j){ + int id = i*3+j; + res += HH[id]* SHCoeffs[id]; + shadow += prt[i][j] * SHCoeffs[id]; + unshadow += HC[id] * SHCoeffs[id]; + } + } + vec3 ratio = clamp(shadow/unshadow,0.0,1.0); + res = ratio * res; + + return res; +} + +vec3 evaluateLightingModelPRT(mat3 prt) +{ + vec3 res = vec3(0.0); + for(int i = 0; i < 3; ++i){ + for(int j = 0; j < 3; ++j){ + res += prt[i][j] * SHCoeffs[i*3+j]; + } + } + + return res; +} + +void main() +{ + vec2 uv = VertexIn.Texcoord; + vec3 nM = normalize(VertexIn.ModelNormal); + vec3 nC = normalize(VertexIn.CameraNormal); + vec3 nml = nC; + mat3 prt = mat3(VertexIn.PRT1, VertexIn.PRT2, VertexIn.PRT3); + + vec4 albedo, shading; + if(hasAlbedoMap == uint(0)) + albedo = vec4(1.0); + else + albedo = texture(AlbedoMap, uv);//gammaCorrection(texture(AlbedoMap, uv), 1.0/2.2); + + if(hasNormalMap == uint(0)) + { + if(analytic == uint(0)) + shading = vec4(evaluateLightingModelPRT(prt), 1.0f); + else + shading = vec4(evaluateLightingModel(nC), 1.0f); + } + else + { + vec3 n_tan = normalize(texture(NormalMap, uv).rgb*2.0-vec3(1.0)); + + mat3 TBN = mat3(normalize(VertexIn.Tangent),normalize(VertexIn.Bitangent),nC); + vec3 nH = normalize(TBN * n_tan); + + if(analytic == uint(0)) + shading = vec4(evaluateLightingModelHybrid(nC,nH,prt),1.0f); + else + shading = vec4(evaluateLightingModel(nH), 1.0f); + + nml = nH; + } + + shading = gammaCorrection(shading, 2.2); + FragColor = clamp(albedo * shading, 0.0, 1.0); + FragPosition = vec4(VertexIn.Position,1.0); + FragNormal = vec4(0.5*(nM+vec3(1.0)),1.0); +} \ No newline at end of file diff --git a/lib/renderer/gl/data/prt_uv.vs b/lib/renderer/gl/data/prt_uv.vs new file mode 100755 index 0000000000000000000000000000000000000000..22a03564bd95158c3fb9edf513c0717975b93ee0 --- /dev/null +++ b/lib/renderer/gl/data/prt_uv.vs @@ -0,0 +1,168 @@ +#version 330 + +layout (location = 0) in vec3 a_Position; +layout (location = 1) in vec3 a_Normal; +layout (location = 2) in vec2 a_TextureCoord; +layout (location = 3) in vec3 a_Tangent; +layout (location = 4) in vec3 a_Bitangent; +layout (location = 5) in vec3 a_PRT1; +layout (location = 6) in vec3 a_PRT2; +layout (location = 7) in vec3 a_PRT3; + +out VertexData { + vec3 Position; + vec3 ModelNormal; + vec3 CameraNormal; + vec2 Texcoord; + vec3 Tangent; + vec3 Bitangent; + vec3 PRT1; + vec3 PRT2; + vec3 PRT3; +} VertexOut; + +uniform mat3 RotMat; +uniform mat4 NormMat; +uniform mat4 ModelMat; +uniform mat4 PerspMat; + +#define pi 3.1415926535897932384626433832795 + +float s_c3 = 0.94617469575; // (3*sqrt(5))/(4*sqrt(pi)) +float s_c4 = -0.31539156525;// (-sqrt(5))/(4*sqrt(pi)) +float s_c5 = 0.54627421529; // (sqrt(15))/(4*sqrt(pi)) + +float s_c_scale = 1.0/0.91529123286551084; +float s_c_scale_inv = 0.91529123286551084; + +float s_rc2 = 1.5853309190550713*s_c_scale; +float s_c4_div_c3 = s_c4/s_c3; +float s_c4_div_c3_x2 = (s_c4/s_c3)*2.0; + +float s_scale_dst2 = s_c3 * s_c_scale_inv; +float s_scale_dst4 = s_c5 * s_c_scale_inv; + +void OptRotateBand0(float x[1], mat3 R, out float dst[1]) +{ + dst[0] = x[0]; +} + +// 9 multiplies +void OptRotateBand1(float x[3], mat3 R, out float dst[3]) +{ + // derived from SlowRotateBand1 + dst[0] = ( R[1][1])*x[0] + (-R[1][2])*x[1] + ( R[1][0])*x[2]; + dst[1] = (-R[2][1])*x[0] + ( R[2][2])*x[1] + (-R[2][0])*x[2]; + dst[2] = ( R[0][1])*x[0] + (-R[0][2])*x[1] + ( R[0][0])*x[2]; +} + +// 48 multiplies +void OptRotateBand2(float x[5], mat3 R, out float dst[5]) +{ + // Sparse matrix multiply + float sh0 = x[3] + x[4] + x[4] - x[1]; + float sh1 = x[0] + s_rc2*x[2] + x[3] + x[4]; + float sh2 = x[0]; + float sh3 = -x[3]; + float sh4 = -x[1]; + + // Rotations. R0 and R1 just use the raw matrix columns + float r2x = R[0][0] + R[0][1]; + float r2y = R[1][0] + R[1][1]; + float r2z = R[2][0] + R[2][1]; + + float r3x = R[0][0] + R[0][2]; + float r3y = R[1][0] + R[1][2]; + float r3z = R[2][0] + R[2][2]; + + float r4x = R[0][1] + R[0][2]; + float r4y = R[1][1] + R[1][2]; + float r4z = R[2][1] + R[2][2]; + + // dense matrix multiplication one column at a time + + // column 0 + float sh0_x = sh0 * R[0][0]; + float sh0_y = sh0 * R[1][0]; + float d0 = sh0_x * R[1][0]; + float d1 = sh0_y * R[2][0]; + float d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3); + float d3 = sh0_x * R[2][0]; + float d4 = sh0_x * R[0][0] - sh0_y * R[1][0]; + + // column 1 + float sh1_x = sh1 * R[0][2]; + float sh1_y = sh1 * R[1][2]; + d0 += sh1_x * R[1][2]; + d1 += sh1_y * R[2][2]; + d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3); + d3 += sh1_x * R[2][2]; + d4 += sh1_x * R[0][2] - sh1_y * R[1][2]; + + // column 2 + float sh2_x = sh2 * r2x; + float sh2_y = sh2 * r2y; + d0 += sh2_x * r2y; + d1 += sh2_y * r2z; + d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2); + d3 += sh2_x * r2z; + d4 += sh2_x * r2x - sh2_y * r2y; + + // column 3 + float sh3_x = sh3 * r3x; + float sh3_y = sh3 * r3y; + d0 += sh3_x * r3y; + d1 += sh3_y * r3z; + d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2); + d3 += sh3_x * r3z; + d4 += sh3_x * r3x - sh3_y * r3y; + + // column 4 + float sh4_x = sh4 * r4x; + float sh4_y = sh4 * r4y; + d0 += sh4_x * r4y; + d1 += sh4_y * r4z; + d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2); + d3 += sh4_x * r4z; + d4 += sh4_x * r4x - sh4_y * r4y; + + // extra multipliers + dst[0] = d0; + dst[1] = -d1; + dst[2] = d2 * s_scale_dst2; + dst[3] = -d3; + dst[4] = d4 * s_scale_dst4; +} + +void main() +{ + // normalization + mat3 R = mat3(ModelMat) * RotMat; + VertexOut.ModelNormal = a_Normal; + VertexOut.CameraNormal = (R * a_Normal); + VertexOut.Position = a_Position; + VertexOut.Texcoord = a_TextureCoord; + VertexOut.Tangent = (R * a_Tangent); + VertexOut.Bitangent = (R * a_Bitangent); + float PRT0, PRT1[3], PRT2[5]; + PRT0 = a_PRT1[0]; + PRT1[0] = a_PRT1[1]; + PRT1[1] = a_PRT1[2]; + PRT1[2] = a_PRT2[0]; + PRT2[0] = a_PRT2[1]; + PRT2[1] = a_PRT2[2]; + PRT2[2] = a_PRT3[0]; + PRT2[3] = a_PRT3[1]; + PRT2[4] = a_PRT3[2]; + + OptRotateBand1(PRT1, R, PRT1); + OptRotateBand2(PRT2, R, PRT2); + + VertexOut.PRT1 = vec3(PRT0,PRT1[0],PRT1[1]); + VertexOut.PRT2 = vec3(PRT1[2],PRT2[0],PRT2[1]); + VertexOut.PRT3 = vec3(PRT2[2],PRT2[3],PRT2[4]); + + gl_Position = vec4(a_TextureCoord, 0.0, 1.0) - vec4(0.5, 0.5, 0, 0); + gl_Position[0] *= 2.0; + gl_Position[1] *= 2.0; +} diff --git a/lib/renderer/gl/data/quad.fs b/lib/renderer/gl/data/quad.fs new file mode 100755 index 0000000000000000000000000000000000000000..f43502f2352ca2adf19d11e809946b51498df5a5 --- /dev/null +++ b/lib/renderer/gl/data/quad.fs @@ -0,0 +1,11 @@ +#version 330 core +out vec4 FragColor; + +in vec2 TexCoord; + +uniform sampler2D screenTexture; + +void main() +{ + FragColor = texture(screenTexture, TexCoord); +} \ No newline at end of file diff --git a/lib/renderer/gl/data/quad.vs b/lib/renderer/gl/data/quad.vs new file mode 100755 index 0000000000000000000000000000000000000000..811044631a1f29f5b45c490b2d40297f3127b6ea --- /dev/null +++ b/lib/renderer/gl/data/quad.vs @@ -0,0 +1,11 @@ +#version 330 core +layout (location = 0) in vec2 aPos; +layout (location = 1) in vec2 aTexCoord; + +out vec2 TexCoord; + +void main() +{ + gl_Position = vec4(aPos.x, aPos.y, 0.0, 1.0); + TexCoord = aTexCoord; +} \ No newline at end of file diff --git a/lib/renderer/gl/framework.py b/lib/renderer/gl/framework.py new file mode 100755 index 0000000000000000000000000000000000000000..3ae6ad0d09e475e89b38633d457c4b3b68c881c7 --- /dev/null +++ b/lib/renderer/gl/framework.py @@ -0,0 +1,95 @@ +# Mario Rosasco, 2016 +# adapted from framework.cpp, Copyright (C) 2010-2012 by Jason L. McKesson +# This file is licensed under the MIT License. +# +# NB: Unlike in the framework.cpp organization, the main loop is contained +# in the tutorial files, not in this framework file. Additionally, a copy of +# this module file must exist in the same directory as the tutorial files +# to be imported properly. + +import os +from OpenGL.GL import * + + +# Function that creates and compiles shaders according to the given type (a GL enum value) and +# shader program (a file containing a GLSL program). +def loadShader(shaderType, shaderFile): + # check if file exists, get full path name + strFilename = findFileOrThrow(shaderFile) + shaderData = None + with open(strFilename, 'r') as f: + shaderData = f.read() + + shader = glCreateShader(shaderType) + glShaderSource( + shader, + shaderData) # note that this is a simpler function call than in C + + # This shader compilation is more explicit than the one used in + # framework.cpp, which relies on a glutil wrapper function. + # This is made explicit here mainly to decrease dependence on pyOpenGL + # utilities and wrappers, which docs caution may change in future versions. + glCompileShader(shader) + + status = glGetShaderiv(shader, GL_COMPILE_STATUS) + if status == GL_FALSE: + # Note that getting the error log is much simpler in Python than in C/C++ + # and does not require explicit handling of the string buffer + strInfoLog = glGetShaderInfoLog(shader) + strShaderType = "" + if shaderType is GL_VERTEX_SHADER: + strShaderType = "vertex" + elif shaderType is GL_GEOMETRY_SHADER: + strShaderType = "geometry" + elif shaderType is GL_FRAGMENT_SHADER: + strShaderType = "fragment" + + print("Compilation failure for " + strShaderType + " shader:\n" + + str(strInfoLog)) + + return shader + + +# Function that accepts a list of shaders, compiles them, and returns a handle to the compiled program +def createProgram(shaderList): + program = glCreateProgram() + + for shader in shaderList: + glAttachShader(program, shader) + + glLinkProgram(program) + + status = glGetProgramiv(program, GL_LINK_STATUS) + if status == GL_FALSE: + # Note that getting the error log is much simpler in Python than in C/C++ + # and does not require explicit handling of the string buffer + strInfoLog = glGetProgramInfoLog(program) + print("Linker failure: \n" + str(strInfoLog)) + + for shader in shaderList: + glDetachShader(program, shader) + + return program + + +# Helper function to locate and open the target file (passed in as a string). +# Returns the full path to the file as a string. +def findFileOrThrow(strBasename): + # Keep constant names in C-style convention, for readability + # when comparing to C(/C++) code. + if os.path.isfile(strBasename): + return strBasename + + LOCAL_FILE_DIR = "data" + os.sep + GLOBAL_FILE_DIR = os.path.dirname( + os.path.abspath(__file__)) + os.sep + "data" + os.sep + + strFilename = LOCAL_FILE_DIR + strBasename + if os.path.isfile(strFilename): + return strFilename + + strFilename = GLOBAL_FILE_DIR + strBasename + if os.path.isfile(strFilename): + return strFilename + + raise IOError('Could not find target file ' + strBasename) diff --git a/lib/renderer/gl/glcontext.py b/lib/renderer/gl/glcontext.py new file mode 100755 index 0000000000000000000000000000000000000000..6a60389bf098fbce59a9a31f05eb9fd1c17ca454 --- /dev/null +++ b/lib/renderer/gl/glcontext.py @@ -0,0 +1,136 @@ +"""Headless GPU-accelerated OpenGL context creation on Google Colaboratory. + +Typical usage: + + # Optional PyOpenGL configuratiopn can be done here. + # import OpenGL + # OpenGL.ERROR_CHECKING = True + + # 'glcontext' must be imported before any OpenGL.* API. + from lucid.misc.gl.glcontext import create_opengl_context + + # Now it's safe to import OpenGL and EGL functions + import OpenGL.GL as gl + + # create_opengl_context() creates a GL context that is attached to an + # offscreen surface of the specified size. Note that rendering to buffers + # of other sizes and formats is still possible with OpenGL Framebuffers. + # + # Users are expected to directly use the EGL API in case more advanced + # context management is required. + width, height = 640, 480 + create_opengl_context((width, height)) + + # OpenGL context is available here. + +""" + +from __future__ import print_function + +# pylint: disable=unused-import,g-import-not-at-top,g-statement-before-imports + +try: + import OpenGL +except: + print('This module depends on PyOpenGL.') + print('Please run "\033[1m!pip install -q pyopengl\033[0m" ' + 'prior importing this module.') + raise + +import ctypes +from ctypes import pointer, util +import os + +os.environ['PYOPENGL_PLATFORM'] = 'egl' + +# OpenGL loading workaround. +# +# * PyOpenGL tries to load libGL, but we need libOpenGL, see [1,2]. +# This could have been solved by a symlink libGL->libOpenGL, but: +# +# * Python 2.7 can't find libGL and linEGL due to a bug (see [3]) +# in ctypes.util, that was only wixed in Python 3.6. +# +# So, the only solution I've found is to monkeypatch ctypes.util +# [1] https://devblogs.nvidia.com/egl-eye-opengl-visualization-without-x-server/ +# [2] https://devblogs.nvidia.com/linking-opengl-server-side-rendering/ +# [3] https://bugs.python.org/issue9998 +_find_library_old = ctypes.util.find_library +try: + + def _find_library_new(name): + return { + 'GL': 'libOpenGL.so', + 'EGL': 'libEGL.so', + }.get(name, _find_library_old(name)) + + util.find_library = _find_library_new + import OpenGL.GL as gl + import OpenGL.EGL as egl +except: + print('Unable to load OpenGL libraries. ' + 'Make sure you use GPU-enabled backend.') + print('Press "Runtime->Change runtime type" and set ' + '"Hardware accelerator" to GPU.') + raise +finally: + util.find_library = _find_library_old + + +def create_opengl_context(surface_size=(640, 480)): + """Create offscreen OpenGL context and make it current. + + Users are expected to directly use EGL API in case more advanced + context management is required. + + Args: + surface_size: (width, height), size of the offscreen rendering surface. + """ + egl_display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY) + + major, minor = egl.EGLint(), egl.EGLint() + egl.eglInitialize(egl_display, pointer(major), pointer(minor)) + + config_attribs = [ + egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT, egl.EGL_BLUE_SIZE, 8, + egl.EGL_GREEN_SIZE, 8, egl.EGL_RED_SIZE, 8, egl.EGL_DEPTH_SIZE, 24, + egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT, egl.EGL_NONE + ] + config_attribs = (egl.EGLint * len(config_attribs))(*config_attribs) + + num_configs = egl.EGLint() + egl_cfg = egl.EGLConfig() + egl.eglChooseConfig(egl_display, config_attribs, pointer(egl_cfg), 1, + pointer(num_configs)) + + width, height = surface_size + pbuffer_attribs = [ + egl.EGL_WIDTH, + width, + egl.EGL_HEIGHT, + height, + egl.EGL_NONE, + ] + pbuffer_attribs = (egl.EGLint * len(pbuffer_attribs))(*pbuffer_attribs) + egl_surf = egl.eglCreatePbufferSurface(egl_display, egl_cfg, + pbuffer_attribs) + + egl.eglBindAPI(egl.EGL_OPENGL_API) + + context_attribs = None + # context_attribs = [ + # egl.EGL_CONTEXT_MAJOR_VERSION, + # 4, + # egl.EGL_CONTEXT_MINOR_VERSION, + # 1, + # egl.EGL_NONE, + # ] + + egl_context = egl.eglCreateContext(egl_display, egl_cfg, + egl.EGL_NO_CONTEXT, context_attribs) + egl.eglMakeCurrent(egl_display, egl_surf, egl_surf, egl_context) + + buffer_type = egl.EGLint() + out = egl.eglQueryContext(egl_display, egl_context, + egl.EGL_CONTEXT_CLIENT_VERSION, buffer_type) + # print(buffer_type) diff --git a/lib/renderer/gl/init_gl.py b/lib/renderer/gl/init_gl.py new file mode 100644 index 0000000000000000000000000000000000000000..613c1d034733cab395617e61c80fc3cc716e3759 --- /dev/null +++ b/lib/renderer/gl/init_gl.py @@ -0,0 +1,24 @@ +_glut_window = None +_context_inited = None + + +def initialize_GL_context(width=512, height=512, egl=False): + ''' + default context uses GLUT + ''' + if not egl: + import OpenGL.GLUT as GLUT + display_mode = GLUT.GLUT_DOUBLE | GLUT.GLUT_RGB | GLUT.GLUT_DEPTH + global _glut_window + if _glut_window is None: + GLUT.glutInit() + GLUT.glutInitDisplayMode(display_mode) + GLUT.glutInitWindowSize(width, height) + GLUT.glutInitWindowPosition(0, 0) + _glut_window = GLUT.glutCreateWindow("My Render.") + else: + from .glcontext import create_opengl_context + global _context_inited + if _context_inited is None: + create_opengl_context((width, height)) + _context_inited = True diff --git a/lib/renderer/gl/norm_render.py b/lib/renderer/gl/norm_render.py new file mode 100755 index 0000000000000000000000000000000000000000..3ba29a678280a0903361b9d03c4d0d05f42d7fbf --- /dev/null +++ b/lib/renderer/gl/norm_render.py @@ -0,0 +1,79 @@ +''' +MIT License + +Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +''' +from OpenGL.GLUT import * + +from .render2 import Render + + +class NormRender(Render): + def __init__(self, + width=1600, + height=1200, + name='Cam Renderer', + program_files=['simple.fs', 'simple.vs'], + color_size=1, + ms_rate=1): + Render.__init__(self, width, height, name, program_files, color_size, + ms_rate) + self.camera = None + + glutDisplayFunc(self.display) + glutKeyboardFunc(self.keyboard) + + def set_camera(self, camera): + self.camera = camera + self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() + + def set_matrices(self, projection, modelview): + self.projection_matrix = projection + self.model_view_matrix = modelview + + def keyboard(self, key, x, y): + # up + eps = 1 + # print(key) + if key == b'w': + self.camera.center += eps * self.camera.direction + elif key == b's': + self.camera.center -= eps * self.camera.direction + if key == b'a': + self.camera.center -= eps * self.camera.right + elif key == b'd': + self.camera.center += eps * self.camera.right + if key == b' ': + self.camera.center += eps * self.camera.up + elif key == b'x': + self.camera.center -= eps * self.camera.up + elif key == b'i': + self.camera.near += 0.1 * eps + self.camera.far += 0.1 * eps + elif key == b'o': + self.camera.near -= 0.1 * eps + self.camera.far -= 0.1 * eps + + self.projection_matrix, self.model_view_matrix = self.camera.get_gl_matrix( + ) + + def show(self): + glutMainLoop() diff --git a/lib/renderer/gl/normal_render.py b/lib/renderer/gl/normal_render.py new file mode 100644 index 0000000000000000000000000000000000000000..e4eb41c1564719636e606714cc01bbed9356ff21 --- /dev/null +++ b/lib/renderer/gl/normal_render.py @@ -0,0 +1,97 @@ +''' +MIT License + +Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +''' +import numpy as np +import math + +from .framework import * +from .norm_render import NormRender + + +class NormalRender(NormRender): + def __init__(self, width=1600, height=1200, name='Normal Renderer'): + NormRender.__init__(self, + width, + height, + name, + program_files=['normal.vs', 'normal.fs']) + + self.norm_buffer = glGenBuffers(1) + + self.norm_data = None + + def set_normal_mesh(self, vertices, faces, norms, face_normals): + NormRender.set_mesh(self, vertices, faces) + + self.norm_data = norms[face_normals.reshape([-1])] + + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer) + glBufferData(GL_ARRAY_BUFFER, self.norm_data, GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def euler_to_rot_mat(self, r_x, r_y, r_z): + R_x = np.array([[1, 0, 0], [0, math.cos(r_x), -math.sin(r_x)], + [0, math.sin(r_x), math.cos(r_x)]]) + + R_y = np.array([[math.cos(r_y), 0, math.sin(r_y)], [0, 1, 0], + [-math.sin(r_y), 0, math.cos(r_y)]]) + + R_z = np.array([[math.cos(r_z), -math.sin(r_z), 0], + [math.sin(r_z), math.cos(r_z), 0], [0, 0, 1]]) + + R = np.dot(R_z, np.dot(R_y, R_x)) + + return R + + def draw(self): + self.draw_init() + + glUseProgram(self.program) + glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, + self.model_view_matrix.transpose()) + glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, + self.projection_matrix.transpose()) + + # Handle vertex buffer + glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) + + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle normal buffer + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer) + + glEnableVertexAttribArray(1) + glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) + + glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) + + glDisableVertexAttribArray(1) + glDisableVertexAttribArray(0) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + glUseProgram(0) + + self.draw_end() diff --git a/lib/renderer/gl/prt_render.py b/lib/renderer/gl/prt_render.py new file mode 100755 index 0000000000000000000000000000000000000000..a8b3ce290c6d8b7e0a3eb894230a99b0329a7818 --- /dev/null +++ b/lib/renderer/gl/prt_render.py @@ -0,0 +1,450 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np +import random + +from .framework import * +from .cam_render import CamRender + + +class PRTRender(CamRender): + def __init__(self, + width=1600, + height=1200, + name='PRT Renderer', + uv_mode=False, + ms_rate=1, + egl=False): + program_files = ['prt.vs', 'prt.fs' + ] if not uv_mode else ['prt_uv.vs', 'prt_uv.fs'] + CamRender.__init__(self, + width, + height, + name, + program_files=program_files, + color_size=8, + ms_rate=ms_rate, + egl=egl) + + # WARNING: this differs from vertex_buffer and vertex_data in Render + self.vert_buffer = {} + self.vert_data = {} + + self.vert_label_buffer = {} + self.vert_label_data = {} + + self.norm_buffer = {} + self.norm_data = {} + + self.tan_buffer = {} + self.tan_data = {} + + self.btan_buffer = {} + self.btan_data = {} + + self.prt1_buffer = {} + self.prt1_data = {} + + self.prt2_buffer = {} + self.prt2_data = {} + + self.prt3_buffer = {} + self.prt3_data = {} + + self.uv_buffer = {} + self.uv_data = {} + + self.render_texture_mat = {} + + self.vertex_dim = {} + self.n_vertices = {} + self.label_dim = {} + + self.norm_mat_unif = glGetUniformLocation(self.program, 'NormMat') + self.normalize_matrix = np.eye(4) + + self.shcoeff_unif = glGetUniformLocation(self.program, 'SHCoeffs') + self.shcoeffs = np.zeros((9, 3)) + self.shcoeffs[0, :] = 1.0 + #self.shcoeffs[1:,:] = np.random.rand(8,3) + + self.hasAlbedoUnif = glGetUniformLocation(self.program, 'hasAlbedoMap') + self.hasNormalUnif = glGetUniformLocation(self.program, 'hasNormalMap') + + self.analyticUnif = glGetUniformLocation(self.program, 'analytic') + self.analytic = False + + self.rot_mat_unif = glGetUniformLocation(self.program, 'RotMat') + self.rot_matrix = np.eye(3) + + def set_texture(self, mat_name, smplr_name, texture): + # texture_image: H x W x 3 + width = texture.shape[1] + height = texture.shape[0] + texture = np.flip(texture, 0) + img_data = np.fromstring(texture.tostring(), np.uint8) + + if mat_name not in self.render_texture_mat: + self.render_texture_mat[mat_name] = {} + if smplr_name in self.render_texture_mat[mat_name].keys(): + glDeleteTextures([self.render_texture_mat[mat_name][smplr_name]]) + del self.render_texture_mat[mat_name][smplr_name] + + self.render_texture_mat[mat_name][smplr_name] = glGenTextures(1) + glActiveTexture(GL_TEXTURE0) + + glPixelStorei(GL_UNPACK_ALIGNMENT, 1) + glBindTexture(GL_TEXTURE_2D, + self.render_texture_mat[mat_name][smplr_name]) + + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, + GL_UNSIGNED_BYTE, img_data) + + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, 3) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_LINEAR_MIPMAP_LINEAR) + + glGenerateMipmap(GL_TEXTURE_2D) + + def set_albedo(self, texture_image, mat_name='all'): + self.set_texture(mat_name, 'AlbedoMap', texture_image) + + def set_normal_map(self, texture_image, mat_name='all'): + self.set_texture(mat_name, 'NormalMap', texture_image) + + def set_mesh(self, + vertices, + faces, + norms, + faces_nml, + uvs, + faces_uvs, + prt, + faces_prt, + tans, + bitans, + verts_label=None, + mat_name='all'): + + self.vert_data[mat_name] = vertices[faces.reshape([-1])] + self.vert_label_data[mat_name] = verts_label[faces.reshape([-1])] + self.n_vertices[mat_name] = self.vert_data[mat_name].shape[0] + self.vertex_dim[mat_name] = self.vert_data[mat_name].shape[1] + self.label_dim[mat_name] = self.vert_label_data[mat_name].shape[1] + + if mat_name not in self.vert_buffer.keys(): + self.vert_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.vert_data[mat_name], GL_STATIC_DRAW) + + if mat_name not in self.vert_label_buffer.keys(): + self.vert_label_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.vert_label_data[mat_name], + GL_STATIC_DRAW) + + self.uv_data[mat_name] = uvs[faces_uvs.reshape([-1])] + if mat_name not in self.uv_buffer.keys(): + self.uv_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.uv_data[mat_name], GL_STATIC_DRAW) + + self.norm_data[mat_name] = norms[faces_nml.reshape([-1])] + if mat_name not in self.norm_buffer.keys(): + self.norm_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.norm_data[mat_name], GL_STATIC_DRAW) + + self.tan_data[mat_name] = tans[faces_nml.reshape([-1])] + if mat_name not in self.tan_buffer.keys(): + self.tan_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.tan_data[mat_name], GL_STATIC_DRAW) + + self.btan_data[mat_name] = bitans[faces_nml.reshape([-1])] + if mat_name not in self.btan_buffer.keys(): + self.btan_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.btan_data[mat_name], GL_STATIC_DRAW) + + self.prt1_data[mat_name] = prt[faces_prt.reshape([-1])][:, :3] + self.prt2_data[mat_name] = prt[faces_prt.reshape([-1])][:, 3:6] + self.prt3_data[mat_name] = prt[faces_prt.reshape([-1])][:, 6:] + + if mat_name not in self.prt1_buffer.keys(): + self.prt1_buffer[mat_name] = glGenBuffers(1) + if mat_name not in self.prt2_buffer.keys(): + self.prt2_buffer[mat_name] = glGenBuffers(1) + if mat_name not in self.prt3_buffer.keys(): + self.prt3_buffer[mat_name] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.prt1_data[mat_name], GL_STATIC_DRAW) + glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.prt2_data[mat_name], GL_STATIC_DRAW) + glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[mat_name]) + glBufferData(GL_ARRAY_BUFFER, self.prt3_data[mat_name], GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def set_mesh_mtl(self, + vertices, + faces, + norms, + faces_nml, + uvs, + faces_uvs, + tans, + bitans, + prt, + verts_label=None): + for key in faces: + self.vert_data[key] = vertices[faces[key].reshape([-1])] + self.vert_label_data[key] = verts_label[faces[key].reshape([-1])] + self.n_vertices[key] = self.vert_data[key].shape[0] + self.vertex_dim[key] = self.vert_data[key].shape[1] + self.label_dim[key] = self.vert_label_data[key].shape[1] + + if key not in self.vert_buffer.keys(): + self.vert_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.vert_data[key], GL_STATIC_DRAW) + + if key not in self.vert_label_buffer.keys(): + self.vert_label_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.vert_label_data[key], + GL_STATIC_DRAW) + + self.uv_data[key] = uvs[faces_uvs[key].reshape([-1])] + if key not in self.uv_buffer.keys(): + self.uv_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.uv_data[key], GL_STATIC_DRAW) + + self.norm_data[key] = norms[faces_nml[key].reshape([-1])] + if key not in self.norm_buffer.keys(): + self.norm_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.norm_data[key], GL_STATIC_DRAW) + + self.tan_data[key] = tans[faces_nml[key].reshape([-1])] + if key not in self.tan_buffer.keys(): + self.tan_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.tan_data[key], GL_STATIC_DRAW) + + self.btan_data[key] = bitans[faces_nml[key].reshape([-1])] + if key not in self.btan_buffer.keys(): + self.btan_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.btan_data[key], GL_STATIC_DRAW) + + self.prt1_data[key] = prt[faces[key].reshape([-1])][:, :3] + self.prt2_data[key] = prt[faces[key].reshape([-1])][:, 3:6] + self.prt3_data[key] = prt[faces[key].reshape([-1])][:, 6:] + + if key not in self.prt1_buffer.keys(): + self.prt1_buffer[key] = glGenBuffers(1) + if key not in self.prt2_buffer.keys(): + self.prt2_buffer[key] = glGenBuffers(1) + if key not in self.prt3_buffer.keys(): + self.prt3_buffer[key] = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.prt1_data[key], GL_STATIC_DRAW) + glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.prt2_data[key], GL_STATIC_DRAW) + glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[key]) + glBufferData(GL_ARRAY_BUFFER, self.prt3_data[key], GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def cleanup(self): + + glBindBuffer(GL_ARRAY_BUFFER, 0) + for key in self.vert_data: + glDeleteBuffers(1, [self.vert_buffer[key]]) + glDeleteBuffers(1, [self.norm_buffer[key]]) + glDeleteBuffers(1, [self.uv_buffer[key]]) + glDeleteBuffers(1, [self.vert_label_buffer[key]]) + + glDeleteBuffers(1, [self.tan_buffer[key]]) + glDeleteBuffers(1, [self.btan_buffer[key]]) + glDeleteBuffers(1, [self.prt1_buffer[key]]) + glDeleteBuffers(1, [self.prt2_buffer[key]]) + glDeleteBuffers(1, [self.prt3_buffer[key]]) + + glDeleteBuffers(1, []) + + for smplr in self.render_texture_mat[key]: + glDeleteTextures([self.render_texture_mat[key][smplr]]) + + self.vert_buffer = {} + self.vert_data = {} + + self.vert_label_buffer = {} + self.vert_label_data = {} + + self.norm_buffer = {} + self.norm_data = {} + + self.tan_buffer = {} + self.tan_data = {} + + self.btan_buffer = {} + self.btan_data = {} + + self.prt1_buffer = {} + self.prt1_data = {} + + self.prt2_buffer = {} + self.prt2_data = {} + + self.prt3_buffer = {} + self.prt3_data = {} + + self.uv_buffer = {} + self.uv_data = {} + + self.render_texture_mat = {} + + self.vertex_dim = {} + self.n_vertices = {} + self.label_dim = {} + + def randomize_sh(self): + self.shcoeffs[0, :] = 0.8 + self.shcoeffs[1:, :] = 1.0 * np.random.rand(8, 3) + + def set_sh(self, sh): + self.shcoeffs = sh + + def set_norm_mat(self, scale, center): + N = np.eye(4) + N[:3, :3] = scale * np.eye(3) + N[:3, 3] = -scale * center + + self.normalize_matrix = N + + def draw(self): + self.draw_init() + + glDisable(GL_BLEND) + #glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) + glEnable(GL_MULTISAMPLE) + + glUseProgram(self.program) + glUniformMatrix4fv(self.norm_mat_unif, 1, GL_FALSE, + self.normalize_matrix.transpose()) + glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, + self.model_view_matrix.transpose()) + glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, + self.projection_matrix.transpose()) + + if 'AlbedoMap' in self.render_texture_mat['all']: + glUniform1ui(self.hasAlbedoUnif, GLuint(1)) + else: + glUniform1ui(self.hasAlbedoUnif, GLuint(0)) + + if 'NormalMap' in self.render_texture_mat['all']: + glUniform1ui(self.hasNormalUnif, GLuint(1)) + else: + glUniform1ui(self.hasNormalUnif, GLuint(0)) + + glUniform1ui(self.analyticUnif, + GLuint(1) if self.analytic else GLuint(0)) + + glUniform3fv(self.shcoeff_unif, 9, self.shcoeffs) + + glUniformMatrix3fv(self.rot_mat_unif, 1, GL_FALSE, + self.rot_matrix.transpose()) + + for mat in self.vert_buffer: + # Handle vertex buffer + glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat]) + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, self.vertex_dim[mat], GL_DOUBLE, GL_FALSE, + 0, None) + + # Handle normal buffer + glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat]) + glEnableVertexAttribArray(1) + glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle uv buffer + glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[mat]) + glEnableVertexAttribArray(2) + glVertexAttribPointer(2, 2, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle tan buffer + glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[mat]) + glEnableVertexAttribArray(3) + glVertexAttribPointer(3, 3, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle btan buffer + glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[mat]) + glEnableVertexAttribArray(4) + glVertexAttribPointer(4, 3, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle PTR buffer + glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[mat]) + glEnableVertexAttribArray(5) + glVertexAttribPointer(5, 3, GL_DOUBLE, GL_FALSE, 0, None) + + glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[mat]) + glEnableVertexAttribArray(6) + glVertexAttribPointer(6, 3, GL_DOUBLE, GL_FALSE, 0, None) + + glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[mat]) + glEnableVertexAttribArray(7) + glVertexAttribPointer(7, 3, GL_DOUBLE, GL_FALSE, 0, None) + + # Handle vertex label buffer + glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[mat]) + glEnableVertexAttribArray(8) + glVertexAttribPointer(8, self.label_dim[mat], GL_DOUBLE, GL_FALSE, + 0, None) + + for i, smplr in enumerate(self.render_texture_mat[mat]): + glActiveTexture(GL_TEXTURE0 + i) + glBindTexture(GL_TEXTURE_2D, + self.render_texture_mat[mat][smplr]) + glUniform1i(glGetUniformLocation(self.program, smplr), i) + + glDrawArrays(GL_TRIANGLES, 0, self.n_vertices[mat]) + + glDisableVertexAttribArray(8) + glDisableVertexAttribArray(7) + glDisableVertexAttribArray(6) + glDisableVertexAttribArray(5) + glDisableVertexAttribArray(4) + glDisableVertexAttribArray(3) + glDisableVertexAttribArray(2) + glDisableVertexAttribArray(1) + glDisableVertexAttribArray(0) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + glUseProgram(0) + + glDisable(GL_BLEND) + glDisable(GL_MULTISAMPLE) + + self.draw_end() diff --git a/lib/renderer/gl/render.py b/lib/renderer/gl/render.py new file mode 100755 index 0000000000000000000000000000000000000000..94a530a04c4e4229df3d77331f01804a41d247b9 --- /dev/null +++ b/lib/renderer/gl/render.py @@ -0,0 +1,380 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from ctypes import * + +import numpy as np +from .framework import * + +GLUT = None + + +# NOTE: Render class assumes GL context is created already. +class Render: + def __init__(self, + width=1600, + height=1200, + name='GL Renderer', + program_files=['simple.fs', 'simple.vs'], + color_size=1, + ms_rate=1, + egl=False): + self.width = width + self.height = height + self.name = name + self.use_inverse_depth = False + self.egl = egl + + glEnable(GL_DEPTH_TEST) + + glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE) + glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE) + glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE) + + # init program + shader_list = [] + + for program_file in program_files: + _, ext = os.path.splitext(program_file) + if ext == '.vs': + shader_list.append(loadShader(GL_VERTEX_SHADER, program_file)) + elif ext == '.fs': + shader_list.append(loadShader(GL_FRAGMENT_SHADER, + program_file)) + elif ext == '.gs': + shader_list.append(loadShader(GL_GEOMETRY_SHADER, + program_file)) + + self.program = createProgram(shader_list) + + for shader in shader_list: + glDeleteShader(shader) + + # Init uniform variables + self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat') + self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat') + + self.vertex_buffer = glGenBuffers(1) + + # Init screen quad program and buffer + self.quad_program, self.quad_buffer = self.init_quad_program() + + # Configure frame buffer + self.frame_buffer = glGenFramebuffers(1) + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + + self.intermediate_fbo = None + if ms_rate > 1: + # Configure texture buffer to render to + self.color_buffer = [] + for i in range(color_size): + color_buffer = glGenTextures(1) + multi_sample_rate = ms_rate + glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, color_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_LINEAR) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_LINEAR) + glTexImage2DMultisample(GL_TEXTURE_2D_MULTISAMPLE, + multi_sample_rate, GL_RGBA32F, + self.width, self.height, GL_TRUE) + glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, 0) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, + GL_TEXTURE_2D_MULTISAMPLE, color_buffer, + 0) + self.color_buffer.append(color_buffer) + + self.render_buffer = glGenRenderbuffers(1) + glBindRenderbuffer(GL_RENDERBUFFER, self.render_buffer) + glRenderbufferStorageMultisample(GL_RENDERBUFFER, + multi_sample_rate, + GL_DEPTH24_STENCIL8, self.width, + self.height) + glBindRenderbuffer(GL_RENDERBUFFER, 0) + glFramebufferRenderbuffer(GL_FRAMEBUFFER, + GL_DEPTH_STENCIL_ATTACHMENT, + GL_RENDERBUFFER, self.render_buffer) + + attachments = [] + for i in range(color_size): + attachments.append(GL_COLOR_ATTACHMENT0 + i) + glDrawBuffers(color_size, attachments) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + self.intermediate_fbo = glGenFramebuffers(1) + glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo) + + self.screen_texture = [] + for i in range(color_size): + screen_texture = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, screen_texture) + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, + self.height, 0, GL_RGBA, GL_FLOAT, None) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_LINEAR) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_LINEAR) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, + screen_texture, 0) + self.screen_texture.append(screen_texture) + + glDrawBuffers(color_size, attachments) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + else: + self.color_buffer = [] + for i in range(color_size): + color_buffer = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, color_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_NEAREST) + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, + self.height, 0, GL_RGBA, GL_FLOAT, None) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, + color_buffer, 0) + self.color_buffer.append(color_buffer) + + # Configure depth texture map to render to + self.depth_buffer = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, self.depth_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_DEPTH_TEXTURE_MODE, GL_INTENSITY) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, + GL_COMPARE_R_TO_TEXTURE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL) + glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.width, + self.height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None) + glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, + GL_TEXTURE_2D, self.depth_buffer, 0) + + attachments = [] + for i in range(color_size): + attachments.append(GL_COLOR_ATTACHMENT0 + i) + glDrawBuffers(color_size, attachments) + self.screen_texture = self.color_buffer + + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + # Configure texture buffer if needed + self.render_texture = None + + # NOTE: original render_texture only support one input + # this is tentative member of this issue + self.render_texture_v2 = {} + + # Inner storage for buffer data + self.vertex_data = None + self.vertex_dim = None + self.n_vertices = None + + self.model_view_matrix = None + self.projection_matrix = None + + if not egl: + global GLUT + import OpenGL.GLUT as GLUT + GLUT.glutDisplayFunc(self.display) + + def init_quad_program(self): + shader_list = [] + + shader_list.append(loadShader(GL_VERTEX_SHADER, "quad.vs")) + shader_list.append(loadShader(GL_FRAGMENT_SHADER, "quad.fs")) + + the_program = createProgram(shader_list) + + for shader in shader_list: + glDeleteShader(shader) + + # vertex attributes for a quad that fills the entire screen in Normalized Device Coordinates. + # positions # texCoords + quad_vertices = np.array([ + -1.0, 1.0, 0.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0, 1.0, 0.0, + -1.0, 1.0, 0.0, 1.0, 1.0, -1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0 + ]) + + quad_buffer = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, quad_buffer) + glBufferData(GL_ARRAY_BUFFER, quad_vertices, GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + return the_program, quad_buffer + + def set_mesh(self, vertices, faces): + self.vertex_data = vertices[faces.reshape([-1])] + self.vertex_dim = self.vertex_data.shape[1] + self.n_vertices = self.vertex_data.shape[0] + + glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) + glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def set_viewpoint(self, projection, model_view): + self.projection_matrix = projection + self.model_view_matrix = model_view + + def draw_init(self): + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + glEnable(GL_DEPTH_TEST) + + glClearColor(0.0, 0.0, 0.0, 0.0) + if self.use_inverse_depth: + glDepthFunc(GL_GREATER) + glClearDepth(0.0) + else: + glDepthFunc(GL_LESS) + glClearDepth(1.0) + glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) + + def draw_end(self): + if self.intermediate_fbo is not None: + for i in range(len(self.color_buffer)): + glBindFramebuffer(GL_READ_FRAMEBUFFER, self.frame_buffer) + glReadBuffer(GL_COLOR_ATTACHMENT0 + i) + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) + glDrawBuffer(GL_COLOR_ATTACHMENT0 + i) + glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, + self.width, self.height, GL_COLOR_BUFFER_BIT, + GL_NEAREST) + + glBindFramebuffer(GL_FRAMEBUFFER, 0) + glDepthFunc(GL_LESS) + glClearDepth(1.0) + + def draw(self): + self.draw_init() + + glUseProgram(self.program) + glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, + self.model_view_matrix.transpose()) + glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, + self.projection_matrix.transpose()) + + glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) + + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) + + glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) + + glDisableVertexAttribArray(0) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + glUseProgram(0) + + self.draw_end() + + def get_color(self, color_id=0): + glBindFramebuffer( + GL_FRAMEBUFFER, self.intermediate_fbo + if self.intermediate_fbo is not None else self.frame_buffer) + glReadBuffer(GL_COLOR_ATTACHMENT0 + color_id) + data = glReadPixels(0, + 0, + self.width, + self.height, + GL_RGBA, + GL_FLOAT, + outputType=None) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + rgb = data.reshape(self.height, self.width, -1) + rgb = np.flip(rgb, 0) + return rgb + + def get_z_value(self): + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + data = glReadPixels(0, + 0, + self.width, + self.height, + GL_DEPTH_COMPONENT, + GL_FLOAT, + outputType=None) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + z = data.reshape(self.height, self.width) + z = np.flip(z, 0) + return z + + def display(self): + self.draw() + + if not self.egl: + # First we draw a scene. + # Notice the result is stored in the texture buffer. + + # Then we return to the default frame buffer since we will display on the screen. + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + # Do the clean-up. + glClearColor(0.0, 0.0, 0.0, 0.0) + glClear(GL_COLOR_BUFFER_BIT) + + # We draw a rectangle which covers the whole screen. + glUseProgram(self.quad_program) + glBindBuffer(GL_ARRAY_BUFFER, self.quad_buffer) + + size_of_double = 8 + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, 2, GL_DOUBLE, GL_FALSE, + 4 * size_of_double, None) + glEnableVertexAttribArray(1) + glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, + 4 * size_of_double, + c_void_p(2 * size_of_double)) + + glDisable(GL_DEPTH_TEST) + + # The stored texture is then mapped to this rectangle. + # properly assing color buffer texture + glActiveTexture(GL_TEXTURE0) + glBindTexture(GL_TEXTURE_2D, self.screen_texture[0]) + glUniform1i( + glGetUniformLocation(self.quad_program, 'screenTexture'), 0) + + glDrawArrays(GL_TRIANGLES, 0, 6) + + glDisableVertexAttribArray(1) + glDisableVertexAttribArray(0) + + glEnable(GL_DEPTH_TEST) + glBindBuffer(GL_ARRAY_BUFFER, 0) + glUseProgram(0) + + GLUT.glutSwapBuffers() + GLUT.glutPostRedisplay() + + def show(self): + if not self.egl: + GLUT.glutMainLoop() diff --git a/lib/renderer/gl/render2.py b/lib/renderer/gl/render2.py new file mode 100755 index 0000000000000000000000000000000000000000..b7f38fc80ebe6e7a07cedcdd90206c2255172429 --- /dev/null +++ b/lib/renderer/gl/render2.py @@ -0,0 +1,388 @@ +''' +MIT License + +Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +''' +import numpy as np +from OpenGL.GLUT import * +from .framework import * + +_glut_window = None + + +class Render: + def __init__(self, + width=1600, + height=1200, + name='GL Renderer', + program_files=['simple.fs', 'simple.vs'], + color_size=1, + ms_rate=1): + self.width = width + self.height = height + self.name = name + self.display_mode = GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH + self.use_inverse_depth = False + + global _glut_window + if _glut_window is None: + glutInit() + glutInitDisplayMode(self.display_mode) + glutInitWindowSize(self.width, self.height) + glutInitWindowPosition(0, 0) + _glut_window = glutCreateWindow("My Render.") + + # glEnable(GL_DEPTH_CLAMP) + glEnable(GL_DEPTH_TEST) + + glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE) + glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE) + glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE) + + # init program + shader_list = [] + + for program_file in program_files: + _, ext = os.path.splitext(program_file) + if ext == '.vs': + shader_list.append(loadShader(GL_VERTEX_SHADER, program_file)) + elif ext == '.fs': + shader_list.append(loadShader(GL_FRAGMENT_SHADER, + program_file)) + elif ext == '.gs': + shader_list.append(loadShader(GL_GEOMETRY_SHADER, + program_file)) + + self.program = createProgram(shader_list) + + for shader in shader_list: + glDeleteShader(shader) + + # Init uniform variables + self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat') + self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat') + + self.vertex_buffer = glGenBuffers(1) + + # Init screen quad program and buffer + self.quad_program, self.quad_buffer = self.init_quad_program() + + # Configure frame buffer + self.frame_buffer = glGenFramebuffers(1) + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + + self.intermediate_fbo = None + if ms_rate > 1: + # Configure texture buffer to render to + self.color_buffer = [] + for i in range(color_size): + color_buffer = glGenTextures(1) + multi_sample_rate = ms_rate + glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, color_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_LINEAR) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_LINEAR) + glTexImage2DMultisample(GL_TEXTURE_2D_MULTISAMPLE, + multi_sample_rate, GL_RGBA32F, + self.width, self.height, GL_TRUE) + glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, 0) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, + GL_TEXTURE_2D_MULTISAMPLE, color_buffer, + 0) + self.color_buffer.append(color_buffer) + + self.render_buffer = glGenRenderbuffers(1) + glBindRenderbuffer(GL_RENDERBUFFER, self.render_buffer) + glRenderbufferStorageMultisample(GL_RENDERBUFFER, + multi_sample_rate, + GL_DEPTH24_STENCIL8, self.width, + self.height) + glBindRenderbuffer(GL_RENDERBUFFER, 0) + glFramebufferRenderbuffer(GL_FRAMEBUFFER, + GL_DEPTH_STENCIL_ATTACHMENT, + GL_RENDERBUFFER, self.render_buffer) + + attachments = [] + for i in range(color_size): + attachments.append(GL_COLOR_ATTACHMENT0 + i) + glDrawBuffers(color_size, attachments) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + self.intermediate_fbo = glGenFramebuffers(1) + glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo) + + self.screen_texture = [] + for i in range(color_size): + screen_texture = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, screen_texture) + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, + self.height, 0, GL_RGBA, GL_FLOAT, None) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_LINEAR) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_LINEAR) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, + screen_texture, 0) + self.screen_texture.append(screen_texture) + + glDrawBuffers(color_size, attachments) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + else: + self.color_buffer = [] + for i in range(color_size): + color_buffer = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, color_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, + GL_CLAMP_TO_EDGE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, + GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, + GL_NEAREST) + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, + self.height, 0, GL_RGBA, GL_FLOAT, None) + glFramebufferTexture2D(GL_FRAMEBUFFER, + GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, + color_buffer, 0) + self.color_buffer.append(color_buffer) + + # Configure depth texture map to render to + self.depth_buffer = glGenTextures(1) + glBindTexture(GL_TEXTURE_2D, self.depth_buffer) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) + glTexParameteri(GL_TEXTURE_2D, GL_DEPTH_TEXTURE_MODE, GL_INTENSITY) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, + GL_COMPARE_R_TO_TEXTURE) + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL) + glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.width, + self.height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None) + glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, + GL_TEXTURE_2D, self.depth_buffer, 0) + + attachments = [] + for i in range(color_size): + attachments.append(GL_COLOR_ATTACHMENT0 + i) + glDrawBuffers(color_size, attachments) + self.screen_texture = self.color_buffer + + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + # Configure texture buffer if needed + self.render_texture = None + + # NOTE: original render_texture only support one input + # this is tentative member of this issue + self.render_texture_v2 = {} + + # Inner storage for buffer data + self.vertex_data = None + self.vertex_dim = None + self.n_vertices = None + + self.model_view_matrix = None + self.projection_matrix = None + + glutDisplayFunc(self.display) + + def init_quad_program(self): + shader_list = [] + + shader_list.append(loadShader(GL_VERTEX_SHADER, "quad.vs")) + shader_list.append(loadShader(GL_FRAGMENT_SHADER, "quad.fs")) + + the_program = createProgram(shader_list) + + for shader in shader_list: + glDeleteShader(shader) + + # vertex attributes for a quad that fills the entire screen in Normalized Device Coordinates. + # positions # texCoords + quad_vertices = np.array([ + -1.0, 1.0, 0.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0, 1.0, 0.0, + -1.0, 1.0, 0.0, 1.0, 1.0, -1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0 + ]) + + quad_buffer = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, quad_buffer) + glBufferData(GL_ARRAY_BUFFER, quad_vertices, GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + return the_program, quad_buffer + + def set_mesh(self, vertices, faces): + self.vertex_data = vertices[faces.reshape([-1])] + self.vertex_dim = self.vertex_data.shape[1] + self.n_vertices = self.vertex_data.shape[0] + + glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) + glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + def set_viewpoint(self, projection, model_view): + self.projection_matrix = projection + self.model_view_matrix = model_view + + def draw_init(self): + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + glEnable(GL_DEPTH_TEST) + + # glClearColor(0.0, 0.0, 0.0, 0.0) + glClearColor(1.0, 1.0, 1.0, 0.0) # Black background + + if self.use_inverse_depth: + glDepthFunc(GL_GREATER) + glClearDepth(0.0) + else: + glDepthFunc(GL_LESS) + glClearDepth(1.0) + glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) + + def draw_end(self): + if self.intermediate_fbo is not None: + for i in range(len(self.color_buffer)): + glBindFramebuffer(GL_READ_FRAMEBUFFER, self.frame_buffer) + glReadBuffer(GL_COLOR_ATTACHMENT0 + i) + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) + glDrawBuffer(GL_COLOR_ATTACHMENT0 + i) + glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, + self.width, self.height, GL_COLOR_BUFFER_BIT, + GL_NEAREST) + + glBindFramebuffer(GL_FRAMEBUFFER, 0) + glDepthFunc(GL_LESS) + glClearDepth(1.0) + + def draw(self): + self.draw_init() + + glUseProgram(self.program) + glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, + self.model_view_matrix.transpose()) + glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, + self.projection_matrix.transpose()) + + glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) + + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) + + glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) + + glDisableVertexAttribArray(0) + + glBindBuffer(GL_ARRAY_BUFFER, 0) + + glUseProgram(0) + + self.draw_end() + + def get_color(self, color_id=0): + glBindFramebuffer( + GL_FRAMEBUFFER, self.intermediate_fbo + if self.intermediate_fbo is not None else self.frame_buffer) + glReadBuffer(GL_COLOR_ATTACHMENT0 + color_id) + data = glReadPixels(0, + 0, + self.width, + self.height, + GL_RGBA, + GL_FLOAT, + outputType=None) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + rgb = data.reshape(self.height, self.width, -1) + rgb = np.flip(rgb, 0) + return rgb + + def get_z_value(self): + glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) + data = glReadPixels(0, + 0, + self.width, + self.height, + GL_DEPTH_COMPONENT, + GL_FLOAT, + outputType=None) + glBindFramebuffer(GL_FRAMEBUFFER, 0) + z = data.reshape(self.height, self.width) + z = np.flip(z, 0) + return z + + def display(self): + # First we draw a scene. + # Notice the result is stored in the texture buffer. + self.draw() + + # Then we return to the default frame buffer since we will display on the screen. + glBindFramebuffer(GL_FRAMEBUFFER, 0) + + # Do the clean-up. + # glClearColor(0.0, 0.0, 0.0, 0.0) #Black background + glClearColor(1.0, 1.0, 1.0, 0.0) # Black background + glClear(GL_COLOR_BUFFER_BIT) + + # We draw a rectangle which covers the whole screen. + glUseProgram(self.quad_program) + glBindBuffer(GL_ARRAY_BUFFER, self.quad_buffer) + + size_of_double = 8 + glEnableVertexAttribArray(0) + glVertexAttribPointer(0, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, + None) + glEnableVertexAttribArray(1) + glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, + c_void_p(2 * size_of_double)) + + glDisable(GL_DEPTH_TEST) + + # The stored texture is then mapped to this rectangle. + # properly assing color buffer texture + glActiveTexture(GL_TEXTURE0) + glBindTexture(GL_TEXTURE_2D, self.screen_texture[0]) + glUniform1i(glGetUniformLocation(self.quad_program, 'screenTexture'), + 0) + + glDrawArrays(GL_TRIANGLES, 0, 6) + + glDisableVertexAttribArray(1) + glDisableVertexAttribArray(0) + + glEnable(GL_DEPTH_TEST) + glBindBuffer(GL_ARRAY_BUFFER, 0) + glUseProgram(0) + + glutSwapBuffers() + glutPostRedisplay() + + def show(self): + glutMainLoop() diff --git a/lib/renderer/glm.py b/lib/renderer/glm.py new file mode 100755 index 0000000000000000000000000000000000000000..65b1407e6edba36ac5883166a45bbe3ad8fadcce --- /dev/null +++ b/lib/renderer/glm.py @@ -0,0 +1,143 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import numpy as np + + +def vec3(x, y, z): + return np.array([x, y, z], dtype=np.float32) + + +def radians(v): + return np.radians(v) + + +def identity(): + return np.identity(4, dtype=np.float32) + + +def empty(): + return np.zeros([4, 4], dtype=np.float32) + + +def magnitude(v): + return np.linalg.norm(v) + + +def normalize(v): + m = magnitude(v) + return v if m == 0 else v / m + + +def dot(u, v): + return np.sum(u * v) + + +def cross(u, v): + res = vec3(0, 0, 0) + res[0] = u[1] * v[2] - u[2] * v[1] + res[1] = u[2] * v[0] - u[0] * v[2] + res[2] = u[0] * v[1] - u[1] * v[0] + return res + + +# below functions can be optimized + + +def translate(m, v): + res = np.copy(m) + res[:, 3] = m[:, 0] * v[0] + m[:, 1] * v[1] + m[:, 2] * v[2] + m[:, 3] + return res + + +def rotate(m, angle, v): + a = angle + c = np.cos(a) + s = np.sin(a) + + axis = normalize(v) + temp = (1 - c) * axis + + rot = empty() + rot[0][0] = c + temp[0] * axis[0] + rot[0][1] = temp[0] * axis[1] + s * axis[2] + rot[0][2] = temp[0] * axis[2] - s * axis[1] + + rot[1][0] = temp[1] * axis[0] - s * axis[2] + rot[1][1] = c + temp[1] * axis[1] + rot[1][2] = temp[1] * axis[2] + s * axis[0] + + rot[2][0] = temp[2] * axis[0] + s * axis[1] + rot[2][1] = temp[2] * axis[1] - s * axis[0] + rot[2][2] = c + temp[2] * axis[2] + + res = empty() + res[:, 0] = m[:, 0] * rot[0][0] + m[:, 1] * rot[0][1] + m[:, 2] * rot[0][2] + res[:, 1] = m[:, 0] * rot[1][0] + m[:, 1] * rot[1][1] + m[:, 2] * rot[1][2] + res[:, 2] = m[:, 0] * rot[2][0] + m[:, 1] * rot[2][1] + m[:, 2] * rot[2][2] + res[:, 3] = m[:, 3] + return res + + +def perspective(fovy, aspect, zNear, zFar): + tanHalfFovy = np.tan(fovy / 2) + + res = empty() + res[0][0] = 1 / (aspect * tanHalfFovy) + res[1][1] = 1 / (tanHalfFovy) + res[2][3] = -1 + res[2][2] = -(zFar + zNear) / (zFar - zNear) + res[3][2] = -(2 * zFar * zNear) / (zFar - zNear) + + return res.T + + +def ortho(left, right, bottom, top, zNear, zFar): + # res = np.ones([4, 4], dtype=np.float32) + res = identity() + res[0][0] = 2 / (right - left) + res[1][1] = 2 / (top - bottom) + res[2][2] = -2 / (zFar - zNear) + res[3][0] = -(right + left) / (right - left) + res[3][1] = -(top + bottom) / (top - bottom) + res[3][2] = -(zFar + zNear) / (zFar - zNear) + return res.T + + +def lookat(eye, center, up): + f = normalize(center - eye) + s = normalize(cross(f, up)) + u = cross(s, f) + + res = identity() + res[0][0] = s[0] + res[1][0] = s[1] + res[2][0] = s[2] + res[0][1] = u[0] + res[1][1] = u[1] + res[2][1] = u[2] + res[0][2] = -f[0] + res[1][2] = -f[1] + res[2][2] = -f[2] + res[3][0] = -dot(s, eye) + res[3][1] = -dot(u, eye) + res[3][2] = -dot(f, eye) + return res.T + + +def transform(d, m): + return np.dot(m, d.T).T diff --git a/lib/renderer/mesh.py b/lib/renderer/mesh.py new file mode 100755 index 0000000000000000000000000000000000000000..437467f4e0ebe93bb18b4bf5d91df73db59328b1 --- /dev/null +++ b/lib/renderer/mesh.py @@ -0,0 +1,526 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.dataset.mesh_util import SMPLX +from lib.common.render_utils import face_vertices +import numpy as np +import smplx +import trimesh +import torch +import torch.nn.functional as F + +model_init_params = dict( + gender='male', + model_type='smplx', + model_path=SMPLX().model_dir, + create_global_orient=False, + create_body_pose=False, + create_betas=False, + create_left_hand_pose=False, + create_right_hand_pose=False, + create_expression=False, + create_jaw_pose=False, + create_leye_pose=False, + create_reye_pose=False, + create_transl=False, + num_pca_comps=12) + + +def get_smpl_model(model_type, gender): return smplx.create( + **model_init_params) + + +def normalization(data): + _range = np.max(data) - np.min(data) + return ((data - np.min(data)) / _range) + + +def sigmoid(x): + z = 1 / (1 + np.exp(-x)) + return z + + +def load_fit_body(fitted_path, scale, smpl_type='smplx', smpl_gender='neutral', noise_dict=None): + + param = np.load(fitted_path, allow_pickle=True) + for key in param.keys(): + param[key] = torch.as_tensor(param[key]) + + smpl_model = get_smpl_model(smpl_type, smpl_gender) + model_forward_params = dict(betas=param['betas'], + global_orient=param['global_orient'], + body_pose=param['body_pose'], + left_hand_pose=param['left_hand_pose'], + right_hand_pose=param['right_hand_pose'], + jaw_pose=param['jaw_pose'], + leye_pose=param['leye_pose'], + reye_pose=param['reye_pose'], + expression=param['expression'], + return_verts=True) + + if noise_dict is not None: + model_forward_params.update(noise_dict) + + smpl_out = smpl_model(**model_forward_params) + + smpl_verts = ( + (smpl_out.vertices[0] * param['scale'] + param['translation']) * scale).detach() + smpl_joints = ( + (smpl_out.joints[0] * param['scale'] + param['translation']) * scale).detach() + smpl_mesh = trimesh.Trimesh(smpl_verts, + smpl_model.faces, + process=False, maintain_order=True) + + return smpl_mesh, smpl_joints + + +def load_ori_fit_body(fitted_path, smpl_type='smplx', smpl_gender='neutral'): + + param = np.load(fitted_path, allow_pickle=True) + for key in param.keys(): + param[key] = torch.as_tensor(param[key]) + + smpl_model = get_smpl_model(smpl_type, smpl_gender) + model_forward_params = dict(betas=param['betas'], + global_orient=param['global_orient'], + body_pose=param['body_pose'], + left_hand_pose=param['left_hand_pose'], + right_hand_pose=param['right_hand_pose'], + jaw_pose=param['jaw_pose'], + leye_pose=param['leye_pose'], + reye_pose=param['reye_pose'], + expression=param['expression'], + return_verts=True) + + smpl_out = smpl_model(**model_forward_params) + + smpl_verts = smpl_out.vertices[0].detach() + smpl_mesh = trimesh.Trimesh(smpl_verts, + smpl_model.faces, + process=False, maintain_order=True) + + return smpl_mesh + + +def save_obj_mesh(mesh_path, verts, faces): + file = open(mesh_path, 'w') + for v in verts: + file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) + for f in faces: + f_plus = f + 1 + file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) + file.close() + + +# https://github.com/ratcave/wavefront_reader +def read_mtlfile(fname): + materials = {} + with open(fname) as f: + lines = f.read().splitlines() + + for line in lines: + if line: + split_line = line.strip().split(' ', 1) + if len(split_line) < 2: + continue + + prefix, data = split_line[0], split_line[1] + if 'newmtl' in prefix: + material = {} + materials[data] = material + elif materials: + if data: + split_data = data.strip().split(' ') + + # assume texture maps are in the same level + # WARNING: do not include space in your filename!! + if 'map' in prefix: + material[prefix] = split_data[-1].split('\\')[-1] + elif len(split_data) > 1: + material[prefix] = tuple(float(d) for d in split_data) + else: + try: + material[prefix] = int(data) + except ValueError: + material[prefix] = float(data) + + return materials + + +def load_obj_mesh_mtl(mesh_file): + vertex_data = [] + norm_data = [] + uv_data = [] + + face_data = [] + face_norm_data = [] + face_uv_data = [] + + # face per material + face_data_mat = {} + face_norm_data_mat = {} + face_uv_data_mat = {} + + # current material name + mtl_data = None + cur_mat = None + + if isinstance(mesh_file, str): + f = open(mesh_file, "r") + else: + f = mesh_file + for line in f: + if isinstance(line, bytes): + line = line.decode("utf-8") + if line.startswith('#'): + continue + values = line.split() + if not values: + continue + + if values[0] == 'v': + v = list(map(float, values[1:4])) + vertex_data.append(v) + elif values[0] == 'vn': + vn = list(map(float, values[1:4])) + norm_data.append(vn) + elif values[0] == 'vt': + vt = list(map(float, values[1:3])) + uv_data.append(vt) + elif values[0] == 'mtllib': + mtl_data = read_mtlfile( + mesh_file.replace(mesh_file.split('/')[-1], values[1])) + elif values[0] == 'usemtl': + cur_mat = values[1] + elif values[0] == 'f': + # local triangle data + l_face_data = [] + l_face_uv_data = [] + l_face_norm_data = [] + + # quad mesh + if len(values) > 4: + f = list( + map( + lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) + < 0 else int(x.split('/')[0]) - 1, values[1:4])) + l_face_data.append(f) + f = list( + map( + lambda x: int(x.split('/')[0]) + if int(x.split('/')[0]) < 0 else int(x.split('/')[0]) - + 1, [values[3], values[4], values[1]])) + l_face_data.append(f) + # tri mesh + else: + f = list( + map( + lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) + < 0 else int(x.split('/')[0]) - 1, values[1:4])) + l_face_data.append(f) + # deal with texture + if len(values[1].split('/')) >= 2: + # quad mesh + if len(values) > 4: + f = list( + map( + lambda x: int(x.split('/')[1]) + if int(x.split('/')[1]) < 0 else int( + x.split('/')[1]) - 1, values[1:4])) + l_face_uv_data.append(f) + f = list( + map( + lambda x: int(x.split('/')[1]) + if int(x.split('/')[1]) < 0 else int( + x.split('/')[1]) - 1, + [values[3], values[4], values[1]])) + l_face_uv_data.append(f) + # tri mesh + elif len(values[1].split('/')[1]) != 0: + f = list( + map( + lambda x: int(x.split('/')[1]) + if int(x.split('/')[1]) < 0 else int( + x.split('/')[1]) - 1, values[1:4])) + l_face_uv_data.append(f) + # deal with normal + if len(values[1].split('/')) == 3: + # quad mesh + if len(values) > 4: + f = list( + map( + lambda x: int(x.split('/')[2]) + if int(x.split('/')[2]) < 0 else int( + x.split('/')[2]) - 1, values[1:4])) + l_face_norm_data.append(f) + f = list( + map( + lambda x: int(x.split('/')[2]) + if int(x.split('/')[2]) < 0 else int( + x.split('/')[2]) - 1, + [values[3], values[4], values[1]])) + l_face_norm_data.append(f) + # tri mesh + elif len(values[1].split('/')[2]) != 0: + f = list( + map( + lambda x: int(x.split('/')[2]) + if int(x.split('/')[2]) < 0 else int( + x.split('/')[2]) - 1, values[1:4])) + l_face_norm_data.append(f) + + face_data += l_face_data + face_uv_data += l_face_uv_data + face_norm_data += l_face_norm_data + + if cur_mat is not None: + if cur_mat not in face_data_mat.keys(): + face_data_mat[cur_mat] = [] + if cur_mat not in face_uv_data_mat.keys(): + face_uv_data_mat[cur_mat] = [] + if cur_mat not in face_norm_data_mat.keys(): + face_norm_data_mat[cur_mat] = [] + face_data_mat[cur_mat] += l_face_data + face_uv_data_mat[cur_mat] += l_face_uv_data + face_norm_data_mat[cur_mat] += l_face_norm_data + + vertices = np.array(vertex_data) + faces = np.array(face_data) + + norms = np.array(norm_data) + norms = normalize_v3(norms) + face_normals = np.array(face_norm_data) + + uvs = np.array(uv_data) + face_uvs = np.array(face_uv_data) + + out_tuple = (vertices, faces, norms, face_normals, uvs, face_uvs) + + if cur_mat is not None and mtl_data is not None: + for key in face_data_mat: + face_data_mat[key] = np.array(face_data_mat[key]) + face_uv_data_mat[key] = np.array(face_uv_data_mat[key]) + face_norm_data_mat[key] = np.array(face_norm_data_mat[key]) + + out_tuple += (face_data_mat, face_norm_data_mat, face_uv_data_mat, + mtl_data) + + return out_tuple + + +def load_scan(mesh_file, with_normal=False, with_texture=False): + vertex_data = [] + norm_data = [] + uv_data = [] + + face_data = [] + face_norm_data = [] + face_uv_data = [] + + if isinstance(mesh_file, str): + f = open(mesh_file, "r") + else: + f = mesh_file + for line in f: + if isinstance(line, bytes): + line = line.decode("utf-8") + if line.startswith('#'): + continue + values = line.split() + if not values: + continue + + if values[0] == 'v': + v = list(map(float, values[1:4])) + vertex_data.append(v) + elif values[0] == 'vn': + vn = list(map(float, values[1:4])) + norm_data.append(vn) + elif values[0] == 'vt': + vt = list(map(float, values[1:3])) + uv_data.append(vt) + + elif values[0] == 'f': + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + f = list( + map(lambda x: int(x.split('/')[0]), + [values[3], values[4], values[1]])) + face_data.append(f) + # tri mesh + else: + f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) + face_data.append(f) + + # deal with texture + if len(values[1].split('/')) >= 2: + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) + face_uv_data.append(f) + f = list( + map(lambda x: int(x.split('/')[1]), + [values[3], values[4], values[1]])) + face_uv_data.append(f) + # tri mesh + elif len(values[1].split('/')[1]) != 0: + f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) + face_uv_data.append(f) + # deal with normal + if len(values[1].split('/')) == 3: + # quad mesh + if len(values) > 4: + f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) + face_norm_data.append(f) + f = list( + map(lambda x: int(x.split('/')[2]), + [values[3], values[4], values[1]])) + face_norm_data.append(f) + # tri mesh + elif len(values[1].split('/')[2]) != 0: + f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) + face_norm_data.append(f) + + vertices = np.array(vertex_data) + faces = np.array(face_data) - 1 + + if with_texture and with_normal: + uvs = np.array(uv_data) + face_uvs = np.array(face_uv_data) - 1 + norms = np.array(norm_data) + if norms.shape[0] == 0: + norms = compute_normal(vertices, faces) + face_normals = faces + else: + norms = normalize_v3(norms) + face_normals = np.array(face_norm_data) - 1 + return vertices, faces, norms, face_normals, uvs, face_uvs + + if with_texture: + uvs = np.array(uv_data) + face_uvs = np.array(face_uv_data) - 1 + return vertices, faces, uvs, face_uvs + + if with_normal: + norms = np.array(norm_data) + norms = normalize_v3(norms) + face_normals = np.array(face_norm_data) - 1 + return vertices, faces, norms, face_normals + + return vertices, faces + + +def normalize_v3(arr): + ''' Normalize a numpy array of 3 component vectors shape=(n,3) ''' + lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2) + eps = 0.00000001 + lens[lens < eps] = eps + arr[:, 0] /= lens + arr[:, 1] /= lens + arr[:, 2] /= lens + return arr + + +def compute_normal(vertices, faces): + # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal + norm = np.zeros(vertices.shape, dtype=vertices.dtype) + # Create an indexed view into the vertex array using the array of three indices for triangles + tris = vertices[faces] + # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle + n = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0]) + # n is now an array of normals per triangle. The length of each normal is dependent the vertices, + # we need to normalize these, so that our next step weights each normal equally. + normalize_v3(n) + # now we have a normalized array of normals, one per triangle, i.e., per triangle normals. + # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle, + # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards. + # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array + norm[faces[:, 0]] += n + norm[faces[:, 1]] += n + norm[faces[:, 2]] += n + normalize_v3(norm) + + return norm + + +def compute_normal_batch(vertices, faces): + + bs, nv = vertices.shape[:2] + bs, nf = faces.shape[:2] + + vert_norm = torch.zeros(bs * nv, 3).type_as(vertices) + tris = face_vertices(vertices, faces) + face_norm = F.normalize(torch.cross(tris[:, :, 1] - tris[:, :, 0], + tris[:, :, 2] - tris[:, :, 0]), + dim=-1) + + faces = (faces + + (torch.arange(bs).type_as(faces) * nv)[:, None, None]).view( + -1, 3) + + vert_norm[faces[:, 0]] += face_norm.view(-1, 3) + vert_norm[faces[:, 1]] += face_norm.view(-1, 3) + vert_norm[faces[:, 2]] += face_norm.view(-1, 3) + + vert_norm = F.normalize(vert_norm, dim=-1).view(bs, nv, 3) + + return vert_norm + + +# compute tangent and bitangent +def compute_tangent(vertices, faces, normals, uvs, faceuvs): + # NOTE: this could be numerically unstable around [0,0,1] + # but other current solutions are pretty freaky somehow + c1 = np.cross(normals, np.array([0, 1, 0.0])) + tan = c1 + normalize_v3(tan) + btan = np.cross(normals, tan) + + # NOTE: traditional version is below + + # pts_tris = vertices[faces] + # uv_tris = uvs[faceuvs] + + # W = np.stack([pts_tris[::, 1] - pts_tris[::, 0], pts_tris[::, 2] - pts_tris[::, 0]],2) + # UV = np.stack([uv_tris[::, 1] - uv_tris[::, 0], uv_tris[::, 2] - uv_tris[::, 0]], 1) + + # for i in range(W.shape[0]): + # W[i,::] = W[i,::].dot(np.linalg.inv(UV[i,::])) + + # tan = np.zeros(vertices.shape, dtype=vertices.dtype) + # tan[faces[:,0]] += W[:,:,0] + # tan[faces[:,1]] += W[:,:,0] + # tan[faces[:,2]] += W[:,:,0] + + # btan = np.zeros(vertices.shape, dtype=vertices.dtype) + # btan[faces[:,0]] += W[:,:,1] + # btan[faces[:,1]] += W[:,:,1] + # btan[faces[:,2]] += W[:,:,1] + + # normalize_v3(tan) + + # ndott = np.sum(normals*tan, 1, keepdims=True) + # tan = tan - ndott * normals + + # normalize_v3(btan) + # normalize_v3(tan) + + # tan[np.sum(np.cross(normals, tan) * btan, 1) < 0,:] *= -1.0 + + return tan, btan diff --git a/lib/renderer/opengl_util.py b/lib/renderer/opengl_util.py new file mode 100644 index 0000000000000000000000000000000000000000..a9f8801d8e6e47ee5591f1a860dbc9ced79ba4c1 --- /dev/null +++ b/lib/renderer/opengl_util.py @@ -0,0 +1,370 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +from lib.renderer.mesh import load_scan, compute_tangent, compute_normal, load_obj_mesh_mtl +from lib.dataset.mesh_util import projection +from lib.renderer.gl.prt_render import PRTRender +from lib.renderer.camera import Camera +import os +import cv2 +import math +import random +import numpy as np + + +def render_result(rndr, shader_id, path, mask=False): + + cam_render = rndr.get_color(shader_id) + cam_render = cv2.cvtColor(cam_render, cv2.COLOR_RGBA2BGRA) + + os.makedirs(os.path.dirname(path), exist_ok=True) + if shader_id != 2: + cv2.imwrite(path, np.uint8(255.0 * cam_render)) + else: + cam_render[:, :, -1] -= 0.5 + cam_render[:, :, -1] *= 2.0 + if not mask: + cv2.imwrite(path, np.uint8(255.0 / 2.0 * (cam_render + 1.0))) + else: + cv2.imwrite(path, np.uint8(-1.0 * cam_render[:, :, [3]])) + + +def make_rotate(rx, ry, rz): + sinX = np.sin(rx) + sinY = np.sin(ry) + sinZ = np.sin(rz) + + cosX = np.cos(rx) + cosY = np.cos(ry) + cosZ = np.cos(rz) + + Rx = np.zeros((3, 3)) + Rx[0, 0] = 1.0 + Rx[1, 1] = cosX + Rx[1, 2] = -sinX + Rx[2, 1] = sinX + Rx[2, 2] = cosX + + Ry = np.zeros((3, 3)) + Ry[0, 0] = cosY + Ry[0, 2] = sinY + Ry[1, 1] = 1.0 + Ry[2, 0] = -sinY + Ry[2, 2] = cosY + + Rz = np.zeros((3, 3)) + Rz[0, 0] = cosZ + Rz[0, 1] = -sinZ + Rz[1, 0] = sinZ + Rz[1, 1] = cosZ + Rz[2, 2] = 1.0 + + R = np.matmul(np.matmul(Rz, Ry), Rx) + return R + + +def rotateSH(SH, R): + SHn = SH + + # 1st order + SHn[1] = R[1, 1] * SH[1] - R[1, 2] * SH[2] + R[1, 0] * SH[3] + SHn[2] = -R[2, 1] * SH[1] + R[2, 2] * SH[2] - R[2, 0] * SH[3] + SHn[3] = R[0, 1] * SH[1] - R[0, 2] * SH[2] + R[0, 0] * SH[3] + + # 2nd order + SHn[4:, 0] = rotateBand2(SH[4:, 0], R) + SHn[4:, 1] = rotateBand2(SH[4:, 1], R) + SHn[4:, 2] = rotateBand2(SH[4:, 2], R) + + return SHn + + +def rotateBand2(x, R): + s_c3 = 0.94617469575 + s_c4 = -0.31539156525 + s_c5 = 0.54627421529 + + s_c_scale = 1.0 / 0.91529123286551084 + s_c_scale_inv = 0.91529123286551084 + + s_rc2 = 1.5853309190550713 * s_c_scale + s_c4_div_c3 = s_c4 / s_c3 + s_c4_div_c3_x2 = (s_c4 / s_c3) * 2.0 + + s_scale_dst2 = s_c3 * s_c_scale_inv + s_scale_dst4 = s_c5 * s_c_scale_inv + + sh0 = x[3] + x[4] + x[4] - x[1] + sh1 = x[0] + s_rc2 * x[2] + x[3] + x[4] + sh2 = x[0] + sh3 = -x[3] + sh4 = -x[1] + + r2x = R[0][0] + R[0][1] + r2y = R[1][0] + R[1][1] + r2z = R[2][0] + R[2][1] + + r3x = R[0][0] + R[0][2] + r3y = R[1][0] + R[1][2] + r3z = R[2][0] + R[2][2] + + r4x = R[0][1] + R[0][2] + r4y = R[1][1] + R[1][2] + r4z = R[2][1] + R[2][2] + + sh0_x = sh0 * R[0][0] + sh0_y = sh0 * R[1][0] + d0 = sh0_x * R[1][0] + d1 = sh0_y * R[2][0] + d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3) + d3 = sh0_x * R[2][0] + d4 = sh0_x * R[0][0] - sh0_y * R[1][0] + + sh1_x = sh1 * R[0][2] + sh1_y = sh1 * R[1][2] + d0 += sh1_x * R[1][2] + d1 += sh1_y * R[2][2] + d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3) + d3 += sh1_x * R[2][2] + d4 += sh1_x * R[0][2] - sh1_y * R[1][2] + + sh2_x = sh2 * r2x + sh2_y = sh2 * r2y + d0 += sh2_x * r2y + d1 += sh2_y * r2z + d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2) + d3 += sh2_x * r2z + d4 += sh2_x * r2x - sh2_y * r2y + + sh3_x = sh3 * r3x + sh3_y = sh3 * r3y + d0 += sh3_x * r3y + d1 += sh3_y * r3z + d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2) + d3 += sh3_x * r3z + d4 += sh3_x * r3x - sh3_y * r3y + + sh4_x = sh4 * r4x + sh4_y = sh4 * r4y + d0 += sh4_x * r4y + d1 += sh4_y * r4z + d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2) + d3 += sh4_x * r4z + d4 += sh4_x * r4x - sh4_y * r4y + + dst = x + dst[0] = d0 + dst[1] = -d1 + dst[2] = d2 * s_scale_dst2 + dst[3] = -d3 + dst[4] = d4 * s_scale_dst4 + + return dst + + +def load_calib(param, render_size=512): + # pixel unit / world unit + ortho_ratio = param['ortho_ratio'] + # world unit / model unit + scale = param['scale'] + # camera center world coordinate + center = param['center'] + # model rotation + R = param['R'] + + translate = -np.matmul(R, center).reshape(3, 1) + extrinsic = np.concatenate([R, translate], axis=1) + extrinsic = np.concatenate( + [extrinsic, np.array([0, 0, 0, 1]).reshape(1, 4)], 0) + # Match camera space to image pixel space + scale_intrinsic = np.identity(4) + scale_intrinsic[0, 0] = scale / ortho_ratio + scale_intrinsic[1, 1] = -scale / ortho_ratio + scale_intrinsic[2, 2] = scale / ortho_ratio + # Match image pixel space to image uv space + uv_intrinsic = np.identity(4) + uv_intrinsic[0, 0] = 1.0 / float(render_size // 2) + uv_intrinsic[1, 1] = 1.0 / float(render_size // 2) + uv_intrinsic[2, 2] = 1.0 / float(render_size // 2) + + intrinsic = np.matmul(uv_intrinsic, scale_intrinsic) + calib = np.concatenate([extrinsic, intrinsic], axis=0) + return calib + + +def render_prt_ortho(out_path, + folder_name, + subject_name, + shs, + rndr, + rndr_uv, + im_size, + angl_step=4, + n_light=1, + pitch=[0]): + cam = Camera(width=im_size, height=im_size) + cam.ortho_ratio = 0.4 * (512 / im_size) + cam.near = -100 + cam.far = 100 + cam.sanity_check() + + # set path for obj, prt + mesh_file = os.path.join(folder_name, subject_name + '_100k.obj') + if not os.path.exists(mesh_file): + print('ERROR: obj file does not exist!!', mesh_file) + return + prt_file = os.path.join(folder_name, 'bounce', 'bounce0.txt') + if not os.path.exists(prt_file): + print('ERROR: prt file does not exist!!!', prt_file) + return + face_prt_file = os.path.join(folder_name, 'bounce', 'face.npy') + if not os.path.exists(face_prt_file): + print('ERROR: face prt file does not exist!!!', prt_file) + return + text_file = os.path.join(folder_name, 'tex', subject_name + '_dif_2k.jpg') + if not os.path.exists(text_file): + print('ERROR: dif file does not exist!!', text_file) + return + + texture_image = cv2.imread(text_file) + texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) + + vertices, faces, normals, faces_normals, textures, face_textures = load_scan( + mesh_file, with_normal=True, with_texture=True) + vmin = vertices.min(0) + vmax = vertices.max(0) + up_axis = 1 if (vmax - vmin).argmax() == 1 else 2 + + vmed = np.median(vertices, 0) + vmed[up_axis] = 0.5 * (vmax[up_axis] + vmin[up_axis]) + y_scale = 180 / (vmax[up_axis] - vmin[up_axis]) + + rndr.set_norm_mat(y_scale, vmed) + rndr_uv.set_norm_mat(y_scale, vmed) + + tan, bitan = compute_tangent(vertices, faces, normals, textures, + face_textures) + prt = np.loadtxt(prt_file) + face_prt = np.load(face_prt_file) + rndr.set_mesh(vertices, faces, normals, faces_normals, textures, + face_textures, prt, face_prt, tan, bitan) + rndr.set_albedo(texture_image) + + rndr_uv.set_mesh(vertices, faces, normals, faces_normals, textures, + face_textures, prt, face_prt, tan, bitan) + rndr_uv.set_albedo(texture_image) + + os.makedirs(os.path.join(out_path, 'GEO', 'OBJ', subject_name), + exist_ok=True) + os.makedirs(os.path.join(out_path, 'PARAM', subject_name), exist_ok=True) + os.makedirs(os.path.join(out_path, 'RENDER', subject_name), exist_ok=True) + os.makedirs(os.path.join(out_path, 'MASK', subject_name), exist_ok=True) + os.makedirs(os.path.join(out_path, 'UV_RENDER', subject_name), + exist_ok=True) + os.makedirs(os.path.join(out_path, 'UV_MASK', subject_name), exist_ok=True) + os.makedirs(os.path.join(out_path, 'UV_POS', subject_name), exist_ok=True) + os.makedirs(os.path.join(out_path, 'UV_NORMAL', subject_name), + exist_ok=True) + + if not os.path.exists(os.path.join(out_path, 'val.txt')): + f = open(os.path.join(out_path, 'val.txt'), 'w') + f.close() + + # copy obj file + cmd = 'cp %s %s' % (mesh_file, + os.path.join(out_path, 'GEO', 'OBJ', subject_name)) + print(cmd) + os.system(cmd) + + for p in pitch: + for y in tqdm(range(0, 360, angl_step)): + R = np.matmul(make_rotate(math.radians(p), 0, 0), + make_rotate(0, math.radians(y), 0)) + if up_axis == 2: + R = np.matmul(R, make_rotate(math.radians(90), 0, 0)) + + rndr.rot_matrix = R + rndr_uv.rot_matrix = R + rndr.set_camera(cam) + rndr_uv.set_camera(cam) + + for j in range(n_light): + sh_id = random.randint(0, shs.shape[0] - 1) + sh = shs[sh_id] + sh_angle = 0.2 * np.pi * (random.random() - 0.5) + sh = rotateSH(sh, make_rotate(0, sh_angle, 0).T) + + dic = { + 'sh': sh, + 'ortho_ratio': cam.ortho_ratio, + 'scale': y_scale, + 'center': vmed, + 'R': R + } + + rndr.set_sh(sh) + rndr.analytic = False + rndr.use_inverse_depth = False + rndr.display() + + out_all_f = rndr.get_color(0) + out_mask = out_all_f[:, :, 3] + out_all_f = cv2.cvtColor(out_all_f, cv2.COLOR_RGBA2BGR) + + np.save( + os.path.join(out_path, 'PARAM', subject_name, + '%d_%d_%02d.npy' % (y, p, j)), dic) + cv2.imwrite( + os.path.join(out_path, 'RENDER', subject_name, + '%d_%d_%02d.jpg' % (y, p, j)), + 255.0 * out_all_f) + cv2.imwrite( + os.path.join(out_path, 'MASK', subject_name, + '%d_%d_%02d.png' % (y, p, j)), + 255.0 * out_mask) + + rndr_uv.set_sh(sh) + rndr_uv.analytic = False + rndr_uv.use_inverse_depth = False + rndr_uv.display() + + uv_color = rndr_uv.get_color(0) + uv_color = cv2.cvtColor(uv_color, cv2.COLOR_RGBA2BGR) + cv2.imwrite( + os.path.join(out_path, 'UV_RENDER', subject_name, + '%d_%d_%02d.jpg' % (y, p, j)), + 255.0 * uv_color) + + if y == 0 and j == 0 and p == pitch[0]: + uv_pos = rndr_uv.get_color(1) + uv_mask = uv_pos[:, :, 3] + cv2.imwrite( + os.path.join(out_path, 'UV_MASK', subject_name, + '00.png'), 255.0 * uv_mask) + + data = { + 'default': uv_pos[:, :, :3] + } # default is a reserved name + pyexr.write( + os.path.join(out_path, 'UV_POS', subject_name, + '00.exr'), data) + + uv_nml = rndr_uv.get_color(2) + uv_nml = cv2.cvtColor(uv_nml, cv2.COLOR_RGBA2BGR) + cv2.imwrite( + os.path.join(out_path, 'UV_NORMAL', subject_name, + '00.png'), 255.0 * uv_nml) diff --git a/lib/renderer/prt_util.py b/lib/renderer/prt_util.py new file mode 100644 index 0000000000000000000000000000000000000000..d021af079b13b2680c8e0214e36288bf81be2c76 --- /dev/null +++ b/lib/renderer/prt_util.py @@ -0,0 +1,199 @@ + +# -*- coding: utf-8 -*- + +# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is +# holder of all proprietary rights on this computer program. +# You can only use this computer program if you have closed +# a license agreement with MPG or you get the right to use the computer +# program from someone who is authorized to grant you that right. +# Any use of the computer program without a valid license is prohibited and +# liable to prosecution. +# +# Copyright©2019 Max-Planck-Gesellschaft zur Förderung +# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +# for Intelligent Systems. All rights reserved. +# +# Contact: ps-license@tuebingen.mpg.de + +import os +import trimesh +import numpy as np +import math +from scipy.special import sph_harm +import argparse +from tqdm import tqdm +from trimesh.util import bounds_tree + + +def factratio(N, D): + if N >= D: + prod = 1.0 + for i in range(D + 1, N + 1): + prod *= i + return prod + else: + prod = 1.0 + for i in range(N + 1, D + 1): + prod *= i + return 1.0 / prod + + +def KVal(M, L): + return math.sqrt(((2 * L + 1) / (4 * math.pi)) * (factratio(L - M, L + M))) + + +def AssociatedLegendre(M, L, x): + if M < 0 or M > L or np.max(np.abs(x)) > 1.0: + return np.zeros_like(x) + + pmm = np.ones_like(x) + if M > 0: + somx2 = np.sqrt((1.0 + x) * (1.0 - x)) + fact = 1.0 + for i in range(1, M + 1): + pmm = -pmm * fact * somx2 + fact = fact + 2 + + if L == M: + return pmm + else: + pmmp1 = x * (2 * M + 1) * pmm + if L == M + 1: + return pmmp1 + else: + pll = np.zeros_like(x) + for i in range(M + 2, L + 1): + pll = (x * (2 * i - 1) * pmmp1 - (i + M - 1) * pmm) / (i - M) + pmm = pmmp1 + pmmp1 = pll + return pll + + +def SphericalHarmonic(M, L, theta, phi): + if M > 0: + return math.sqrt(2.0) * KVal(M, L) * np.cos( + M * phi) * AssociatedLegendre(M, L, np.cos(theta)) + elif M < 0: + return math.sqrt(2.0) * KVal(-M, L) * np.sin( + -M * phi) * AssociatedLegendre(-M, L, np.cos(theta)) + else: + return KVal(0, L) * AssociatedLegendre(0, L, np.cos(theta)) + + +def save_obj(mesh_path, verts): + file = open(mesh_path, 'w') + for v in verts: + file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) + file.close() + + +def sampleSphericalDirections(n): + xv = np.random.rand(n, n) + yv = np.random.rand(n, n) + theta = np.arccos(1 - 2 * xv) + phi = 2.0 * math.pi * yv + + phi = phi.reshape(-1) + theta = theta.reshape(-1) + + vx = -np.sin(theta) * np.cos(phi) + vy = -np.sin(theta) * np.sin(phi) + vz = np.cos(theta) + return np.stack([vx, vy, vz], 1), phi, theta + + +def getSHCoeffs(order, phi, theta): + shs = [] + for n in range(0, order + 1): + for m in range(-n, n + 1): + s = SphericalHarmonic(m, n, theta, phi) + shs.append(s) + + return np.stack(shs, 1) + + +def computePRT(mesh_path, scale, n, order): + + prt_dir = os.path.join(os.path.dirname(mesh_path), "prt") + bounce_path = os.path.join(prt_dir, "bounce.npy") + face_path = os.path.join(prt_dir, "face.npy") + + os.makedirs(prt_dir, exist_ok=True) + + PRT = None + F = None + + if os.path.exists(bounce_path) and os.path.exists(face_path): + + PRT = np.load(bounce_path) + F = np.load(face_path) + + else: + + mesh = trimesh.load(mesh_path, + skip_materials=True, + process=False, + maintain_order=True) + mesh.vertices *= scale + + vectors_orig, phi, theta = sampleSphericalDirections(n) + SH_orig = getSHCoeffs(order, phi, theta) + + w = 4.0 * math.pi / (n * n) + + origins = mesh.vertices + normals = mesh.vertex_normals + n_v = origins.shape[0] + + origins = np.repeat(origins[:, None], n, axis=1).reshape(-1, 3) + normals = np.repeat(normals[:, None], n, axis=1).reshape(-1, 3) + PRT_all = None + for i in range(n): + SH = np.repeat(SH_orig[None, (i * n):((i + 1) * n)], n_v, + axis=0).reshape(-1, SH_orig.shape[1]) + vectors = np.repeat(vectors_orig[None, (i * n):((i + 1) * n)], + n_v, + axis=0).reshape(-1, 3) + + dots = (vectors * normals).sum(1) + front = (dots > 0.0) + + delta = 1e-3 * min(mesh.bounding_box.extents) + + hits = mesh.ray.intersects_any(origins + delta * normals, vectors) + nohits = np.logical_and(front, np.logical_not(hits)) + + PRT = (nohits.astype(np.float) * dots)[:, None] * SH + + if PRT_all is not None: + PRT_all += (PRT.reshape(-1, n, SH.shape[1]).sum(1)) + else: + PRT_all = (PRT.reshape(-1, n, SH.shape[1]).sum(1)) + + PRT = w * PRT_all + F = mesh.faces + + np.save(bounce_path, PRT) + np.save(face_path, F) + + # NOTE: trimesh sometimes break the original vertex order, but topology will not change. + # when loading PRT in other program, use the triangle list from trimesh. + + return PRT, F + + +def testPRT(obj_path, n=40): + + os.makedirs(os.path.join(os.path.dirname(obj_path), + f'../bounce/{os.path.basename(obj_path)[:-4]}'), + exist_ok=True) + + PRT, F = computePRT(obj_path, n, 2) + np.savetxt( + os.path.join(os.path.dirname(obj_path), + f'../bounce/{os.path.basename(obj_path)[:-4]}', + 'bounce.npy'), PRT) + np.save( + os.path.join(os.path.dirname(obj_path), + f'../bounce/{os.path.basename(obj_path)[:-4]}', + 'face.npy'), F) diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000000000000000000000000000000000000..496bac0380dffa74dc3416a0294380ced58f4d8f --- /dev/null +++ b/packages.txt @@ -0,0 +1,9 @@ +libgl1 +freeglut3-dev +unzip +ffmpeg +libsm6 +libxext6 +libgl1-mesa-dri +libegl1-mesa +libgbm1 \ No newline at end of file diff --git a/requirement.txt b/requirement.txt new file mode 100644 index 0000000000000000000000000000000000000000..ca34dfc68926c5a63d2e2de705de9b5ba4980ed4 --- /dev/null +++ b/requirement.txt @@ -0,0 +1,37 @@ +iopath +fvcore +pyembree +matplotlib==3.5.0 +Pillow==9.0.0 +PyOpenGL +PyOpenGL_accelerate +PyYAML>=6.0 +yacs>=0.1.8 +scikit-image==0.19.1 +termcolor +tqdm==4.62.3 +loguru==0.5.3 +trimesh==3.9.35 +flatten_dict==0.4.2 +ipykernel==5.3.4 +ipywidgets==7.6.5 +jpeg4py +shapely==1.7.1 +vedo==2020.4.2 +rtree==0.9.7 +pytorch_lightning==1.2.5 +PyMCubes +kornia +chumpy +cython==0.29.20 +rembg>=2.0.3 +opencv-python +opencv_contrib_python +scikit-learn +protobuf==3.20.0 +gdown==4.4.0 +git+https://github.com/Project-Splinter/human_det +git+https://github.com/YuliangXiu/smplx.git +git+https://github.com/facebookresearch/pytorch3d.git +git+https://github.com/YuliangXiu/neural_voxelization_layer.git +git+https://github.com/NVIDIAGameWorks/kaolin.git