diff --git a/.gitignore b/.gitignore index bd973b50e5620cc6ca0a87c6ba5fe9790f643f64..f660fd6c046a92a58b99b25b3a84e0634467daf3 100644 --- a/.gitignore +++ b/.gitignore @@ -16,4 +16,5 @@ build dist *egg-info *.so -run.sh \ No newline at end of file +run.sh +*.log \ No newline at end of file diff --git a/README.md b/README.md index 8066a33a5ad4f0c8d395639b723127751b02cf0b..83f9681e81c28a2678137d2ee916284a38c2e69b 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Lightning cupy Twitter + discord invitation link

Google Colab Docker @@ -35,7 +36,6 @@ Project Page - youtube views

diff --git a/apps/IFGeo.py b/apps/IFGeo.py index 8cb033d8d3fbd597ac80526d3d0c691451975685..45e892a655fd894d59ee925f32a62d378370bd49 100644 --- a/apps/IFGeo.py +++ b/apps/IFGeo.py @@ -14,11 +14,12 @@ # # Contact: ps-license@tuebingen.mpg.de -from lib.common.seg3d_lossless import Seg3dLossless -from lib.common.train_util import * -import torch import numpy as np import pytorch_lightning as pl +import torch + +from lib.common.seg3d_lossless import Seg3dLossless +from lib.common.train_util import * torch.backends.cudnn.benchmark = True diff --git a/apps/Normal.py b/apps/Normal.py index 235c0aef05914ef040f1495843d339c758ebd9f3..94176f9996753680c3ae8608a485093deae14822 100644 --- a/apps/Normal.py +++ b/apps/Normal.py @@ -1,9 +1,10 @@ -from lib.net import NormalNet -from lib.common.train_util import batch_mean -import torch import numpy as np -from skimage.transform import resize import pytorch_lightning as pl +import torch +from skimage.transform import resize + +from lib.common.train_util import batch_mean +from lib.net import NormalNet class Normal(pl.LightningModule): diff --git a/apps/avatarizer.py b/apps/avatarizer.py index 10c35fa3d31d844685c1e17d3c6e8714426891c0..cc48a600f177f7892b68bb9e3a0b8e70ad733d56 100644 --- a/apps/avatarizer.py +++ b/apps/avatarizer.py @@ -1,17 +1,25 @@ -import numpy as np -import trimesh -import torch import argparse +import os import os.path as osp -import lib.smplx as smplx + +import numpy as np +import torch +import trimesh from pytorch3d.ops import SubdivideMeshes from pytorch3d.structures import Meshes - -from lib.smplx.lbs import general_lbs -from lib.dataset.mesh_util import keep_largest, poisson from scipy.spatial import cKDTree -from lib.dataset.mesh_util import SMPLX + +import lib.smplx as smplx from lib.common.local_affine import register +from lib.dataset.mesh_util import ( + SMPLX, + export_obj, + keep_largest, + o3d_ransac, + poisson, + remesh_laplacian, +) +from lib.smplx.lbs import general_lbs # loading cfg file parser = argparse.ArgumentParser() @@ -22,12 +30,18 @@ args = parser.parse_args() smplx_container = SMPLX() device = torch.device(f"cuda:{args.gpu}") +# loading SMPL-X and econ objs inferred with ECON prefix = f"./results/econ/obj/{args.name}" smpl_path = f"{prefix}_smpl_00.npy" -econ_path = f"{prefix}_0_full.obj" - smplx_param = np.load(smpl_path, allow_pickle=True).item() + +# export econ obj with pre-computed normals +econ_path = f"{prefix}_0_full.obj" econ_obj = trimesh.load(econ_path) +assert (econ_obj.vertex_normals.shape[1] == 3) +econ_obj.export(f"{prefix}_econ_raw.ply") + +# align econ with SMPL-X econ_obj.vertices *= np.array([1.0, -1.0, -1.0]) econ_obj.vertices /= smplx_param["scale"].cpu().numpy() econ_obj.vertices -= smplx_param["transl"].cpu().numpy() @@ -49,6 +63,7 @@ smpl_model = smplx.create( smpl_out_lst = [] +# obtain the pose params of T-pose, DA-pose, and the original pose for pose_type in ["t-pose", "da-pose", "pose"]: smpl_out_lst.append( smpl_model( @@ -67,6 +82,12 @@ for pose_type in ["t-pose", "da-pose", "pose"]: ) ) +# -------------------------- align econ and SMPL-X in DA-pose space ------------------------- # +# 1. find the vertex-correspondence between SMPL-X and econ +# 2. ECON + SMPL-X: posed space --> T-pose space --> DA-pose space +# 3. ECON (w/o hands & over-streched faces) + SMPL-X (w/ hands & registered inpainting parts) +# ------------------------------------------------------------------------------------------- # + smpl_verts = smpl_out_lst[2].vertices.detach()[0] smpl_tree = cKDTree(smpl_verts.cpu().numpy()) dist, idx = smpl_tree.query(econ_obj.vertices, k=5) @@ -143,14 +164,25 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da. smpl_da_body.remove_unreferenced_vertices() smpl_hand = smpl_da.copy() - smpl_hand.update_faces(smplx_container.smplx_mano_vertex_mask.numpy()[smpl_hand.faces].all(axis=1)) + smpl_hand.update_faces( + smplx_container.smplx_mano_vertex_mask.numpy()[smpl_hand.faces].all(axis=1) + ) smpl_hand.remove_unreferenced_vertices() econ_da = sum([smpl_hand, smpl_da_body, econ_da_body]) - econ_da = poisson(econ_da, f"{prefix}_econ_da.obj", depth=10, decimation=False) + econ_da = poisson(econ_da, f"{prefix}_econ_da.obj", depth=10, face_count=50000) + econ_da = remesh_laplacian(econ_da, f"{prefix}_econ_da.obj") else: econ_da = trimesh.load(f"{prefix}_econ_da.obj") smpl_da = trimesh.load(f"{prefix}_smpl_da.obj", maintain_orders=True, process=False) +# ---------------------- SMPL-X compatible ECON ---------------------- # +# 1. Find the new vertex-correspondence between NEW ECON and SMPL-X +# 2. Build the new J_regressor, lbs_weights, posedirs +# 3. canonicalize the NEW ECON +# ------------------------------------------------------------------- # + +print("Start building the SMPL-X compatible ECON model...") + smpl_tree = cKDTree(smpl_da.vertices) dist, idx = smpl_tree.query(econ_da.vertices, k=5) knn_weights = np.exp(-dist**2) @@ -167,19 +199,137 @@ econ_posedirs = ( econ_J_regressor /= econ_J_regressor.sum(dim=1, keepdims=True).clip(min=1e-10) econ_lbs_weights /= econ_lbs_weights.sum(dim=1, keepdims=True) -# re-compute da-pose rot_mat for ECON rot_mat_da = smpl_out_lst[1].vertex_transformation.detach()[0][idx[:, 0]] econ_da_verts = torch.tensor(econ_da.vertices).float() -econ_cano_verts = torch.inverse(rot_mat_da) @ torch.cat( - [econ_da_verts, torch.ones_like(econ_da_verts)[..., :1]], dim=1 -).unsqueeze(-1) +econ_cano_verts = torch.inverse(rot_mat_da) @ torch.cat([ + econ_da_verts, torch.ones_like(econ_da_verts)[..., :1] +], + dim=1).unsqueeze(-1) econ_cano_verts = econ_cano_verts[:, :3, 0].double() # ---------------------------------------------------- -# use any SMPL-X pose to animate ECON reconstruction +# use original pose to animate ECON reconstruction # ---------------------------------------------------- new_pose = smpl_out_lst[2].full_pose +# new_pose[:, :3] = 0. + +posed_econ_verts, _ = general_lbs( + pose=new_pose, + v_template=econ_cano_verts.unsqueeze(0), + posedirs=econ_posedirs, + J_regressor=econ_J_regressor, + parents=smpl_model.parents, + lbs_weights=econ_lbs_weights +) + +aligned_econ_verts = posed_econ_verts[0].detach().cpu().numpy() +aligned_econ_verts += smplx_param["transl"].cpu().numpy() +aligned_econ_verts *= smplx_param["scale"].cpu().numpy() * np.array([1.0, -1.0, -1.0]) +econ_pose = trimesh.Trimesh(aligned_econ_verts, econ_da.faces) +assert (econ_pose.vertex_normals.shape[1] == 3) +econ_pose.export(f"{prefix}_econ_pose.ply") + +# ------------------------------------------------------------------------- +# Align posed ECON with original ECON, for pixel-aligned texture extraction +# ------------------------------------------------------------------------- + +print("Start ICP registration between posed & original ECON...") +import open3d as o3d + +source = o3d.io.read_point_cloud(f"{prefix}_econ_pose.ply") +target = o3d.io.read_point_cloud(f"{prefix}_econ_raw.ply") +trans_init = o3d_ransac(source, target) +icp_criteria = o3d.pipelines.registration.ICPConvergenceCriteria( + relative_fitness=0.000001, relative_rmse=0.000001, max_iteration=100 +) + +reg_p2l = o3d.pipelines.registration.registration_icp( + source, + target, + 0.1, + trans_init, + o3d.pipelines.registration.TransformationEstimationPointToPlane(), + criteria=icp_criteria +) +econ_pose.apply_transform(reg_p2l.transformation) + +cache_path = f"{prefix.replace('obj','cache')}" +os.makedirs(cache_path, exist_ok=True) + +# ----------------------------------------------------------------- +# create UV texture (.obj .mtl .png) from posed ECON reconstruction +# ----------------------------------------------------------------- + +print("Start Color mapping...") +from PIL import Image +from torchvision import transforms + +from lib.common.render import query_color +from lib.common.render_utils import Pytorch3dRasterizer + +if not osp.exists(f"{prefix}_econ_icp_rgb.ply"): + masked_image = f"./results/econ/png/{args.name}_cloth.png" + tensor_image = transforms.ToTensor()(Image.open(masked_image))[:, :, :512] + final_colors = query_color( + torch.tensor(econ_pose.vertices).float(), + torch.tensor(econ_pose.faces).long(), + ((tensor_image - 0.5) * 2.0).unsqueeze(0).to(device), + device=device, + paint_normal=False, + ) + final_colors[final_colors == tensor_image[:, 0, 0] * 255.0] = 0.0 + final_colors = final_colors.detach().cpu().numpy() + econ_pose.visual.vertex_colors = final_colors + econ_pose.export(f"{prefix}_econ_icp_rgb.ply") +else: + mesh = trimesh.load(f"{prefix}_econ_icp_rgb.ply") + final_colors = mesh.visual.vertex_colors[:, :3] + +print("Start UV texture generation...") + +# Generate UV coords +v_np = econ_pose.vertices +f_np = econ_pose.faces + +vt_cache = osp.join(cache_path, "vt.pt") +ft_cache = osp.join(cache_path, "ft.pt") + +if osp.exists(vt_cache) and osp.exists(ft_cache): + vt = torch.load(vt_cache).to(device) + ft = torch.load(ft_cache).to(device) +else: + import xatlas + atlas = xatlas.Atlas() + atlas.add_mesh(v_np, f_np) + chart_options = xatlas.ChartOptions() + chart_options.max_iterations = 4 + atlas.generate(chart_options=chart_options) + vmapping, ft_np, vt_np = atlas[0] + + vt = torch.from_numpy(vt_np.astype(np.float32)).float().to(device) + ft = torch.from_numpy(ft_np.astype(np.int64)).int().to(device) + torch.save(vt.cpu(), vt_cache) + torch.save(ft.cpu(), ft_cache) + +# UV texture rendering +uv_rasterizer = Pytorch3dRasterizer(image_size=512, device=device) +texture_npy = uv_rasterizer.get_texture( + torch.cat([(vt - 0.5) * 2.0, torch.ones_like(vt[:, :1])], dim=1), + ft, + torch.tensor(v_np).unsqueeze(0).float(), + torch.tensor(f_np).unsqueeze(0).long(), + torch.tensor(final_colors).unsqueeze(0).float() / 255.0, +) + +Image.fromarray((texture_npy * 255.0).astype(np.uint8)).save(f"{cache_path}/texture.png") + +# UV mask for TEXTure (https://readpaper.com/paper/4720151447010820097) +texture_npy[texture_npy.sum(axis=2) == 0.0] = 1.0 +Image.fromarray((texture_npy * 255.0).astype(np.uint8)).save(f"{cache_path}/mask.png") + +# generate da-pose vertices +new_pose = smpl_out_lst[1].full_pose new_pose[:, :3] = 0. posed_econ_verts, _ = general_lbs( @@ -191,5 +341,8 @@ posed_econ_verts, _ = general_lbs( lbs_weights=econ_lbs_weights ) -econ_pose = trimesh.Trimesh(posed_econ_verts[0].detach(), econ_da.faces) -econ_pose.export(f"{prefix}_econ_pose.obj") +# export mtl file +mtl_string = f"newmtl mat0 \nKa 1.000000 1.000000 1.000000 \nKd 1.000000 1.000000 1.000000 \nKs 0.000000 0.000000 0.000000 \nTr 1.000000 \nillum 1 \nNs 0.000000\nmap_Kd texture.png" +with open(f"{cache_path}/material.mtl", 'w') as file: + file.write(mtl_string) +export_obj(posed_econ_verts[0].detach().cpu().numpy(), f_np, vt, ft, f"{cache_path}/mesh.obj") diff --git a/apps/benchmark.py b/apps/benchmark.py index 3c95b0e209439af0c221742a3f32168e1fb5c280..504e2537e7c38794390345a0b0dca3b3667540db 100644 --- a/apps/benchmark.py +++ b/apps/benchmark.py @@ -14,28 +14,29 @@ # # Contact: ps-license@tuebingen.mpg.de -import warnings import logging +import warnings warnings.filterwarnings("ignore") logging.getLogger("lightning").setLevel(logging.ERROR) logging.getLogger("trimesh").setLevel(logging.ERROR) -import torch import argparse import os +import torch from termcolor import colored from tqdm.auto import tqdm -from apps.Normal import Normal + from apps.IFGeo import IFGeo -from lib.common.config import cfg +from apps.Normal import Normal from lib.common.BNI import BNI from lib.common.BNI_utils import save_normal_tensor +from lib.common.config import cfg +from lib.common.voxelize import VoxelGrid from lib.dataset.EvalDataset import EvalDataset from lib.dataset.Evaluator import Evaluator from lib.dataset.mesh_util import * -from lib.common.voxelize import VoxelGrid torch.backends.cudnn.benchmark = True speed_analysis = False @@ -62,8 +63,14 @@ if __name__ == "__main__": device = torch.device("cuda:0") cfg_test_list = [ - "dataset.rotation_num", 3, "bni.use_smpl", ["hand"], "bni.use_ifnet", args.ifnet, - "bni.cut_intersection", True, + "dataset.rotation_num", + 3, + "bni.use_smpl", + ["hand"], + "bni.use_ifnet", + args.ifnet, + "bni.cut_intersection", + True, ] # # if w/ RenderPeople+CAPE @@ -176,12 +183,10 @@ if __name__ == "__main__": # mesh completion via IF-net in_tensor.update( - dataset.depth_to_voxel( - { - "depth_F": BNI_object.F_depth.unsqueeze(0).to(device), - "depth_B": BNI_object.B_depth.unsqueeze(0).to(device) - } - ) + dataset.depth_to_voxel({ + "depth_F": BNI_object.F_depth.unsqueeze(0).to(device), "depth_B": + BNI_object.B_depth.unsqueeze(0).to(device) + }) ) occupancies = VoxelGrid.from_mesh(side_mesh, cfg.vol_res, loc=[ diff --git a/apps/infer.py b/apps/infer.py index 97b80c437270711cccf0df2ff14779ecd1fc0ef7..3a9d85c6d92f6e92cbb4acd5f2f905163f079c4f 100644 --- a/apps/infer.py +++ b/apps/infer.py @@ -14,35 +14,37 @@ # # Contact: ps-license@tuebingen.mpg.de -import warnings import logging +import warnings warnings.filterwarnings("ignore") logging.getLogger("lightning").setLevel(logging.ERROR) logging.getLogger("trimesh").setLevel(logging.ERROR) -import torch, torchvision -import trimesh -import numpy as np import argparse import os +import numpy as np +import torch +import torchvision +import trimesh +from pytorch3d.ops import SubdivideMeshes from termcolor import colored from tqdm.auto import tqdm -from apps.Normal import Normal + from apps.IFGeo import IFGeo -from pytorch3d.ops import SubdivideMeshes -from lib.common.config import cfg -from lib.common.render import query_color -from lib.common.train_util import init_loss, Format -from lib.common.imutils import blend_rgb_norm +from apps.Normal import Normal from lib.common.BNI import BNI from lib.common.BNI_utils import save_normal_tensor -from lib.dataset.TestDataset import TestDataset +from lib.common.config import cfg +from lib.common.imutils import blend_rgb_norm from lib.common.local_affine import register -from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis -from lib.dataset.mesh_util import * +from lib.common.render import query_color +from lib.common.train_util import Format, init_loss from lib.common.voxelize import VoxelGrid +from lib.dataset.mesh_util import * +from lib.dataset.TestDataset import TestDataset +from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis torch.backends.cudnn.benchmark = True @@ -146,9 +148,8 @@ if __name__ == "__main__": os.makedirs(osp.join(args.out_dir, cfg.name, "obj"), exist_ok=True) in_tensor = { - "smpl_faces": data["smpl_faces"], - "image": data["img_icon"].to(device), - "mask": data["img_mask"].to(device) + "smpl_faces": data["smpl_faces"], "image": data["img_icon"].to(device), "mask": + data["img_mask"].to(device) } # The optimizer and variables @@ -157,9 +158,11 @@ if __name__ == "__main__": optimed_betas = data["betas"].requires_grad_(True) optimed_orient = data["global_orient"].requires_grad_(True) - optimizer_smpl = torch.optim.Adam( - [optimed_pose, optimed_trans, optimed_betas, optimed_orient], lr=1e-2, amsgrad=True - ) + optimizer_smpl = torch.optim.Adam([ + optimed_pose, optimed_trans, optimed_betas, optimed_orient + ], + lr=1e-2, + amsgrad=True) scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer_smpl, mode="min", @@ -234,9 +237,9 @@ if __name__ == "__main__": ) smpl_verts = (smpl_verts + optimed_trans) * data["scale"] - smpl_joints = (smpl_joints + optimed_trans) * data["scale"] * torch.tensor( - [1.0, 1.0, -1.0] - ).to(device) + smpl_joints = (smpl_joints + optimed_trans) * data["scale"] * torch.tensor([ + 1.0, 1.0, -1.0 + ]).to(device) # landmark errors smpl_joints_3d = ( @@ -280,13 +283,11 @@ if __name__ == "__main__": # BUG: PyTorch3D silhouette renderer generates dilated mask bg_value = in_tensor["T_normal_F"][0, 0, 0, 0] - smpl_arr_fake = torch.cat( - [ - in_tensor["T_normal_F"][:, 0].ne(bg_value).float(), - in_tensor["T_normal_B"][:, 0].ne(bg_value).float() - ], - dim=-1 - ) + smpl_arr_fake = torch.cat([ + in_tensor["T_normal_F"][:, 0].ne(bg_value).float(), + in_tensor["T_normal_B"][:, 0].ne(bg_value).float() + ], + dim=-1) body_overlap = (gt_arr * smpl_arr_fake.gt(0.0) ).sum(dim=[1, 2]) / smpl_arr_fake.gt(0.0).sum(dim=[1, 2]) @@ -322,22 +323,18 @@ if __name__ == "__main__": # save intermediate results if (i == args.loop_smpl - 1) and (not args.novis): - per_loop_lst.extend( - [ - in_tensor["image"], - in_tensor["T_normal_F"], - in_tensor["normal_F"], - diff_S[:, :, :512].unsqueeze(1).repeat(1, 3, 1, 1), - ] - ) - per_loop_lst.extend( - [ - in_tensor["image"], - in_tensor["T_normal_B"], - in_tensor["normal_B"], - diff_S[:, :, 512:].unsqueeze(1).repeat(1, 3, 1, 1), - ] - ) + per_loop_lst.extend([ + in_tensor["image"], + in_tensor["T_normal_F"], + in_tensor["normal_F"], + diff_S[:, :, :512].unsqueeze(1).repeat(1, 3, 1, 1), + ]) + per_loop_lst.extend([ + in_tensor["image"], + in_tensor["T_normal_B"], + in_tensor["normal_B"], + diff_S[:, :, 512:].unsqueeze(1).repeat(1, 3, 1, 1), + ]) per_data_lst.append( get_optim_grid_image(per_loop_lst, None, nrow=N_body * 2, type="smpl") ) @@ -357,13 +354,11 @@ if __name__ == "__main__": if not args.novis: img_crop_path = osp.join(args.out_dir, cfg.name, "png", f"{data['name']}_crop.png") torchvision.utils.save_image( - torch.cat( - [ - data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5, - (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5 - ], - dim=3 - ), img_crop_path + torch.cat([ + data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5, + (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5 + ], + dim=3), img_crop_path ) rgb_norm_F = blend_rgb_norm(in_tensor["normal_F"], data) @@ -392,27 +387,25 @@ if __name__ == "__main__": smpl_obj.export(smpl_obj_path) smpl_info = { "betas": - optimed_betas[idx].detach().cpu().unsqueeze(0), + optimed_betas[idx].detach().cpu().unsqueeze(0), "body_pose": - rotation_matrix_to_angle_axis(optimed_pose_mat[idx].detach() - ).cpu().unsqueeze(0), + rotation_matrix_to_angle_axis(optimed_pose_mat[idx].detach() + ).cpu().unsqueeze(0), "global_orient": - rotation_matrix_to_angle_axis(optimed_orient_mat[idx].detach() - ).cpu().unsqueeze(0), + rotation_matrix_to_angle_axis(optimed_orient_mat[idx].detach() + ).cpu().unsqueeze(0), "transl": - optimed_trans[idx].detach().cpu(), + optimed_trans[idx].detach().cpu(), "expression": - data["exp"][idx].cpu().unsqueeze(0), + data["exp"][idx].cpu().unsqueeze(0), "jaw_pose": - rotation_matrix_to_angle_axis(data["jaw_pose"][idx]).cpu().unsqueeze(0), + rotation_matrix_to_angle_axis(data["jaw_pose"][idx]).cpu().unsqueeze(0), "left_hand_pose": - rotation_matrix_to_angle_axis(data["left_hand_pose"][idx] - ).cpu().unsqueeze(0), + rotation_matrix_to_angle_axis(data["left_hand_pose"][idx]).cpu().unsqueeze(0), "right_hand_pose": - rotation_matrix_to_angle_axis(data["right_hand_pose"][idx] - ).cpu().unsqueeze(0), + rotation_matrix_to_angle_axis(data["right_hand_pose"][idx]).cpu().unsqueeze(0), "scale": - data["scale"][idx].cpu(), + data["scale"][idx].cpu(), } np.save( smpl_obj_path.replace(".obj", ".npy"), @@ -434,8 +427,8 @@ if __name__ == "__main__": per_data_lst = [] - batch_smpl_verts = in_tensor["smpl_verts"].detach( - ) * torch.tensor([1.0, -1.0, 1.0], device=device) + batch_smpl_verts = in_tensor["smpl_verts"].detach() * torch.tensor([1.0, -1.0, 1.0], + device=device) batch_smpl_faces = in_tensor["smpl_faces"].detach()[:, :, [0, 2, 1]] in_tensor["depth_F"], in_tensor["depth_B"] = dataset.render_depth( @@ -491,12 +484,10 @@ if __name__ == "__main__": # mesh completion via IF-net in_tensor.update( - dataset.depth_to_voxel( - { - "depth_F": BNI_object.F_depth.unsqueeze(0), - "depth_B": BNI_object.B_depth.unsqueeze(0) - } - ) + dataset.depth_to_voxel({ + "depth_F": BNI_object.F_depth.unsqueeze(0), "depth_B": + BNI_object.B_depth.unsqueeze(0) + }) ) occupancies = VoxelGrid.from_mesh(side_mesh, cfg.vol_res, loc=[ diff --git a/apps/multi_render.py b/apps/multi_render.py index 4088440757ce81137aaad7685d9df4b53b1c1383..10bdf4e52b3b69e756906d402124d2359c9403f0 100644 --- a/apps/multi_render.py +++ b/apps/multi_render.py @@ -1,7 +1,9 @@ -from lib.common.render import Render -import torch import argparse +import torch + +from lib.common.render import Render + root = "./results/econ/vid" # loading cfg file diff --git a/configs/econ.yaml b/configs/econ.yaml index f96784da9a111dbc8e1139d3f6e9af641401e9e1..5283393be72b3dc50ed7f88021c4c76515d19a1d 100644 --- a/configs/econ.yaml +++ b/configs/econ.yaml @@ -28,7 +28,7 @@ bni: lambda1: 1e-4 boundary_consist: 1e-6 poisson_depth: 10 - use_smpl: ["hand", "face"] + use_smpl: ["hand"] use_ifnet: False use_poisson: True hand_thres: 8e-2 diff --git a/docs/tricks.md b/docs/tricks.md index 8fa32714b3f77561705a9c83052a6367ed7cae58..17ba7e98f502d385fc075ae756656a639b0502a5 100644 --- a/docs/tricks.md +++ b/docs/tricks.md @@ -2,7 +2,7 @@ ### If the reconstructed geometry is not satisfying, play with the adjustable parameters in _config/econ.yaml_ -- `use_smpl: ["hand", "face"]` +- `use_smpl: ["hand"]` - [ ]: don't use either hands or face parts from SMPL-X - ["hand"]: only use the **visible** hands from SMPL-X - ["hand", "face"]: use both **visible** hands and face from SMPL-X diff --git a/lib/common/BNI.py b/lib/common/BNI.py index 2e4365e012c3c862b5f009bb879c4dda45490db4..1ae381db6786bbf3c07c1b423467e39457fe651a 100644 --- a/lib/common/BNI.py +++ b/lib/common/BNI.py @@ -1,10 +1,12 @@ -from lib.common.BNI_utils import ( - verts_inverse_transform, depth_inverse_transform, double_side_bilateral_normal_integration -) - import torch import trimesh +from lib.common.BNI_utils import ( + depth_inverse_transform, + double_side_bilateral_normal_integration, + verts_inverse_transform, +) + class BNI: def __init__(self, dir_path, name, BNI_dict, cfg, device): @@ -84,8 +86,9 @@ class BNI: if __name__ == "__main__": - import numpy as np import os.path as osp + + import numpy as np from tqdm import tqdm root = "/home/yxiu/Code/ECON/results/examples/BNI" diff --git a/lib/common/BNI_utils.py b/lib/common/BNI_utils.py index 6fa64c44c9049d0a7a37fb3e6d252549e180a931..235910fb51c263c51410afc247795f224456dd78 100644 --- a/lib/common/BNI_utils.py +++ b/lib/common/BNI_utils.py @@ -1,13 +1,23 @@ -import torch -import trimesh -import cv2, os -from PIL import Image +import os import os.path as osp + import cupy as cp +import cv2 import numpy as np -from cupyx.scipy.sparse import csr_matrix, vstack, hstack, spdiags, diags, coo_matrix +import torch +import trimesh +from cupyx.scipy.sparse import ( + coo_matrix, + csr_matrix, + diags, + hstack, + spdiags, + vstack, +) from cupyx.scipy.sparse.linalg import cg +from PIL import Image from tqdm.auto import tqdm + from lib.dataset.mesh_util import clean_floats @@ -68,13 +78,11 @@ def mean_value_cordinates(inner_pts, contour_pts): body_edges_c = np.roll(body_edges_a, shift=-1, axis=1) body_edges_b = np.sqrt(((contour_pts - np.roll(contour_pts, shift=-1, axis=0))**2).sum(axis=1)) - body_edges = np.concatenate( - [ - body_edges_a[..., None], body_edges_c[..., None], - np.repeat(body_edges_b[None, :, None], axis=0, repeats=len(inner_pts)) - ], - axis=-1 - ) + body_edges = np.concatenate([ + body_edges_a[..., None], body_edges_c[..., None], + np.repeat(body_edges_b[None, :, None], axis=0, repeats=len(inner_pts)) + ], + axis=-1) body_cos = (body_edges[:, :, 0]**2 + body_edges[:, :, 1]**2 - body_edges[:, :, 2]**2) / (2 * body_edges[:, :, 0] * body_edges[:, :, 1]) @@ -167,9 +175,9 @@ def verts_transform(t, depth_scale): t_copy = t.clone() t_copy *= depth_scale * 0.5 t_copy += depth_scale * 0.5 - t_copy = t_copy[:, [1, 0, 2]] * torch.Tensor([2.0, 2.0, -2.0]) + torch.Tensor( - [0.0, 0.0, depth_scale] - ) + t_copy = t_copy[:, [1, 0, 2]] * torch.Tensor([2.0, 2.0, -2.0]) + torch.Tensor([ + 0.0, 0.0, depth_scale + ]) return t_copy @@ -342,15 +350,13 @@ def construct_facets_from(mask): facet_bottom_left_mask = move_bottom(facet_top_left_mask) facet_bottom_right_mask = move_bottom_right(facet_top_left_mask) - return cp.hstack( - ( - 4 * cp.ones((cp.sum(facet_top_left_mask).item(), 1)), - idx[facet_top_left_mask][:, None], - idx[facet_bottom_left_mask][:, None], - idx[facet_bottom_right_mask][:, None], - idx[facet_top_right_mask][:, None], - ) - ).astype(int) + return cp.hstack(( + 4 * cp.ones((cp.sum(facet_top_left_mask).item(), 1)), + idx[facet_top_left_mask][:, None], + idx[facet_bottom_left_mask][:, None], + idx[facet_bottom_right_mask][:, None], + idx[facet_top_right_mask][:, None], + )).astype(int) def map_depth_map_to_point_clouds(depth_map, mask, K=None, step_size=1): @@ -614,7 +620,7 @@ def double_side_bilateral_normal_integration( energy_list.append(energy) relative_energy = cp.abs(energy - energy_old) / energy_old - + # print(f"step {i + 1}/{max_iter} energy: {energy:.3e}" # f" relative energy: {relative_energy:.3e}") @@ -640,13 +646,11 @@ def double_side_bilateral_normal_integration( B_verts = verts_inverse_transform(torch.as_tensor(vertices_back).float(), 256.0) F_B_verts = torch.cat((F_verts, B_verts), dim=0) - F_B_faces = torch.cat( - ( - torch.as_tensor(faces_front_).long(), - torch.as_tensor(faces_back_).long() + faces_front_.max() + 1 - ), - dim=0 - ) + F_B_faces = torch.cat(( + torch.as_tensor(faces_front_).long(), + torch.as_tensor(faces_back_).long() + faces_front_.max() + 1 + ), + dim=0) front_surf = trimesh.Trimesh(F_verts, faces_front_) back_surf = trimesh.Trimesh(B_verts, faces_back_) @@ -690,12 +694,12 @@ def double_side_bilateral_normal_integration( back_mesh = clean_floats(trimesh.Trimesh(vertices_back, faces_back)) result = { - "F_verts": torch.as_tensor(front_mesh.vertices).float(), - "F_faces": torch.as_tensor(front_mesh.faces).long(), - "B_verts": torch.as_tensor(back_mesh.vertices).float(), - "B_faces": torch.as_tensor(back_mesh.faces).long(), - "F_depth": torch.as_tensor(depth_map_front_est).float(), - "B_depth": torch.as_tensor(depth_map_back_est).float() + "F_verts": torch.as_tensor(front_mesh.vertices).float(), "F_faces": torch.as_tensor( + front_mesh.faces + ).long(), "B_verts": torch.as_tensor(back_mesh.vertices).float(), "B_faces": + torch.as_tensor(back_mesh.faces).long(), "F_depth": + torch.as_tensor(depth_map_front_est).float(), "B_depth": + torch.as_tensor(depth_map_back_est).float() } return result diff --git a/lib/common/blender_utils.py b/lib/common/blender_utils.py deleted file mode 100644 index a02260cc722bd9729dfbeb153543ac5f648deacf..0000000000000000000000000000000000000000 --- a/lib/common/blender_utils.py +++ /dev/null @@ -1,383 +0,0 @@ -import bpy -import sys, os -from math import radians -import mathutils -import bmesh - -print(sys.exec_prefix) -from tqdm import tqdm -import numpy as np - -################################################## -# Globals -################################################## - -views = 120 - -render = 'eevee' -cycles_gpu = False - -quality_preview = False -samples_preview = 16 -samples_final = 256 - -resolution_x = 512 -resolution_y = 512 - -shadows = False - -# diffuse_color = (57.0/255.0, 108.0/255.0, 189.0/255.0, 1.0) -# diffuse_color = (18/255., 139/255., 142/255.,1) #correct -# diffuse_color = (251/255., 60/255., 60/255.,1) #wrong - -smooth = False - -wireframe = False -line_thickness = 0.1 -quads = False - -object_transparent = False -mouth_transparent = False - -compositor_background_image = False -compositor_image_scale = 1.0 -compositor_alpha = 0.7 - -################################################## -# Helper functions -################################################## - - -def blender_print(*args, **kwargs): - print(*args, **kwargs, file=sys.stderr) - - -def using_app(): - ''' Returns if script is running through Blender application (GUI or background processing)''' - return (not sys.argv[0].endswith('.py')) - - -def setup_diffuse_transparent_material(target, color, object_transparent, backface_transparent): - ''' Sets up diffuse/transparent material with backface culling in cycles''' - - mat = target.active_material - if mat is None: - # Create material - mat = bpy.data.materials.new(name='Material') - target.data.materials.append(mat) - - mat.use_nodes = True - nodes = mat.node_tree.nodes - for node in nodes: - nodes.remove(node) - - node_geometry = nodes.new('ShaderNodeNewGeometry') - - node_diffuse = nodes.new('ShaderNodeBsdfDiffuse') - node_diffuse.inputs[0].default_value = color - - node_transparent = nodes.new('ShaderNodeBsdfTransparent') - node_transparent.inputs[0].default_value = (1.0, 1.0, 1.0, 1.0) - - node_emission = nodes.new('ShaderNodeEmission') - node_emission.inputs[0].default_value = (0.0, 0.0, 0.0, 1.0) - - node_mix = nodes.new(type='ShaderNodeMixShader') - if object_transparent: - node_mix.inputs[0].default_value = 1.0 - else: - node_mix.inputs[0].default_value = 0.0 - - node_mix_mouth = nodes.new(type='ShaderNodeMixShader') - if object_transparent or backface_transparent: - node_mix_mouth.inputs[0].default_value = 1.0 - else: - node_mix_mouth.inputs[0].default_value = 0.0 - - node_mix_backface = nodes.new(type='ShaderNodeMixShader') - - node_output = nodes.new(type='ShaderNodeOutputMaterial') - - links = mat.node_tree.links - - links.new(node_geometry.outputs[6], node_mix_backface.inputs[0]) - - links.new(node_diffuse.outputs[0], node_mix.inputs[1]) - links.new(node_transparent.outputs[0], node_mix.inputs[2]) - links.new(node_mix.outputs[0], node_mix_backface.inputs[1]) - - links.new(node_emission.outputs[0], node_mix_mouth.inputs[1]) - links.new(node_transparent.outputs[0], node_mix_mouth.inputs[2]) - links.new(node_mix_mouth.outputs[0], node_mix_backface.inputs[2]) - - links.new(node_mix_backface.outputs[0], node_output.inputs[0]) - return - - -################################################## - - -def setup_scene(): - global render - global cycles_gpu - global quality_preview - global resolution_x - global resolution_y - global shadows - global wireframe - global line_thickness - global compositor_background_image - - # Remove default cube - if 'Cube' in bpy.data.objects: - bpy.data.objects['Cube'].select_set(True) - bpy.ops.object.delete() - - scene = bpy.data.scenes['Scene'] - - # Setup render engine - if render == 'cycles': - scene.render.engine = 'CYCLES' - else: - scene.render.engine = 'BLENDER_EEVEE' - - scene.render.resolution_x = resolution_x - scene.render.resolution_y = resolution_y - scene.render.resolution_percentage = 100 - scene.render.film_transparent = True - if quality_preview: - scene.cycles.samples = samples_preview - else: - scene.cycles.samples = samples_final - - # Setup Cycles CUDA GPU acceleration if requested - if render == 'cycles': - if cycles_gpu: - print('Activating GPU acceleration') - bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA' - - if bpy.app.version[0] >= 3: - cuda_devices = bpy.context.preferences.addons[ - 'cycles'].preferences.get_devices_for_type(compute_device_type='CUDA') - else: - (cuda_devices, opencl_devices - ) = bpy.context.preferences.addons['cycles'].preferences.get_devices() - - if (len(cuda_devices) < 1): - print('ERROR: CUDA GPU acceleration not available') - sys.exit(1) - - for cuda_device in cuda_devices: - if cuda_device.type == 'CUDA': - cuda_device.use = True - print('Using CUDA device: ' + str(cuda_device.name)) - else: - cuda_device.use = False - print('Igoring CUDA device: ' + str(cuda_device.name)) - - scene.cycles.device = 'GPU' - if bpy.app.version[0] < 3: - scene.render.tile_x = 256 - scene.render.tile_y = 256 - else: - scene.cycles.device = 'CPU' - if bpy.app.version[0] < 3: - scene.render.tile_x = 64 - scene.render.tile_y = 64 - - # Disable Blender 3 denoiser to properly measure Cycles render speed - if bpy.app.version[0] >= 3: - scene.cycles.use_denoising = False - - # Setup camera - camera = bpy.data.objects['Camera'] - camera.location = (0.0, -3, 1.8) - camera.rotation_euler = (radians(74), 0.0, 0) - bpy.data.cameras['Camera'].lens = 55 - - # Setup light - - # Setup lights - light = bpy.data.objects['Light'] - light.location = (-2, -3.0, 0.0) - light.rotation_euler = (radians(90.0), 0.0, 0.0) - bpy.data.lights['Light'].type = 'POINT' - bpy.data.lights['Light'].energy = 2 - light.data.cycles.cast_shadow = False - - if 'Sun' not in bpy.data.objects: - bpy.ops.object.light_add(type='SUN') - light_sun = bpy.context.active_object - light_sun.location = (0.0, -3, 0.0) - light_sun.rotation_euler = (radians(45.0), 0.0, radians(30)) - bpy.data.lights['Sun'].energy = 2 - light_sun.data.cycles.cast_shadow = shadows - else: - light_sun = bpy.data.objects['Sun'] - - if shadows: - # Setup shadow catcher - bpy.ops.mesh.primitive_plane_add() - plane = bpy.context.active_object - plane.scale = (5.0, 5.0, 1) - - plane.cycles.is_shadow_catcher = True - - # Exclude plane from diffuse cycles contribution to avoid bright pixel noise in body rendering - # plane.cycles_visibility.diffuse = False - - if wireframe: - # Unmark freestyle edges - bpy.ops.object.mode_set(mode='EDIT') - bpy.ops.mesh.mark_freestyle_edge(clear=True) - bpy.ops.object.mode_set(mode='OBJECT') - - # Setup freestyle mode for wireframe overlay rendering - if wireframe: - scene.render.use_freestyle = True - scene.render.line_thickness = line_thickness - bpy.context.view_layer.freestyle_settings.linesets[0].select_edge_mark = True - - # Disable border edges so that we don't see contour of shadow catcher plane - bpy.context.view_layer.freestyle_settings.linesets[0].select_border = False - else: - scene.render.use_freestyle = False - - if compositor_background_image: - # Setup compositing when using background image - setup_compositing() - else: - # Output transparent image when no background is used - scene.render.image_settings.color_mode = 'RGBA' - - -################################################## - - -def setup_compositing(): - - global compositor_image_scale - global compositor_alpha - - # Node editor compositing setup - bpy.context.scene.use_nodes = True - tree = bpy.context.scene.node_tree - - # Create input image node - image_node = tree.nodes.new(type='CompositorNodeImage') - - scale_node = tree.nodes.new(type='CompositorNodeScale') - scale_node.inputs[1].default_value = compositor_image_scale - scale_node.inputs[2].default_value = compositor_image_scale - - blend_node = tree.nodes.new(type='CompositorNodeAlphaOver') - blend_node.inputs[0].default_value = compositor_alpha - - # Link nodes - links = tree.links - links.new(image_node.outputs[0], scale_node.inputs[0]) - - links.new(scale_node.outputs[0], blend_node.inputs[1]) - links.new(tree.nodes['Render Layers'].outputs[0], blend_node.inputs[2]) - - links.new(blend_node.outputs[0], tree.nodes['Composite'].inputs[0]) - - -def render_file(input_file, input_dir, output_file, output_dir, yaw, correct): - '''Render image of given model file''' - global smooth - global object_transparent - global mouth_transparent - global compositor_background_image - global quads - - path = input_dir + input_file - - # Import object into scene - bpy.ops.import_scene.obj(filepath=path) - object = bpy.context.selected_objects[0] - - object.rotation_euler = (radians(90.0), 0.0, radians(yaw)) - z_bottom = np.min(np.array([vert.co for vert in object.data.vertices])[:, 1]) - # z_top = np.max(np.array([vert.co for vert in object.data.vertices])[:,1]) - # blender_print(radians(90.0), z_bottom, z_top) - object.location -= mathutils.Vector((0.0, 0.0, z_bottom)) - - if quads: - bpy.context.view_layer.objects.active = object - bpy.ops.object.mode_set(mode='EDIT') - bpy.ops.mesh.tris_convert_to_quads() - bpy.ops.object.mode_set(mode='OBJECT') - - if smooth: - bpy.ops.object.shade_smooth() - - # Mark freestyle edges - bpy.context.view_layer.objects.active = object - bpy.ops.object.mode_set(mode='EDIT') - bpy.ops.mesh.mark_freestyle_edge(clear=False) - bpy.ops.object.mode_set(mode='OBJECT') - - if correct: - diffuse_color = (18 / 255., 139 / 255., 142 / 255., 1) #correct - else: - diffuse_color = (251 / 255., 60 / 255., 60 / 255., 1) #wrong - - setup_diffuse_transparent_material(object, diffuse_color, object_transparent, mouth_transparent) - - if compositor_background_image: - # Set background image - image_path = input_dir + input_file.replace('.obj', '_original.png') - bpy.context.scene.node_tree.nodes['Image'].image = bpy.data.images.load(image_path) - - # Render - bpy.context.scene.render.filepath = os.path.join(output_dir, output_file) - - # Silence console output of bpy.ops.render.render by redirecting stdout to file - # Note: Does not actually write the output to file (Windows 7) - sys.stdout.flush() - old = os.dup(1) - os.close(1) - os.open('blender_render.log', os.O_WRONLY | os.O_CREAT) - - # Render - bpy.ops.render.render(write_still=True) - - # Remove temporary output redirection - # sys.stdout.flush() - # os.close(1) - # os.dup(old) - # os.close(old) - - # Delete last selected object from scene - object.select_set(True) - bpy.ops.object.delete() - - -def process_file(input_file, input_dir, output_file, output_dir, correct=True): - global views - global quality_preview - - if not input_file.endswith('.obj'): - print('ERROR: Invalid input: ' + input_file) - return - - print('Processing: ' + input_file) - if output_file == '': - output_file = input_file[:-4] - - if quality_preview: - output_file = output_file.replace('.png', '-preview.png') - - angle = 360.0 / views - pbar = tqdm(range(0, views)) - for view in pbar: - pbar.set_description(f"{os.path.basename(output_file)} | View:{str(view)}") - yaw = view * angle - output_file_view = f"{output_file}/{view:03d}.png" - if not os.path.exists(os.path.join(output_dir, output_file_view)): - render_file(input_file, input_dir, output_file_view, output_dir, yaw, correct) - - cmd = "ffmpeg -loglevel quiet -r 30 -f lavfi -i color=c=white:s=512x512 -i " + os.path.join(output_dir, output_file, '%3d.png') + \ - " -shortest -filter_complex \"[0:v][1:v]overlay=shortest=1,format=yuv420p[out]\" -map \"[out]\" -y " + output_dir+"/"+output_file+".mp4" - os.system(cmd) diff --git a/lib/common/cloth_extraction.py b/lib/common/cloth_extraction.py index 612a96787e1aa836b097971e7aaf55b284ef178a..3e6f67234e25eaa2124c3d20692d138b49d20f8d 100644 --- a/lib/common/cloth_extraction.py +++ b/lib/common/cloth_extraction.py @@ -1,10 +1,11 @@ -import numpy as np +import itertools import json import os -import itertools +from collections import Counter + +import numpy as np import trimesh from matplotlib.path import Path -from collections import Counter from sklearn.neighbors import KNeighborsClassifier @@ -36,13 +37,11 @@ def load_segmentation(path, shape): xy = np.vstack((x, y)).T coordinates.append(xy) - segmentations.append( - { - "type": val["category_name"], - "type_id": val["category_id"], - "coordinates": coordinates, - } - ) + segmentations.append({ + "type": val["category_name"], + "type_id": val["category_id"], + "coordinates": coordinates, + }) return segmentations diff --git a/lib/common/config.py b/lib/common/config.py index 3b42fc73cea714259e2a350ee3bd7dc14dbf9375..6d946c115f2b482d598b06d405608f3004b75e78 100644 --- a/lib/common/config.py +++ b/lib/common/config.py @@ -14,9 +14,10 @@ # # Contact: ps-license@tuebingen.mpg.de -from yacs.config import CfgNode as CN import os +from yacs.config import CfgNode as CN + _C = CN(new_allowed=True) # needed by trainer diff --git a/lib/common/imutils.py b/lib/common/imutils.py index 39f61cff3d38d9b50ed245b1a87655ceb2c0eae3..287ab057ee20ac6d28ae8f7a83a77f681af040db 100644 --- a/lib/common/imutils.py +++ b/lib/common/imutils.py @@ -1,17 +1,18 @@ import os -os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" + +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" import cv2 import mediapipe as mp -import torch import numpy as np +import torch import torch.nn.functional as F +from kornia.geometry.transform import get_affine_matrix2d, warp_affine from PIL import Image -from lib.pymafx.core import constants - from rembg import remove from rembg.session_factory import new_session from torchvision import transforms -from kornia.geometry.transform import get_affine_matrix2d, warp_affine + +from lib.pymafx.core import constants def transform_to_tensor(res, mean=None, std=None, is_tensor=False): @@ -40,13 +41,14 @@ def get_affine_matrix_box(boxes, w2, h2): # boxes [left, top, right, bottom] width = boxes[:, 2] - boxes[:, 0] #(N,) height = boxes[:, 3] - boxes[:, 1] #(N,) - center = torch.tensor( - [(boxes[:, 0] + boxes[:, 2]) / 2.0, (boxes[:, 1] + boxes[:, 3]) / 2.0] - ).T #(N,2) + center = torch.tensor([(boxes[:, 0] + boxes[:, 2]) / 2.0, + (boxes[:, 1] + boxes[:, 3]) / 2.0]).T #(N,2) scale = torch.min(torch.tensor([w2 / width, h2 / height]), dim=0)[0].unsqueeze(1).repeat(1, 2) * 0.9 #(N,2) - transl = torch.cat([w2 / 2.0 - center[:, 0:1], h2 / 2.0 - center[:, 1:2]], dim=1) #(N,2) - M = get_affine_matrix2d(transl, center, scale, angle=torch.tensor([0.,]*transl.shape[0])) + transl = torch.cat([w2 / 2.0 - center[:, 0:1], h2 / 2.0 - center[:, 1:2]], dim=1) #(N,2) + M = get_affine_matrix2d(transl, center, scale, angle=torch.tensor([ + 0., + ] * transl.shape[0])) return M @@ -54,12 +56,12 @@ def get_affine_matrix_box(boxes, w2, h2): def load_img(img_file): if img_file.endswith("exr"): - img = cv2.imread(img_file, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) - else : + img = cv2.imread(img_file, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) + else: img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED) # considering non 8-bit image - if img.dtype != np.uint8 : + if img.dtype != np.uint8: img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U) if len(img.shape) == 2: @@ -112,8 +114,8 @@ def get_pymafx(image, landmarks): # image [3,512,512] item = { - 'img_body': - F.interpolate(image.unsqueeze(0), size=224, mode='bicubic', align_corners=True)[0] + 'img_body': F.interpolate(image.unsqueeze(0), size=224, mode='bicubic', + align_corners=True)[0] } for part in ['lhand', 'rhand', 'face']: @@ -211,11 +213,8 @@ def process_image(img_file, hps_type, single, input_res, detector): img_pymafx_lst = [] uncrop_param = { - "ori_shape": [in_height, in_width], - "box_shape": [input_res, input_res], - "square_shape": [tgt_res, tgt_res], - "M_square": M_square, - "M_crop": M_crop + "ori_shape": [in_height, in_width], "box_shape": [input_res, input_res], "square_shape": + [tgt_res, tgt_res], "M_square": M_square, "M_crop": M_crop } for idx in range(len(boxes)): @@ -226,11 +225,11 @@ def process_image(img_file, hps_type, single, input_res, detector): else: mask_detection = masks[0] * 0. - img_square_rgba = torch.cat( - [img_square.squeeze(0).permute(1, 2, 0), - torch.tensor(mask_detection < 0.4) * 255], - dim=2 - ) + img_square_rgba = torch.cat([ + img_square.squeeze(0).permute(1, 2, 0), + torch.tensor(mask_detection < 0.4) * 255 + ], + dim=2) img_crop = warp_affine( img_square_rgba.unsqueeze(0).permute(0, 3, 1, 2), diff --git a/lib/common/libmesh/inside_mesh.py b/lib/common/libmesh/inside_mesh.py index eaac43c2e6fe103c6a1dd4e182642ff0cc6a024a..24fa682633ce0a751054bdf0bdde3fa1f595d0b7 100644 --- a/lib/common/libmesh/inside_mesh.py +++ b/lib/common/libmesh/inside_mesh.py @@ -1,4 +1,5 @@ import numpy as np + from .triangle_hash import TriangleHash as _TriangleHash @@ -147,8 +148,6 @@ class TriangleIntersector2d: v = (-A[:, 1, 0] * y[:, 0] + A[:, 0, 0] * y[:, 1]) * s_detA sum_uv = u + v - contains[mask] = ( - (0 < u) & (u < abs_detA) & (0 < v) & (v < abs_detA) & (0 < sum_uv) & - (sum_uv < abs_detA) - ) + contains[mask] = ((0 < u) & (u < abs_detA) & (0 < v) & (v < abs_detA) & (0 < sum_uv) & + (sum_uv < abs_detA)) return contains diff --git a/lib/common/libmesh/setup.py b/lib/common/libmesh/setup.py index 38ac162300df4e987134e81306e1a6ad674a5323..65a56afc07363ce47320a69da57883e1906f35ab 100644 --- a/lib/common/libmesh/setup.py +++ b/lib/common/libmesh/setup.py @@ -1,5 +1,5 @@ +import numpy from setuptools import setup from Cython.Build import cythonize -import numpy setup(name='libmesh', ext_modules=cythonize("*.pyx"), include_dirs=[numpy.get_include()]) diff --git a/lib/common/libmesh/triangle_hash.cpp b/lib/common/libmesh/triangle_hash.cpp index 0fc7f44a1db876944eeb72eca57c56b038e80dd6..5f46e5e4f372d1c4dd629743e41063cc6e528adf 100644 --- a/lib/common/libmesh/triangle_hash.cpp +++ b/lib/common/libmesh/triangle_hash.cpp @@ -720,12 +720,12 @@ static CYTHON_INLINE float __PYX_NAN() { /* NumPy API declarations from "numpy/__init__.pxd" */ +#include #include "ios" #include "new" #include "stdexcept" #include "typeinfo" #include -#include #include "pythread.h" #include #include "pystate.h" @@ -1330,8 +1330,8 @@ typedef npy_clongdouble __pyx_t_5numpy_clongdouble_t; */ typedef npy_cdouble __pyx_t_5numpy_complex_t; -/* "triangle_hash.pyx":9 - * from libc.math cimport floor, ceil +/* "triangle_hash.pyx":11 + * * * cdef class TriangleHash: # <<<<<<<<<<<<<< * cdef vector[vector[int]] spatial_hash @@ -1423,8 +1423,8 @@ struct __pyx_memoryviewslice_obj { -/* "triangle_hash.pyx":9 - * from libc.math cimport floor, ceil +/* "triangle_hash.pyx":11 + * * * cdef class TriangleHash: # <<<<<<<<<<<<<< * cdef vector[vector[int]] spatial_hash @@ -2279,6 +2279,10 @@ static PyObject *__pyx_memoryview_assign_item_from_object(struct __pyx_memoryvie static PyObject *__pyx_memoryviewslice_convert_item_to_object(struct __pyx_memoryviewslice_obj *__pyx_v_self, char *__pyx_v_itemp); /* proto*/ static PyObject *__pyx_memoryviewslice_assign_item_from_object(struct __pyx_memoryviewslice_obj *__pyx_v_self, char *__pyx_v_itemp, PyObject *__pyx_v_value); /* proto*/ +/* Module declarations from 'cython.view' */ + +/* Module declarations from 'cython' */ + /* Module declarations from 'cpython.buffer' */ /* Module declarations from 'libc.string' */ @@ -2317,14 +2321,10 @@ static PyTypeObject *__pyx_ptype_5numpy_flexible = 0; static PyTypeObject *__pyx_ptype_5numpy_character = 0; static PyTypeObject *__pyx_ptype_5numpy_ufunc = 0; -/* Module declarations from 'cython.view' */ - -/* Module declarations from 'cython' */ +/* Module declarations from 'libc.math' */ /* Module declarations from 'libcpp.vector' */ -/* Module declarations from 'libc.math' */ - /* Module declarations from 'triangle_hash' */ static PyTypeObject *__pyx_ptype_13triangle_hash_TriangleHash = 0; static PyTypeObject *__pyx_array_type = 0; @@ -2667,7 +2667,7 @@ static PyObject *__pyx_tuple__28; static PyObject *__pyx_codeobj__29; /* Late includes */ -/* "triangle_hash.pyx":13 +/* "triangle_hash.pyx":15 * cdef int resolution * * def __cinit__(self, double[:, :, :] triangles, int resolution): # <<<<<<<<<<<<<< @@ -2709,11 +2709,11 @@ static int __pyx_pw_13triangle_hash_12TriangleHash_1__cinit__(PyObject *__pyx_v_ case 1: if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_resolution)) != 0)) kw_args--; else { - __Pyx_RaiseArgtupleInvalid("__cinit__", 1, 2, 2, 1); __PYX_ERR(0, 13, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("__cinit__", 1, 2, 2, 1); __PYX_ERR(0, 15, __pyx_L3_error) } } if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "__cinit__") < 0)) __PYX_ERR(0, 13, __pyx_L3_error) + if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "__cinit__") < 0)) __PYX_ERR(0, 15, __pyx_L3_error) } } else if (PyTuple_GET_SIZE(__pyx_args) != 2) { goto __pyx_L5_argtuple_error; @@ -2721,12 +2721,12 @@ static int __pyx_pw_13triangle_hash_12TriangleHash_1__cinit__(PyObject *__pyx_v_ values[0] = PyTuple_GET_ITEM(__pyx_args, 0); values[1] = PyTuple_GET_ITEM(__pyx_args, 1); } - __pyx_v_triangles = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_double(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_triangles.memview)) __PYX_ERR(0, 13, __pyx_L3_error) - __pyx_v_resolution = __Pyx_PyInt_As_int(values[1]); if (unlikely((__pyx_v_resolution == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 13, __pyx_L3_error) + __pyx_v_triangles = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_double(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_triangles.memview)) __PYX_ERR(0, 15, __pyx_L3_error) + __pyx_v_resolution = __Pyx_PyInt_As_int(values[1]); if (unlikely((__pyx_v_resolution == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 15, __pyx_L3_error) } goto __pyx_L4_argument_unpacking_done; __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("__cinit__", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 13, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("__cinit__", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 15, __pyx_L3_error) __pyx_L3_error:; __Pyx_AddTraceback("triangle_hash.TriangleHash.__cinit__", __pyx_clineno, __pyx_lineno, __pyx_filename); __Pyx_RefNannyFinishContext(); @@ -2747,7 +2747,7 @@ static int __pyx_pf_13triangle_hash_12TriangleHash___cinit__(struct __pyx_obj_13 int __pyx_clineno = 0; __Pyx_RefNannySetupContext("__cinit__", 0); - /* "triangle_hash.pyx":14 + /* "triangle_hash.pyx":16 * * def __cinit__(self, double[:, :, :] triangles, int resolution): * self.spatial_hash.resize(resolution * resolution) # <<<<<<<<<<<<<< @@ -2758,10 +2758,10 @@ static int __pyx_pf_13triangle_hash_12TriangleHash___cinit__(struct __pyx_obj_13 __pyx_v_self->spatial_hash.resize((__pyx_v_resolution * __pyx_v_resolution)); } catch(...) { __Pyx_CppExn2PyErr(); - __PYX_ERR(0, 14, __pyx_L1_error) + __PYX_ERR(0, 16, __pyx_L1_error) } - /* "triangle_hash.pyx":15 + /* "triangle_hash.pyx":17 * def __cinit__(self, double[:, :, :] triangles, int resolution): * self.spatial_hash.resize(resolution * resolution) * self.resolution = resolution # <<<<<<<<<<<<<< @@ -2770,7 +2770,7 @@ static int __pyx_pf_13triangle_hash_12TriangleHash___cinit__(struct __pyx_obj_13 */ __pyx_v_self->resolution = __pyx_v_resolution; - /* "triangle_hash.pyx":16 + /* "triangle_hash.pyx":18 * self.spatial_hash.resize(resolution * resolution) * self.resolution = resolution * self._build_hash(triangles) # <<<<<<<<<<<<<< @@ -2779,7 +2779,7 @@ static int __pyx_pf_13triangle_hash_12TriangleHash___cinit__(struct __pyx_obj_13 */ (void)(((struct __pyx_vtabstruct_13triangle_hash_TriangleHash *)__pyx_v_self->__pyx_vtab)->_build_hash(__pyx_v_self, __pyx_v_triangles)); - /* "triangle_hash.pyx":13 + /* "triangle_hash.pyx":15 * cdef int resolution * * def __cinit__(self, double[:, :, :] triangles, int resolution): # <<<<<<<<<<<<<< @@ -2799,7 +2799,7 @@ static int __pyx_pf_13triangle_hash_12TriangleHash___cinit__(struct __pyx_obj_13 return __pyx_r; } -/* "triangle_hash.pyx":20 +/* "triangle_hash.pyx":22 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cdef int _build_hash(self, double[:, :, :] triangles): # <<<<<<<<<<<<<< @@ -2839,7 +2839,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 int __pyx_clineno = 0; __Pyx_RefNannySetupContext("_build_hash", 0); - /* "triangle_hash.pyx":21 + /* "triangle_hash.pyx":23 * @cython.wraparound(False) # Deactivate negative indexing. * cdef int _build_hash(self, double[:, :, :] triangles): * assert(triangles.shape[1] == 3) # <<<<<<<<<<<<<< @@ -2850,12 +2850,12 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_triangles.shape[1]) == 3) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 21, __pyx_L1_error) + __PYX_ERR(0, 23, __pyx_L1_error) } } #endif - /* "triangle_hash.pyx":22 + /* "triangle_hash.pyx":24 * cdef int _build_hash(self, double[:, :, :] triangles): * assert(triangles.shape[1] == 3) * assert(triangles.shape[2] == 2) # <<<<<<<<<<<<<< @@ -2866,12 +2866,12 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_triangles.shape[2]) == 2) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 22, __pyx_L1_error) + __PYX_ERR(0, 24, __pyx_L1_error) } } #endif - /* "triangle_hash.pyx":24 + /* "triangle_hash.pyx":26 * assert(triangles.shape[2] == 2) * * cdef int n_tri = triangles.shape[0] # <<<<<<<<<<<<<< @@ -2880,7 +2880,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 */ __pyx_v_n_tri = (__pyx_v_triangles.shape[0]); - /* "triangle_hash.pyx":31 + /* "triangle_hash.pyx":33 * cdef int spatial_idx * * for i_tri in range(n_tri): # <<<<<<<<<<<<<< @@ -2892,7 +2892,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 for (__pyx_t_3 = 0; __pyx_t_3 < __pyx_t_2; __pyx_t_3+=1) { __pyx_v_i_tri = __pyx_t_3; - /* "triangle_hash.pyx":33 + /* "triangle_hash.pyx":35 * for i_tri in range(n_tri): * # Compute bounding box * for j in range(2): # <<<<<<<<<<<<<< @@ -2902,7 +2902,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 for (__pyx_t_4 = 0; __pyx_t_4 < 2; __pyx_t_4+=1) { __pyx_v_j = __pyx_t_4; - /* "triangle_hash.pyx":35 + /* "triangle_hash.pyx":37 * for j in range(2): * bbox_min[j] = min( * triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] # <<<<<<<<<<<<<< @@ -2933,7 +2933,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 __pyx_t_11 = __pyx_t_10; } - /* "triangle_hash.pyx":34 + /* "triangle_hash.pyx":36 * # Compute bounding box * for j in range(2): * bbox_min[j] = min( # <<<<<<<<<<<<<< @@ -2942,7 +2942,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 */ (__pyx_v_bbox_min[__pyx_v_j]) = ((int)__pyx_t_11); - /* "triangle_hash.pyx":38 + /* "triangle_hash.pyx":40 * ) * bbox_max[j] = max( * triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] # <<<<<<<<<<<<<< @@ -2973,7 +2973,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 __pyx_t_10 = __pyx_t_9; } - /* "triangle_hash.pyx":37 + /* "triangle_hash.pyx":39 * triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] * ) * bbox_max[j] = max( # <<<<<<<<<<<<<< @@ -2982,7 +2982,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 */ (__pyx_v_bbox_max[__pyx_v_j]) = ((int)__pyx_t_10); - /* "triangle_hash.pyx":40 + /* "triangle_hash.pyx":42 * triangles[i_tri, 0, j], triangles[i_tri, 1, j], triangles[i_tri, 2, j] * ) * bbox_min[j] = min(max(bbox_min[j], 0), self.resolution - 1) # <<<<<<<<<<<<<< @@ -3005,7 +3005,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 } (__pyx_v_bbox_min[__pyx_v_j]) = __pyx_t_15; - /* "triangle_hash.pyx":41 + /* "triangle_hash.pyx":43 * ) * bbox_min[j] = min(max(bbox_min[j], 0), self.resolution - 1) * bbox_max[j] = min(max(bbox_max[j], 0), self.resolution - 1) # <<<<<<<<<<<<<< @@ -3029,7 +3029,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 (__pyx_v_bbox_max[__pyx_v_j]) = __pyx_t_13; } - /* "triangle_hash.pyx":44 + /* "triangle_hash.pyx":46 * * # Find all voxels where bounding box intersects * for x in range(bbox_min[0], bbox_max[0] + 1): # <<<<<<<<<<<<<< @@ -3041,7 +3041,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 for (__pyx_t_4 = (__pyx_v_bbox_min[0]); __pyx_t_4 < __pyx_t_15; __pyx_t_4+=1) { __pyx_v_x = __pyx_t_4; - /* "triangle_hash.pyx":45 + /* "triangle_hash.pyx":47 * # Find all voxels where bounding box intersects * for x in range(bbox_min[0], bbox_max[0] + 1): * for y in range(bbox_min[1], bbox_max[1] + 1): # <<<<<<<<<<<<<< @@ -3053,7 +3053,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 for (__pyx_t_14 = (__pyx_v_bbox_min[1]); __pyx_t_14 < __pyx_t_16; __pyx_t_14+=1) { __pyx_v_y = __pyx_t_14; - /* "triangle_hash.pyx":46 + /* "triangle_hash.pyx":48 * for x in range(bbox_min[0], bbox_max[0] + 1): * for y in range(bbox_min[1], bbox_max[1] + 1): * spatial_idx = self.resolution * x + y # <<<<<<<<<<<<<< @@ -3062,7 +3062,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 */ __pyx_v_spatial_idx = ((__pyx_v_self->resolution * __pyx_v_x) + __pyx_v_y); - /* "triangle_hash.pyx":47 + /* "triangle_hash.pyx":49 * for y in range(bbox_min[1], bbox_max[1] + 1): * spatial_idx = self.resolution * x + y * self.spatial_hash[spatial_idx].push_back(i_tri) # <<<<<<<<<<<<<< @@ -3073,13 +3073,13 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 (__pyx_v_self->spatial_hash[__pyx_v_spatial_idx]).push_back(__pyx_v_i_tri); } catch(...) { __Pyx_CppExn2PyErr(); - __PYX_ERR(0, 47, __pyx_L1_error) + __PYX_ERR(0, 49, __pyx_L1_error) } } } } - /* "triangle_hash.pyx":20 + /* "triangle_hash.pyx":22 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cdef int _build_hash(self, double[:, :, :] triangles): # <<<<<<<<<<<<<< @@ -3098,7 +3098,7 @@ static int __pyx_f_13triangle_hash_12TriangleHash__build_hash(struct __pyx_obj_1 return __pyx_r; } -/* "triangle_hash.pyx":51 +/* "triangle_hash.pyx":53 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef query(self, double[:, :] points): # <<<<<<<<<<<<<< @@ -3155,12 +3155,12 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 if (unlikely(!__Pyx_object_dict_version_matches(((PyObject *)__pyx_v_self), __pyx_tp_dict_version, __pyx_obj_dict_version))) { PY_UINT64_T __pyx_type_dict_guard = __Pyx_get_tp_dict_version(((PyObject *)__pyx_v_self)); #endif - __pyx_t_1 = __Pyx_PyObject_GetAttrStr(((PyObject *)__pyx_v_self), __pyx_n_s_query); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 51, __pyx_L1_error) + __pyx_t_1 = __Pyx_PyObject_GetAttrStr(((PyObject *)__pyx_v_self), __pyx_n_s_query); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 53, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); if (!PyCFunction_Check(__pyx_t_1) || (PyCFunction_GET_FUNCTION(__pyx_t_1) != (PyCFunction)(void*)__pyx_pw_13triangle_hash_12TriangleHash_3query)) { __Pyx_XDECREF(__pyx_r); - if (unlikely(!__pyx_v_points.memview)) { __Pyx_RaiseUnboundLocalError("points"); __PYX_ERR(0, 51, __pyx_L1_error) } - __pyx_t_3 = __pyx_memoryview_fromslice(__pyx_v_points, 2, (PyObject *(*)(char *)) __pyx_memview_get_double, (int (*)(char *, PyObject *)) __pyx_memview_set_double, 0);; if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 51, __pyx_L1_error) + if (unlikely(!__pyx_v_points.memview)) { __Pyx_RaiseUnboundLocalError("points"); __PYX_ERR(0, 53, __pyx_L1_error) } + __pyx_t_3 = __pyx_memoryview_fromslice(__pyx_v_points, 2, (PyObject *(*)(char *)) __pyx_memview_get_double, (int (*)(char *, PyObject *)) __pyx_memview_set_double, 0);; if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 53, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_3); __Pyx_INCREF(__pyx_t_1); __pyx_t_4 = __pyx_t_1; __pyx_t_5 = NULL; @@ -3176,7 +3176,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_t_2 = (__pyx_t_5) ? __Pyx_PyObject_Call2Args(__pyx_t_4, __pyx_t_5, __pyx_t_3) : __Pyx_PyObject_CallOneArg(__pyx_t_4, __pyx_t_3); __Pyx_XDECREF(__pyx_t_5); __pyx_t_5 = 0; __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 51, __pyx_L1_error) + if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 53, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_2); __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; __pyx_r = __pyx_t_2; @@ -3197,7 +3197,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 #endif } - /* "triangle_hash.pyx":52 + /* "triangle_hash.pyx":54 * @cython.wraparound(False) # Deactivate negative indexing. * cpdef query(self, double[:, :] points): * assert(points.shape[1] == 2) # <<<<<<<<<<<<<< @@ -3208,12 +3208,12 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_points.shape[1]) == 2) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 52, __pyx_L1_error) + __PYX_ERR(0, 54, __pyx_L1_error) } } #endif - /* "triangle_hash.pyx":53 + /* "triangle_hash.pyx":55 * cpdef query(self, double[:, :] points): * assert(points.shape[1] == 2) * cdef int n_points = points.shape[0] # <<<<<<<<<<<<<< @@ -3222,7 +3222,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 */ __pyx_v_n_points = (__pyx_v_points.shape[0]); - /* "triangle_hash.pyx":63 + /* "triangle_hash.pyx":65 * cdef int spatial_idx * * for i_point in range(n_points): # <<<<<<<<<<<<<< @@ -3234,7 +3234,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 for (__pyx_t_8 = 0; __pyx_t_8 < __pyx_t_7; __pyx_t_8+=1) { __pyx_v_i_point = __pyx_t_8; - /* "triangle_hash.pyx":64 + /* "triangle_hash.pyx":66 * * for i_point in range(n_points): * x = int(points[i_point, 0]) # <<<<<<<<<<<<<< @@ -3245,7 +3245,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_t_10 = 0; __pyx_v_x = ((int)(*((double *) ( /* dim=1 */ (( /* dim=0 */ (__pyx_v_points.data + __pyx_t_9 * __pyx_v_points.strides[0]) ) + __pyx_t_10 * __pyx_v_points.strides[1]) )))); - /* "triangle_hash.pyx":65 + /* "triangle_hash.pyx":67 * for i_point in range(n_points): * x = int(points[i_point, 0]) * y = int(points[i_point, 1]) # <<<<<<<<<<<<<< @@ -3256,7 +3256,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_t_9 = 1; __pyx_v_y = ((int)(*((double *) ( /* dim=1 */ (( /* dim=0 */ (__pyx_v_points.data + __pyx_t_10 * __pyx_v_points.strides[0]) ) + __pyx_t_9 * __pyx_v_points.strides[1]) )))); - /* "triangle_hash.pyx":66 + /* "triangle_hash.pyx":68 * x = int(points[i_point, 0]) * y = int(points[i_point, 1]) * if not (0 <= x < self.resolution and 0 <= y < self.resolution): # <<<<<<<<<<<<<< @@ -3283,7 +3283,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_t_12 = ((!__pyx_t_11) != 0); if (__pyx_t_12) { - /* "triangle_hash.pyx":67 + /* "triangle_hash.pyx":69 * y = int(points[i_point, 1]) * if not (0 <= x < self.resolution and 0 <= y < self.resolution): * continue # <<<<<<<<<<<<<< @@ -3292,7 +3292,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 */ goto __pyx_L3_continue; - /* "triangle_hash.pyx":66 + /* "triangle_hash.pyx":68 * x = int(points[i_point, 0]) * y = int(points[i_point, 1]) * if not (0 <= x < self.resolution and 0 <= y < self.resolution): # <<<<<<<<<<<<<< @@ -3301,7 +3301,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 */ } - /* "triangle_hash.pyx":69 + /* "triangle_hash.pyx":71 * continue * * spatial_idx = self.resolution * x + y # <<<<<<<<<<<<<< @@ -3310,7 +3310,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 */ __pyx_v_spatial_idx = ((__pyx_v_self->resolution * __pyx_v_x) + __pyx_v_y); - /* "triangle_hash.pyx":70 + /* "triangle_hash.pyx":72 * * spatial_idx = self.resolution * x + y * for i_tri in self.spatial_hash[spatial_idx]: # <<<<<<<<<<<<<< @@ -3325,7 +3325,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 ++__pyx_t_14; __pyx_v_i_tri = __pyx_t_16; - /* "triangle_hash.pyx":71 + /* "triangle_hash.pyx":73 * spatial_idx = self.resolution * x + y * for i_tri in self.spatial_hash[spatial_idx]: * points_indices.push_back(i_point) # <<<<<<<<<<<<<< @@ -3336,10 +3336,10 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_v_points_indices.push_back(__pyx_v_i_point); } catch(...) { __Pyx_CppExn2PyErr(); - __PYX_ERR(0, 71, __pyx_L1_error) + __PYX_ERR(0, 73, __pyx_L1_error) } - /* "triangle_hash.pyx":72 + /* "triangle_hash.pyx":74 * for i_tri in self.spatial_hash[spatial_idx]: * points_indices.push_back(i_point) * tri_indices.push_back(i_tri) # <<<<<<<<<<<<<< @@ -3350,10 +3350,10 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_v_tri_indices.push_back(__pyx_v_i_tri); } catch(...) { __Pyx_CppExn2PyErr(); - __PYX_ERR(0, 72, __pyx_L1_error) + __PYX_ERR(0, 74, __pyx_L1_error) } - /* "triangle_hash.pyx":70 + /* "triangle_hash.pyx":72 * * spatial_idx = self.resolution * x + y * for i_tri in self.spatial_hash[spatial_idx]: # <<<<<<<<<<<<<< @@ -3364,35 +3364,35 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_L3_continue:; } - /* "triangle_hash.pyx":74 + /* "triangle_hash.pyx":76 * tri_indices.push_back(i_tri) * * points_indices_np = np.zeros(points_indices.size(), dtype=np.int32) # <<<<<<<<<<<<<< * tri_indices_np = np.zeros(tri_indices.size(), dtype=np.int32) * */ - __Pyx_GetModuleGlobalName(__pyx_t_1, __pyx_n_s_np); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 74, __pyx_L1_error) + __Pyx_GetModuleGlobalName(__pyx_t_1, __pyx_n_s_np); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); - __pyx_t_2 = __Pyx_PyObject_GetAttrStr(__pyx_t_1, __pyx_n_s_zeros); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_2 = __Pyx_PyObject_GetAttrStr(__pyx_t_1, __pyx_n_s_zeros); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_2); __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - __pyx_t_1 = __Pyx_PyInt_FromSize_t(__pyx_v_points_indices.size()); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_1 = __Pyx_PyInt_FromSize_t(__pyx_v_points_indices.size()); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); - __pyx_t_4 = PyTuple_New(1); if (unlikely(!__pyx_t_4)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_4 = PyTuple_New(1); if (unlikely(!__pyx_t_4)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_4); __Pyx_GIVEREF(__pyx_t_1); PyTuple_SET_ITEM(__pyx_t_4, 0, __pyx_t_1); __pyx_t_1 = 0; - __pyx_t_1 = __Pyx_PyDict_NewPresized(1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_1 = __Pyx_PyDict_NewPresized(1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); - __Pyx_GetModuleGlobalName(__pyx_t_3, __pyx_n_s_np); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 74, __pyx_L1_error) + __Pyx_GetModuleGlobalName(__pyx_t_3, __pyx_n_s_np); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_3); - __pyx_t_5 = __Pyx_PyObject_GetAttrStr(__pyx_t_3, __pyx_n_s_int32); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_5 = __Pyx_PyObject_GetAttrStr(__pyx_t_3, __pyx_n_s_int32); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_5); __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (PyDict_SetItem(__pyx_t_1, __pyx_n_s_dtype, __pyx_t_5) < 0) __PYX_ERR(0, 74, __pyx_L1_error) + if (PyDict_SetItem(__pyx_t_1, __pyx_n_s_dtype, __pyx_t_5) < 0) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_DECREF(__pyx_t_5); __pyx_t_5 = 0; - __pyx_t_5 = __Pyx_PyObject_Call(__pyx_t_2, __pyx_t_4, __pyx_t_1); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 74, __pyx_L1_error) + __pyx_t_5 = __Pyx_PyObject_Call(__pyx_t_2, __pyx_t_4, __pyx_t_1); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 76, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_5); __Pyx_DECREF(__pyx_t_2); __pyx_t_2 = 0; __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; @@ -3400,35 +3400,35 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_v_points_indices_np = __pyx_t_5; __pyx_t_5 = 0; - /* "triangle_hash.pyx":75 + /* "triangle_hash.pyx":77 * * points_indices_np = np.zeros(points_indices.size(), dtype=np.int32) * tri_indices_np = np.zeros(tri_indices.size(), dtype=np.int32) # <<<<<<<<<<<<<< * * cdef int[:] points_indices_view = points_indices_np */ - __Pyx_GetModuleGlobalName(__pyx_t_5, __pyx_n_s_np); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 75, __pyx_L1_error) + __Pyx_GetModuleGlobalName(__pyx_t_5, __pyx_n_s_np); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_5); - __pyx_t_1 = __Pyx_PyObject_GetAttrStr(__pyx_t_5, __pyx_n_s_zeros); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_1 = __Pyx_PyObject_GetAttrStr(__pyx_t_5, __pyx_n_s_zeros); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); __Pyx_DECREF(__pyx_t_5); __pyx_t_5 = 0; - __pyx_t_5 = __Pyx_PyInt_FromSize_t(__pyx_v_tri_indices.size()); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_5 = __Pyx_PyInt_FromSize_t(__pyx_v_tri_indices.size()); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_5); - __pyx_t_4 = PyTuple_New(1); if (unlikely(!__pyx_t_4)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_4 = PyTuple_New(1); if (unlikely(!__pyx_t_4)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_4); __Pyx_GIVEREF(__pyx_t_5); PyTuple_SET_ITEM(__pyx_t_4, 0, __pyx_t_5); __pyx_t_5 = 0; - __pyx_t_5 = __Pyx_PyDict_NewPresized(1); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_5 = __Pyx_PyDict_NewPresized(1); if (unlikely(!__pyx_t_5)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_5); - __Pyx_GetModuleGlobalName(__pyx_t_2, __pyx_n_s_np); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 75, __pyx_L1_error) + __Pyx_GetModuleGlobalName(__pyx_t_2, __pyx_n_s_np); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_2); - __pyx_t_3 = __Pyx_PyObject_GetAttrStr(__pyx_t_2, __pyx_n_s_int32); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_3 = __Pyx_PyObject_GetAttrStr(__pyx_t_2, __pyx_n_s_int32); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_3); __Pyx_DECREF(__pyx_t_2); __pyx_t_2 = 0; - if (PyDict_SetItem(__pyx_t_5, __pyx_n_s_dtype, __pyx_t_3) < 0) __PYX_ERR(0, 75, __pyx_L1_error) + if (PyDict_SetItem(__pyx_t_5, __pyx_n_s_dtype, __pyx_t_3) < 0) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_t_1, __pyx_t_4, __pyx_t_5); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 75, __pyx_L1_error) + __pyx_t_3 = __Pyx_PyObject_Call(__pyx_t_1, __pyx_t_4, __pyx_t_5); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 77, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_3); __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; @@ -3436,31 +3436,31 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_v_tri_indices_np = __pyx_t_3; __pyx_t_3 = 0; - /* "triangle_hash.pyx":77 + /* "triangle_hash.pyx":79 * tri_indices_np = np.zeros(tri_indices.size(), dtype=np.int32) * * cdef int[:] points_indices_view = points_indices_np # <<<<<<<<<<<<<< * cdef int[:] tri_indices_view = tri_indices_np * */ - __pyx_t_17 = __Pyx_PyObject_to_MemoryviewSlice_ds_int(__pyx_v_points_indices_np, PyBUF_WRITABLE); if (unlikely(!__pyx_t_17.memview)) __PYX_ERR(0, 77, __pyx_L1_error) + __pyx_t_17 = __Pyx_PyObject_to_MemoryviewSlice_ds_int(__pyx_v_points_indices_np, PyBUF_WRITABLE); if (unlikely(!__pyx_t_17.memview)) __PYX_ERR(0, 79, __pyx_L1_error) __pyx_v_points_indices_view = __pyx_t_17; __pyx_t_17.memview = NULL; __pyx_t_17.data = NULL; - /* "triangle_hash.pyx":78 + /* "triangle_hash.pyx":80 * * cdef int[:] points_indices_view = points_indices_np * cdef int[:] tri_indices_view = tri_indices_np # <<<<<<<<<<<<<< * * for k in range(points_indices.size()): */ - __pyx_t_17 = __Pyx_PyObject_to_MemoryviewSlice_ds_int(__pyx_v_tri_indices_np, PyBUF_WRITABLE); if (unlikely(!__pyx_t_17.memview)) __PYX_ERR(0, 78, __pyx_L1_error) + __pyx_t_17 = __Pyx_PyObject_to_MemoryviewSlice_ds_int(__pyx_v_tri_indices_np, PyBUF_WRITABLE); if (unlikely(!__pyx_t_17.memview)) __PYX_ERR(0, 80, __pyx_L1_error) __pyx_v_tri_indices_view = __pyx_t_17; __pyx_t_17.memview = NULL; __pyx_t_17.data = NULL; - /* "triangle_hash.pyx":80 + /* "triangle_hash.pyx":82 * cdef int[:] tri_indices_view = tri_indices_np * * for k in range(points_indices.size()): # <<<<<<<<<<<<<< @@ -3472,7 +3472,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 for (__pyx_t_6 = 0; __pyx_t_6 < __pyx_t_19; __pyx_t_6+=1) { __pyx_v_k = __pyx_t_6; - /* "triangle_hash.pyx":81 + /* "triangle_hash.pyx":83 * * for k in range(points_indices.size()): * points_indices_view[k] = points_indices[k] # <<<<<<<<<<<<<< @@ -3483,7 +3483,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 *((int *) ( /* dim=0 */ (__pyx_v_points_indices_view.data + __pyx_t_9 * __pyx_v_points_indices_view.strides[0]) )) = (__pyx_v_points_indices[__pyx_v_k]); } - /* "triangle_hash.pyx":83 + /* "triangle_hash.pyx":85 * points_indices_view[k] = points_indices[k] * * for k in range(tri_indices.size()): # <<<<<<<<<<<<<< @@ -3495,7 +3495,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 for (__pyx_t_6 = 0; __pyx_t_6 < __pyx_t_19; __pyx_t_6+=1) { __pyx_v_k = __pyx_t_6; - /* "triangle_hash.pyx":84 + /* "triangle_hash.pyx":86 * * for k in range(tri_indices.size()): * tri_indices_view[k] = tri_indices[k] # <<<<<<<<<<<<<< @@ -3506,13 +3506,13 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 *((int *) ( /* dim=0 */ (__pyx_v_tri_indices_view.data + __pyx_t_9 * __pyx_v_tri_indices_view.strides[0]) )) = (__pyx_v_tri_indices[__pyx_v_k]); } - /* "triangle_hash.pyx":86 + /* "triangle_hash.pyx":88 * tri_indices_view[k] = tri_indices[k] * * return points_indices_np, tri_indices_np # <<<<<<<<<<<<<< */ __Pyx_XDECREF(__pyx_r); - __pyx_t_3 = PyTuple_New(2); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 86, __pyx_L1_error) + __pyx_t_3 = PyTuple_New(2); if (unlikely(!__pyx_t_3)) __PYX_ERR(0, 88, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_3); __Pyx_INCREF(__pyx_v_points_indices_np); __Pyx_GIVEREF(__pyx_v_points_indices_np); @@ -3524,7 +3524,7 @@ static PyObject *__pyx_f_13triangle_hash_12TriangleHash_query(struct __pyx_obj_1 __pyx_t_3 = 0; goto __pyx_L0; - /* "triangle_hash.pyx":51 + /* "triangle_hash.pyx":53 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef query(self, double[:, :] points): # <<<<<<<<<<<<<< @@ -3563,7 +3563,7 @@ static PyObject *__pyx_pw_13triangle_hash_12TriangleHash_3query(PyObject *__pyx_ __Pyx_RefNannyDeclarations __Pyx_RefNannySetupContext("query (wrapper)", 0); assert(__pyx_arg_points); { - __pyx_v_points = __Pyx_PyObject_to_MemoryviewSlice_dsds_double(__pyx_arg_points, PyBUF_WRITABLE); if (unlikely(!__pyx_v_points.memview)) __PYX_ERR(0, 51, __pyx_L3_error) + __pyx_v_points = __Pyx_PyObject_to_MemoryviewSlice_dsds_double(__pyx_arg_points, PyBUF_WRITABLE); if (unlikely(!__pyx_v_points.memview)) __PYX_ERR(0, 53, __pyx_L3_error) } goto __pyx_L4_argument_unpacking_done; __pyx_L3_error:; @@ -3587,8 +3587,8 @@ static PyObject *__pyx_pf_13triangle_hash_12TriangleHash_2query(struct __pyx_obj int __pyx_clineno = 0; __Pyx_RefNannySetupContext("query", 0); __Pyx_XDECREF(__pyx_r); - if (unlikely(!__pyx_v_points.memview)) { __Pyx_RaiseUnboundLocalError("points"); __PYX_ERR(0, 51, __pyx_L1_error) } - __pyx_t_1 = __pyx_f_13triangle_hash_12TriangleHash_query(__pyx_v_self, __pyx_v_points, 1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 51, __pyx_L1_error) + if (unlikely(!__pyx_v_points.memview)) { __Pyx_RaiseUnboundLocalError("points"); __PYX_ERR(0, 53, __pyx_L1_error) } + __pyx_t_1 = __pyx_f_13triangle_hash_12TriangleHash_query(__pyx_v_self, __pyx_v_points, 1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 53, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); __pyx_r = __pyx_t_1; __pyx_t_1 = 0; @@ -18738,7 +18738,7 @@ static __Pyx_StringTabEntry __pyx_string_tab[] = { {0, 0, 0, 0, 0, 0, 0} }; static CYTHON_SMALL_CODE int __Pyx_InitCachedBuiltins(void) { - __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 31, __pyx_L1_error) + __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 33, __pyx_L1_error) __pyx_builtin_TypeError = __Pyx_GetBuiltinName(__pyx_n_s_TypeError); if (!__pyx_builtin_TypeError) __PYX_ERR(1, 2, __pyx_L1_error) __pyx_builtin_ImportError = __Pyx_GetBuiltinName(__pyx_n_s_ImportError); if (!__pyx_builtin_ImportError) __PYX_ERR(2, 945, __pyx_L1_error) __pyx_builtin_ValueError = __Pyx_GetBuiltinName(__pyx_n_s_ValueError); if (!__pyx_builtin_ValueError) __PYX_ERR(1, 133, __pyx_L1_error) @@ -19118,16 +19118,16 @@ static int __Pyx_modinit_type_init_code(void) { __pyx_vtabptr_13triangle_hash_TriangleHash = &__pyx_vtable_13triangle_hash_TriangleHash; __pyx_vtable_13triangle_hash_TriangleHash._build_hash = (int (*)(struct __pyx_obj_13triangle_hash_TriangleHash *, __Pyx_memviewslice))__pyx_f_13triangle_hash_12TriangleHash__build_hash; __pyx_vtable_13triangle_hash_TriangleHash.query = (PyObject *(*)(struct __pyx_obj_13triangle_hash_TriangleHash *, __Pyx_memviewslice, int __pyx_skip_dispatch))__pyx_f_13triangle_hash_12TriangleHash_query; - if (PyType_Ready(&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 9, __pyx_L1_error) + if (PyType_Ready(&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 11, __pyx_L1_error) #if PY_VERSION_HEX < 0x030800B1 __pyx_type_13triangle_hash_TriangleHash.tp_print = 0; #endif if ((CYTHON_USE_TYPE_SLOTS && CYTHON_USE_PYTYPE_LOOKUP) && likely(!__pyx_type_13triangle_hash_TriangleHash.tp_dictoffset && __pyx_type_13triangle_hash_TriangleHash.tp_getattro == PyObject_GenericGetAttr)) { __pyx_type_13triangle_hash_TriangleHash.tp_getattro = __Pyx_PyObject_GenericGetAttr; } - if (__Pyx_SetVtable(__pyx_type_13triangle_hash_TriangleHash.tp_dict, __pyx_vtabptr_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 9, __pyx_L1_error) - if (PyObject_SetAttr(__pyx_m, __pyx_n_s_TriangleHash, (PyObject *)&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 9, __pyx_L1_error) - if (__Pyx_setup_reduce((PyObject*)&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 9, __pyx_L1_error) + if (__Pyx_SetVtable(__pyx_type_13triangle_hash_TriangleHash.tp_dict, __pyx_vtabptr_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 11, __pyx_L1_error) + if (PyObject_SetAttr(__pyx_m, __pyx_n_s_TriangleHash, (PyObject *)&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 11, __pyx_L1_error) + if (__Pyx_setup_reduce((PyObject*)&__pyx_type_13triangle_hash_TriangleHash) < 0) __PYX_ERR(0, 11, __pyx_L1_error) __pyx_ptype_13triangle_hash_TriangleHash = &__pyx_type_13triangle_hash_TriangleHash; __pyx_vtabptr_array = &__pyx_vtable_array; __pyx_vtable_array.get_memview = (PyObject *(*)(struct __pyx_array_obj *))__pyx_array_get_memview; @@ -19468,7 +19468,7 @@ if (!__Pyx_RefNanny) { * * # distutils: language=c++ * import numpy as np # <<<<<<<<<<<<<< - * cimport numpy as np + * * cimport cython */ __pyx_t_1 = __Pyx_Import(__pyx_n_s_numpy, 0, -1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 3, __pyx_L1_error) @@ -19480,7 +19480,7 @@ if (!__Pyx_RefNanny) { * * # distutils: language=c++ # <<<<<<<<<<<<<< * import numpy as np - * cimport numpy as np + * */ __pyx_t_1 = __Pyx_PyDict_NewPresized(0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 2, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); diff --git a/lib/common/libmesh/triangle_hash.pyx b/lib/common/libmesh/triangle_hash.pyx index 9e3ad59bd7bfc35a3e92e0b40e0349bfb3759a4b..6d6da0a763044307899100e3fa0f849ba2e0444b 100644 --- a/lib/common/libmesh/triangle_hash.pyx +++ b/lib/common/libmesh/triangle_hash.pyx @@ -1,10 +1,12 @@ # distutils: language=c++ import numpy as np -cimport numpy as np + cimport cython +cimport numpy as np +from libc.math cimport ceil, floor from libcpp.vector cimport vector -from libc.math cimport floor, ceil + cdef class TriangleHash: cdef vector[vector[int]] spatial_hash diff --git a/lib/common/libvoxelize/voxelize.c b/lib/common/libvoxelize/voxelize.c index 7320c8246e2cec1892352b4e4208765868af9651..11d0d3b5c00684dd4cd6b5268462adb83d7c89ad 100644 --- a/lib/common/libvoxelize/voxelize.c +++ b/lib/common/libvoxelize/voxelize.c @@ -2115,7 +2115,7 @@ static PyObject *__pyx_tuple__24; static PyObject *__pyx_codeobj__25; /* Late includes */ -/* "voxelize.pyx":12 +/* "voxelize.pyx":13 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): # <<<<<<<<<<<<<< @@ -2138,7 +2138,7 @@ static int __pyx_f_8voxelize_voxelize_mesh_(__Pyx_memviewslice __pyx_v_occ, __Py int __pyx_clineno = 0; __Pyx_RefNannySetupContext("voxelize_mesh_", 0); - /* "voxelize.pyx":13 + /* "voxelize.pyx":14 * @cython.wraparound(False) # Deactivate negative indexing. * cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): * assert(faces.shape[1] == 3) # <<<<<<<<<<<<<< @@ -2149,12 +2149,12 @@ static int __pyx_f_8voxelize_voxelize_mesh_(__Pyx_memviewslice __pyx_v_occ, __Py if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_faces.shape[1]) == 3) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 13, __pyx_L1_error) + __PYX_ERR(0, 14, __pyx_L1_error) } } #endif - /* "voxelize.pyx":14 + /* "voxelize.pyx":15 * cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): * assert(faces.shape[1] == 3) * assert(faces.shape[2] == 3) # <<<<<<<<<<<<<< @@ -2165,12 +2165,12 @@ static int __pyx_f_8voxelize_voxelize_mesh_(__Pyx_memviewslice __pyx_v_occ, __Py if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_faces.shape[2]) == 3) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 14, __pyx_L1_error) + __PYX_ERR(0, 15, __pyx_L1_error) } } #endif - /* "voxelize.pyx":16 + /* "voxelize.pyx":17 * assert(faces.shape[2] == 3) * * n_faces = faces.shape[0] # <<<<<<<<<<<<<< @@ -2179,7 +2179,7 @@ static int __pyx_f_8voxelize_voxelize_mesh_(__Pyx_memviewslice __pyx_v_occ, __Py */ __pyx_v_n_faces = (__pyx_v_faces.shape[0]); - /* "voxelize.pyx":18 + /* "voxelize.pyx":19 * n_faces = faces.shape[0] * cdef int i * for i in range(n_faces): # <<<<<<<<<<<<<< @@ -2191,7 +2191,7 @@ static int __pyx_f_8voxelize_voxelize_mesh_(__Pyx_memviewslice __pyx_v_occ, __Py for (__pyx_t_3 = 0; __pyx_t_3 < __pyx_t_2; __pyx_t_3+=1) { __pyx_v_i = __pyx_t_3; - /* "voxelize.pyx":19 + /* "voxelize.pyx":20 * cdef int i * for i in range(n_faces): * voxelize_triangle_(occ, faces[i]) # <<<<<<<<<<<<<< @@ -2221,7 +2221,7 @@ __pyx_t_4.strides[1] = __pyx_v_faces.strides[2]; __pyx_t_4.data = NULL; } - /* "voxelize.pyx":12 + /* "voxelize.pyx":13 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef int voxelize_mesh_(bint[:, :, :] occ, float[:, :, ::1] faces): # <<<<<<<<<<<<<< @@ -2275,11 +2275,11 @@ static PyObject *__pyx_pw_8voxelize_1voxelize_mesh_(PyObject *__pyx_self, PyObje case 1: if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_faces)) != 0)) kw_args--; else { - __Pyx_RaiseArgtupleInvalid("voxelize_mesh_", 1, 2, 2, 1); __PYX_ERR(0, 12, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("voxelize_mesh_", 1, 2, 2, 1); __PYX_ERR(0, 13, __pyx_L3_error) } } if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "voxelize_mesh_") < 0)) __PYX_ERR(0, 12, __pyx_L3_error) + if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "voxelize_mesh_") < 0)) __PYX_ERR(0, 13, __pyx_L3_error) } } else if (PyTuple_GET_SIZE(__pyx_args) != 2) { goto __pyx_L5_argtuple_error; @@ -2287,12 +2287,12 @@ static PyObject *__pyx_pw_8voxelize_1voxelize_mesh_(PyObject *__pyx_self, PyObje values[0] = PyTuple_GET_ITEM(__pyx_args, 0); values[1] = PyTuple_GET_ITEM(__pyx_args, 1); } - __pyx_v_occ = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_occ.memview)) __PYX_ERR(0, 12, __pyx_L3_error) - __pyx_v_faces = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_faces.memview)) __PYX_ERR(0, 12, __pyx_L3_error) + __pyx_v_occ = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_occ.memview)) __PYX_ERR(0, 13, __pyx_L3_error) + __pyx_v_faces = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_faces.memview)) __PYX_ERR(0, 13, __pyx_L3_error) } goto __pyx_L4_argument_unpacking_done; __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("voxelize_mesh_", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 12, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("voxelize_mesh_", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 13, __pyx_L3_error) __pyx_L3_error:; __Pyx_AddTraceback("voxelize.voxelize_mesh_", __pyx_clineno, __pyx_lineno, __pyx_filename); __Pyx_RefNannyFinishContext(); @@ -2314,9 +2314,9 @@ static PyObject *__pyx_pf_8voxelize_voxelize_mesh_(CYTHON_UNUSED PyObject *__pyx int __pyx_clineno = 0; __Pyx_RefNannySetupContext("voxelize_mesh_", 0); __Pyx_XDECREF(__pyx_r); - if (unlikely(!__pyx_v_occ.memview)) { __Pyx_RaiseUnboundLocalError("occ"); __PYX_ERR(0, 12, __pyx_L1_error) } - if (unlikely(!__pyx_v_faces.memview)) { __Pyx_RaiseUnboundLocalError("faces"); __PYX_ERR(0, 12, __pyx_L1_error) } - __pyx_t_1 = __Pyx_PyInt_From_int(__pyx_f_8voxelize_voxelize_mesh_(__pyx_v_occ, __pyx_v_faces, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 12, __pyx_L1_error) + if (unlikely(!__pyx_v_occ.memview)) { __Pyx_RaiseUnboundLocalError("occ"); __PYX_ERR(0, 13, __pyx_L1_error) } + if (unlikely(!__pyx_v_faces.memview)) { __Pyx_RaiseUnboundLocalError("faces"); __PYX_ERR(0, 13, __pyx_L1_error) } + __pyx_t_1 = __Pyx_PyInt_From_int(__pyx_f_8voxelize_voxelize_mesh_(__pyx_v_occ, __pyx_v_faces, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 13, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); __pyx_r = __pyx_t_1; __pyx_t_1 = 0; @@ -2335,7 +2335,7 @@ static PyObject *__pyx_pf_8voxelize_voxelize_mesh_(CYTHON_UNUSED PyObject *__pyx return __pyx_r; } -/* "voxelize.pyx":24 +/* "voxelize.pyx":25 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef int voxelize_triangle_(bint[:, :, :] occupancies, float[:, ::1] triverts): # <<<<<<<<<<<<<< @@ -2382,7 +2382,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup Py_ssize_t __pyx_t_25; __Pyx_RefNannySetupContext("voxelize_triangle_", 0); - /* "voxelize.pyx":32 + /* "voxelize.pyx":33 * cdef bint intersection * * boxhalfsize[:] = (0.5, 0.5, 0.5) # <<<<<<<<<<<<<< @@ -2397,7 +2397,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup (__pyx_t_1[1]) = __pyx_t_3; (__pyx_t_1[2]) = __pyx_t_4; - /* "voxelize.pyx":34 + /* "voxelize.pyx":35 * boxhalfsize[:] = (0.5, 0.5, 0.5) * * for i in range(3): # <<<<<<<<<<<<<< @@ -2407,7 +2407,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup for (__pyx_t_5 = 0; __pyx_t_5 < 3; __pyx_t_5+=1) { __pyx_v_i = __pyx_t_5; - /* "voxelize.pyx":36 + /* "voxelize.pyx":37 * for i in range(3): * bbox_min[i] = ( * min(triverts[0, i], triverts[1, i], triverts[2, i]) # <<<<<<<<<<<<<< @@ -2435,7 +2435,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup __pyx_t_11 = __pyx_t_10; } - /* "voxelize.pyx":35 + /* "voxelize.pyx":36 * * for i in range(3): * bbox_min[i] = ( # <<<<<<<<<<<<<< @@ -2444,7 +2444,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup */ (__pyx_v_bbox_min[__pyx_v_i]) = ((int)__pyx_t_11); - /* "voxelize.pyx":38 + /* "voxelize.pyx":39 * min(triverts[0, i], triverts[1, i], triverts[2, i]) * ) * bbox_min[i] = min(max(bbox_min[i], 0), occupancies.shape[i] - 1) # <<<<<<<<<<<<<< @@ -2468,7 +2468,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup (__pyx_v_bbox_min[__pyx_v_i]) = __pyx_t_16; } - /* "voxelize.pyx":40 + /* "voxelize.pyx":41 * bbox_min[i] = min(max(bbox_min[i], 0), occupancies.shape[i] - 1) * * for i in range(3): # <<<<<<<<<<<<<< @@ -2478,7 +2478,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup for (__pyx_t_5 = 0; __pyx_t_5 < 3; __pyx_t_5+=1) { __pyx_v_i = __pyx_t_5; - /* "voxelize.pyx":42 + /* "voxelize.pyx":43 * for i in range(3): * bbox_max[i] = ( * max(triverts[0, i], triverts[1, i], triverts[2, i]) # <<<<<<<<<<<<<< @@ -2506,7 +2506,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup __pyx_t_10 = __pyx_t_9; } - /* "voxelize.pyx":41 + /* "voxelize.pyx":42 * * for i in range(3): * bbox_max[i] = ( # <<<<<<<<<<<<<< @@ -2515,7 +2515,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup */ (__pyx_v_bbox_max[__pyx_v_i]) = ((int)__pyx_t_10); - /* "voxelize.pyx":44 + /* "voxelize.pyx":45 * max(triverts[0, i], triverts[1, i], triverts[2, i]) * ) * bbox_max[i] = min(max(bbox_max[i], 0), occupancies.shape[i] - 1) # <<<<<<<<<<<<<< @@ -2539,7 +2539,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup (__pyx_v_bbox_max[__pyx_v_i]) = __pyx_t_12; } - /* "voxelize.pyx":46 + /* "voxelize.pyx":47 * bbox_max[i] = min(max(bbox_max[i], 0), occupancies.shape[i] - 1) * * for i in range(bbox_min[0], bbox_max[0] + 1): # <<<<<<<<<<<<<< @@ -2551,7 +2551,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup for (__pyx_t_5 = (__pyx_v_bbox_min[0]); __pyx_t_5 < __pyx_t_15; __pyx_t_5+=1) { __pyx_v_i = __pyx_t_5; - /* "voxelize.pyx":47 + /* "voxelize.pyx":48 * * for i in range(bbox_min[0], bbox_max[0] + 1): * for j in range(bbox_min[1], bbox_max[1] + 1): # <<<<<<<<<<<<<< @@ -2563,7 +2563,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup for (__pyx_t_14 = (__pyx_v_bbox_min[1]); __pyx_t_14 < __pyx_t_18; __pyx_t_14+=1) { __pyx_v_j = __pyx_t_14; - /* "voxelize.pyx":48 + /* "voxelize.pyx":49 * for i in range(bbox_min[0], bbox_max[0] + 1): * for j in range(bbox_min[1], bbox_max[1] + 1): * for k in range(bbox_min[2], bbox_max[2] + 1): # <<<<<<<<<<<<<< @@ -2575,7 +2575,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup for (__pyx_t_21 = (__pyx_v_bbox_min[2]); __pyx_t_21 < __pyx_t_20; __pyx_t_21+=1) { __pyx_v_k = __pyx_t_21; - /* "voxelize.pyx":49 + /* "voxelize.pyx":50 * for j in range(bbox_min[1], bbox_max[1] + 1): * for k in range(bbox_min[2], bbox_max[2] + 1): * boxcenter[:] = (i + 0.5, j + 0.5, k + 0.5) # <<<<<<<<<<<<<< @@ -2590,7 +2590,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup (__pyx_t_1[1]) = __pyx_t_3; (__pyx_t_1[2]) = __pyx_t_2; - /* "voxelize.pyx":51 + /* "voxelize.pyx":52 * boxcenter[:] = (i + 0.5, j + 0.5, k + 0.5) * intersection = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], * &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) # <<<<<<<<<<<<<< @@ -2604,7 +2604,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup __pyx_t_24 = 2; __pyx_t_25 = 0; - /* "voxelize.pyx":50 + /* "voxelize.pyx":51 * for k in range(bbox_min[2], bbox_max[2] + 1): * boxcenter[:] = (i + 0.5, j + 0.5, k + 0.5) * intersection = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], # <<<<<<<<<<<<<< @@ -2613,7 +2613,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup */ __pyx_v_intersection = triBoxOverlap((&(__pyx_v_boxcenter[0])), (&(__pyx_v_boxhalfsize[0])), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_6 * __pyx_v_triverts.strides[0]) )) + __pyx_t_7)) )))), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_22 * __pyx_v_triverts.strides[0]) )) + __pyx_t_23)) )))), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_24 * __pyx_v_triverts.strides[0]) )) + __pyx_t_25)) ))))); - /* "voxelize.pyx":52 + /* "voxelize.pyx":53 * intersection = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], * &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) * occupancies[i, j, k] |= intersection # <<<<<<<<<<<<<< @@ -2628,7 +2628,7 @@ static int __pyx_f_8voxelize_voxelize_triangle_(__Pyx_memviewslice __pyx_v_occup } } - /* "voxelize.pyx":24 + /* "voxelize.pyx":25 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cpdef int voxelize_triangle_(bint[:, :, :] occupancies, float[:, ::1] triverts): # <<<<<<<<<<<<<< @@ -2676,11 +2676,11 @@ static PyObject *__pyx_pw_8voxelize_3voxelize_triangle_(PyObject *__pyx_self, Py case 1: if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triverts)) != 0)) kw_args--; else { - __Pyx_RaiseArgtupleInvalid("voxelize_triangle_", 1, 2, 2, 1); __PYX_ERR(0, 24, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("voxelize_triangle_", 1, 2, 2, 1); __PYX_ERR(0, 25, __pyx_L3_error) } } if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "voxelize_triangle_") < 0)) __PYX_ERR(0, 24, __pyx_L3_error) + if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "voxelize_triangle_") < 0)) __PYX_ERR(0, 25, __pyx_L3_error) } } else if (PyTuple_GET_SIZE(__pyx_args) != 2) { goto __pyx_L5_argtuple_error; @@ -2688,12 +2688,12 @@ static PyObject *__pyx_pw_8voxelize_3voxelize_triangle_(PyObject *__pyx_self, Py values[0] = PyTuple_GET_ITEM(__pyx_args, 0); values[1] = PyTuple_GET_ITEM(__pyx_args, 1); } - __pyx_v_occupancies = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_occupancies.memview)) __PYX_ERR(0, 24, __pyx_L3_error) - __pyx_v_triverts = __Pyx_PyObject_to_MemoryviewSlice_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_triverts.memview)) __PYX_ERR(0, 24, __pyx_L3_error) + __pyx_v_occupancies = __Pyx_PyObject_to_MemoryviewSlice_dsdsds_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_occupancies.memview)) __PYX_ERR(0, 25, __pyx_L3_error) + __pyx_v_triverts = __Pyx_PyObject_to_MemoryviewSlice_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_triverts.memview)) __PYX_ERR(0, 25, __pyx_L3_error) } goto __pyx_L4_argument_unpacking_done; __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("voxelize_triangle_", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 24, __pyx_L3_error) + __Pyx_RaiseArgtupleInvalid("voxelize_triangle_", 1, 2, 2, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 25, __pyx_L3_error) __pyx_L3_error:; __Pyx_AddTraceback("voxelize.voxelize_triangle_", __pyx_clineno, __pyx_lineno, __pyx_filename); __Pyx_RefNannyFinishContext(); @@ -2715,9 +2715,9 @@ static PyObject *__pyx_pf_8voxelize_2voxelize_triangle_(CYTHON_UNUSED PyObject * int __pyx_clineno = 0; __Pyx_RefNannySetupContext("voxelize_triangle_", 0); __Pyx_XDECREF(__pyx_r); - if (unlikely(!__pyx_v_occupancies.memview)) { __Pyx_RaiseUnboundLocalError("occupancies"); __PYX_ERR(0, 24, __pyx_L1_error) } - if (unlikely(!__pyx_v_triverts.memview)) { __Pyx_RaiseUnboundLocalError("triverts"); __PYX_ERR(0, 24, __pyx_L1_error) } - __pyx_t_1 = __Pyx_PyInt_From_int(__pyx_f_8voxelize_voxelize_triangle_(__pyx_v_occupancies, __pyx_v_triverts, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 24, __pyx_L1_error) + if (unlikely(!__pyx_v_occupancies.memview)) { __Pyx_RaiseUnboundLocalError("occupancies"); __PYX_ERR(0, 25, __pyx_L1_error) } + if (unlikely(!__pyx_v_triverts.memview)) { __Pyx_RaiseUnboundLocalError("triverts"); __PYX_ERR(0, 25, __pyx_L1_error) } + __pyx_t_1 = __Pyx_PyInt_From_int(__pyx_f_8voxelize_voxelize_triangle_(__pyx_v_occupancies, __pyx_v_triverts, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 25, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); __pyx_r = __pyx_t_1; __pyx_t_1 = 0; @@ -2736,7 +2736,7 @@ static PyObject *__pyx_pf_8voxelize_2voxelize_triangle_(CYTHON_UNUSED PyObject * return __pyx_r; } -/* "voxelize.pyx":57 +/* "voxelize.pyx":58 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): # <<<<<<<<<<<<<< @@ -2762,7 +2762,7 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce int __pyx_clineno = 0; __Pyx_RefNannySetupContext("test_triangle_aabb", 0); - /* "voxelize.pyx":58 + /* "voxelize.pyx":59 * @cython.wraparound(False) # Deactivate negative indexing. * cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): * assert(boxcenter.shape[0] == 3) # <<<<<<<<<<<<<< @@ -2773,12 +2773,12 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_boxcenter.shape[0]) == 3) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 58, __pyx_L1_error) + __PYX_ERR(0, 59, __pyx_L1_error) } } #endif - /* "voxelize.pyx":59 + /* "voxelize.pyx":60 * cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): * assert(boxcenter.shape[0] == 3) * assert(boxhalfsize.shape[0] == 3) # <<<<<<<<<<<<<< @@ -2789,12 +2789,12 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce if (unlikely(!Py_OptimizeFlag)) { if (unlikely(!(((__pyx_v_boxhalfsize.shape[0]) == 3) != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 59, __pyx_L1_error) + __PYX_ERR(0, 60, __pyx_L1_error) } } #endif - /* "voxelize.pyx":60 + /* "voxelize.pyx":61 * assert(boxcenter.shape[0] == 3) * assert(boxhalfsize.shape[0] == 3) * assert(triverts.shape[0] == triverts.shape[1] == 3) # <<<<<<<<<<<<<< @@ -2809,12 +2809,12 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce } if (unlikely(!(__pyx_t_1 != 0))) { PyErr_SetNone(PyExc_AssertionError); - __PYX_ERR(0, 60, __pyx_L1_error) + __PYX_ERR(0, 61, __pyx_L1_error) } } #endif - /* "voxelize.pyx":64 + /* "voxelize.pyx":65 * # print(triverts) * # Call functions * cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], # <<<<<<<<<<<<<< @@ -2824,7 +2824,7 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce __pyx_t_2 = 0; __pyx_t_3 = 0; - /* "voxelize.pyx":65 + /* "voxelize.pyx":66 * # Call functions * cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], * &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) # <<<<<<<<<<<<<< @@ -2837,7 +2837,7 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce __pyx_t_8 = 2; __pyx_t_9 = 0; - /* "voxelize.pyx":64 + /* "voxelize.pyx":65 * # print(triverts) * # Call functions * cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], # <<<<<<<<<<<<<< @@ -2846,7 +2846,7 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce */ __pyx_v_result = triBoxOverlap((&(*((float *) ( /* dim=0 */ ((char *) (((float *) __pyx_v_boxcenter.data) + __pyx_t_2)) )))), (&(*((float *) ( /* dim=0 */ ((char *) (((float *) __pyx_v_boxhalfsize.data) + __pyx_t_3)) )))), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_4 * __pyx_v_triverts.strides[0]) )) + __pyx_t_5)) )))), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_6 * __pyx_v_triverts.strides[0]) )) + __pyx_t_7)) )))), (&(*((float *) ( /* dim=1 */ ((char *) (((float *) ( /* dim=0 */ (__pyx_v_triverts.data + __pyx_t_8 * __pyx_v_triverts.strides[0]) )) + __pyx_t_9)) ))))); - /* "voxelize.pyx":66 + /* "voxelize.pyx":67 * cdef int result = triBoxOverlap(&boxcenter[0], &boxhalfsize[0], * &triverts[0, 0], &triverts[1, 0], &triverts[2, 0]) * return result # <<<<<<<<<<<<<< @@ -2854,7 +2854,7 @@ static int __pyx_f_8voxelize_test_triangle_aabb(__Pyx_memviewslice __pyx_v_boxce __pyx_r = __pyx_v_result; goto __pyx_L0; - /* "voxelize.pyx":57 + /* "voxelize.pyx":58 * @cython.boundscheck(False) # Deactivate bounds checking * @cython.wraparound(False) # Deactivate negative indexing. * cdef int test_triangle_aabb(float[::1] boxcenter, float[::1] boxhalfsize, float[:, ::1] triverts): # <<<<<<<<<<<<<< @@ -16757,7 +16757,7 @@ static __Pyx_StringTabEntry __pyx_string_tab[] = { {0, 0, 0, 0, 0, 0, 0} }; static CYTHON_SMALL_CODE int __Pyx_InitCachedBuiltins(void) { - __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 18, __pyx_L1_error) + __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 19, __pyx_L1_error) __pyx_builtin_ValueError = __Pyx_GetBuiltinName(__pyx_n_s_ValueError); if (!__pyx_builtin_ValueError) __PYX_ERR(1, 133, __pyx_L1_error) __pyx_builtin_MemoryError = __Pyx_GetBuiltinName(__pyx_n_s_MemoryError); if (!__pyx_builtin_MemoryError) __PYX_ERR(1, 148, __pyx_L1_error) __pyx_builtin_enumerate = __Pyx_GetBuiltinName(__pyx_n_s_enumerate); if (!__pyx_builtin_enumerate) __PYX_ERR(1, 151, __pyx_L1_error) @@ -17377,8 +17377,8 @@ if (!__Pyx_RefNanny) { /* "voxelize.pyx":1 * cimport cython # <<<<<<<<<<<<<< - * from libc.math cimport floor, ceil * from cython.view cimport array as cvarray + * from libc.math cimport ceil, floor */ __pyx_t_1 = __Pyx_PyDict_NewPresized(0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 1, __pyx_L1_error) __Pyx_GOTREF(__pyx_t_1); diff --git a/lib/common/libvoxelize/voxelize.pyx b/lib/common/libvoxelize/voxelize.pyx index 1ba840295c88f2ae11e2c575b3aece9dc07a930a..8e6a4829ba31b11bfab21ac1ae40bff27a74be65 100644 --- a/lib/common/libvoxelize/voxelize.pyx +++ b/lib/common/libvoxelize/voxelize.pyx @@ -1,6 +1,7 @@ cimport cython -from libc.math cimport floor, ceil from cython.view cimport array as cvarray +from libc.math cimport ceil, floor + cdef extern from "tribox2.h": int triBoxOverlap(float boxcenter[3], float boxhalfsize[3], diff --git a/lib/common/local_affine.py b/lib/common/local_affine.py index a8c7225ba1e0428161425669a19da3a94d0866df..ca23bd61e7c90de4a8ac19a4554c46417c5be87d 100644 --- a/lib/common/local_affine.py +++ b/lib/common/local_affine.py @@ -5,13 +5,14 @@ # file that should have been included as part of this package. import torch -import trimesh import torch.nn as nn -from tqdm import tqdm -from pytorch3d.structures import Meshes +import trimesh from pytorch3d.loss import chamfer_distance -from lib.dataset.mesh_util import update_mesh_shape_prior_losses +from pytorch3d.structures import Meshes +from tqdm import tqdm + from lib.common.train_util import init_loss +from lib.dataset.mesh_util import update_mesh_shape_prior_losses # reference: https://github.com/wuhaozhe/pytorch-nicp @@ -84,11 +85,9 @@ def register(target_mesh, src_mesh, device, verbose=True): src_mesh.verts_padded().shape[0], src_mesh.edges_packed() ).to(device) - optimizer_cloth = torch.optim.Adam( - [{ - 'params': local_affine_model.parameters() - }], lr=1e-2, amsgrad=True - ) + optimizer_cloth = torch.optim.Adam([{'params': local_affine_model.parameters()}], + lr=1e-2, + amsgrad=True) scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer_cloth, mode="min", @@ -104,7 +103,7 @@ def register(target_mesh, src_mesh, device, verbose=True): loop_cloth = tqdm(range(100)) else: loop_cloth = range(100) - + for i in loop_cloth: optimizer_cloth.zero_grad() diff --git a/lib/common/render.py b/lib/common/render.py index 769f6899a72a11bfe51cbd204e8e886b61aa4194..5f89eceeacba0d7ebac9602a800db48a0d406f4e 100644 --- a/lib/common/render.py +++ b/lib/common/render.py @@ -14,35 +14,36 @@ # # Contact: ps-license@tuebingen.mpg.de +import math +import os + +import cv2 +import numpy as np +import torch +from PIL import ImageColor from pytorch3d.renderer import ( + AlphaCompositor, BlendParams, - blending, - look_at_view_transform, FoVOrthographicCameras, - RasterizationSettings, + MeshRasterizer, + MeshRenderer, PointsRasterizationSettings, - PointsRenderer, - AlphaCompositor, PointsRasterizer, - MeshRenderer, - MeshRasterizer, + PointsRenderer, + RasterizationSettings, SoftSilhouetteShader, TexturesVertex, + blending, + look_at_view_transform, ) from pytorch3d.renderer.mesh import TexturesVertex from pytorch3d.structures import Meshes -from lib.dataset.mesh_util import get_visibility -from lib.common.imutils import blend_rgb_norm +from termcolor import colored +from tqdm import tqdm import lib.common.render_utils as util -import torch -import numpy as np -from PIL import ImageColor -from tqdm import tqdm -import os -import cv2 -import math -from termcolor import colored +from lib.common.imutils import blend_rgb_norm +from lib.dataset.mesh_util import get_visibility def image2vid(images, vid_path): @@ -58,7 +59,7 @@ def image2vid(images, vid_path): video.release() -def query_color(verts, faces, image, device): +def query_color(verts, faces, image, device, paint_normal=True): """query colors from points and image Args: @@ -77,16 +78,16 @@ def query_color(verts, faces, image, device): 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] + colors = (( + torch.nn.functional.grid_sample(image, uv, align_corners=True)[0, :, :, 0].permute(1, 0) + + 1.0 + ) * 0.5 * 255.0) + if paint_normal: + 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] + else: + colors[visibility == 0.0] = torch.tensor([0.0, 0.0, 0.0]).to(device) return colors.detach().cpu() @@ -121,31 +122,25 @@ class Render: self.step = 3 self.cam_pos = { - "frontback": - torch.tensor( - [ - (0, self.mesh_y_center, self.dis), - (0, self.mesh_y_center, -self.dis), - ] - ), - "four": - torch.tensor( - [ - (0, self.mesh_y_center, self.dis), - (self.dis, self.mesh_y_center, 0), - (0, self.mesh_y_center, -self.dis), - (-self.dis, self.mesh_y_center, 0), - ] - ), - "around": - torch.tensor( - [ - ( - 100.0 * math.cos(np.pi / 180 * angle), self.mesh_y_center, - 100.0 * math.sin(np.pi / 180 * angle) - ) for angle in range(0, 360, self.step) - ] - ) + "front": + torch.tensor([ + (0, self.mesh_y_center, self.dis), + (0, self.mesh_y_center, -self.dis), + ]), "frontback": + torch.tensor([ + (0, self.mesh_y_center, self.dis), + (0, self.mesh_y_center, -self.dis), + ]), "four": + torch.tensor([ + (0, self.mesh_y_center, self.dis), + (self.dis, self.mesh_y_center, 0), + (0, self.mesh_y_center, -self.dis), + (-self.dis, self.mesh_y_center, 0), + ]), "around": + torch.tensor([( + 100.0 * math.cos(np.pi / 180 * angle), self.mesh_y_center, + 100.0 * math.sin(np.pi / 180 * angle) + ) for angle in range(0, 360, self.step)]) } self.type = "color" @@ -315,7 +310,7 @@ class Render: save_path, fourcc, self.fps, - (width*3, int(height)), + (width * 3, int(height)), ) pbar = tqdm(range(len(self.meshes))) @@ -352,15 +347,13 @@ class Render: for cam_id in pbar: img_raw = data["img_raw"] num_obj = len(mesh_renders) // 2 - img_smpl = blend_rgb_norm( - (torch.stack(mesh_renders)[:num_obj, cam_id] - 0.5) * 2.0, data - ) - img_cloth = blend_rgb_norm( - (torch.stack(mesh_renders)[num_obj:, cam_id] - 0.5) * 2.0, data - ) - final_img = torch.cat( - [img_raw, img_smpl, img_cloth], dim=-1).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8) - + img_smpl = blend_rgb_norm((torch.stack(mesh_renders)[:num_obj, cam_id] - 0.5) * 2.0, + data) + img_cloth = blend_rgb_norm((torch.stack(mesh_renders)[num_obj:, cam_id] - 0.5) * 2.0, + data) + final_img = torch.cat([img_raw, img_smpl, img_cloth], + dim=-1).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8) + video.write(final_img[:, :, ::-1]) video.release() diff --git a/lib/common/render_utils.py b/lib/common/render_utils.py index cb2ca46f420c063c7a1c6a82276d41c42852e451..389eab740fc341b9395e1e008b4fa91e9ef3cf83 100644 --- a/lib/common/render_utils.py +++ b/lib/common/render_utils.py @@ -14,13 +14,15 @@ # # 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 + +import numpy as np +import torch +import trimesh from pytorch3d.renderer.mesh import rasterize_meshes +from pytorch3d.structures import Meshes +from torch import nn Tensor = NewType("Tensor", torch.Tensor) @@ -125,8 +127,6 @@ def batch_contains(verts, faces, points): def dict2obj(d): - # if isinstance(d, list): - # d = [dict2obj(x) for x in d] if not isinstance(d, dict): return d @@ -161,7 +161,9 @@ class Pytorch3dRasterizer(nn.Module): x,y,z are in image space, normalized can only render squared image now """ - def __init__(self, image_size=224, blur_radius=0.0, faces_per_pixel=1): + def __init__( + self, image_size=224, blur_radius=0.0, faces_per_pixel=1, device=torch.device("cuda:0") + ): """ use fixed raster_settings for rendering faces """ @@ -177,6 +179,7 @@ class Pytorch3dRasterizer(nn.Module): } raster_settings = dict2obj(raster_settings) self.raster_settings = raster_settings + self.device = device def forward(self, vertices, faces, attributes=None): fixed_vertices = vertices.clone() @@ -209,3 +212,15 @@ class Pytorch3dRasterizer(nn.Module): pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) pixel_vals = torch.cat([pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) return pixel_vals + + def get_texture(self, uvcoords, uvfaces, verts, faces, verts_color): + + batch_size = verts.shape[0] + uv_verts_color = face_vertices(verts_color, faces.expand(batch_size, -1, + -1)).to(self.device) + uv_map = self.forward( + uvcoords.expand(batch_size, -1, -1), uvfaces.expand(batch_size, -1, -1), uv_verts_color + )[:, :3] + uv_map_npy = np.flip(uv_map.squeeze(0).permute(1, 2, 0).cpu().numpy(), 0) + + return uv_map_npy diff --git a/lib/common/seg3d_lossless.py b/lib/common/seg3d_lossless.py index 4f5cba2a1edb3a5df14d17beabb9d296203865c1..bf086b505a98e0fb8a639394e982cddb048ab8b0 100644 --- a/lib/common/seg3d_lossless.py +++ b/lib/common/seg3d_lossless.py @@ -14,19 +14,16 @@ # # Contact: ps-license@tuebingen.mpg.de -from .seg3d_utils import ( - create_grid3D, - plot_mask3D, - SmoothConv3D, -) +import logging +import numpy as np import torch import torch.nn as nn -import numpy as np import torch.nn.functional as F -import logging from pytorch3d.ops.marching_cubes import marching_cubes +from .seg3d_utils import SmoothConv3D, create_grid3D, plot_mask3D + logging.getLogger("lightning").setLevel(logging.ERROR) @@ -378,10 +375,8 @@ class Seg3dLossless(nn.Module): with torch.no_grad(): # conflicts - conflicts = ( - (occupancys_interp - self.balance_value) * - (occupancys_topk - self.balance_value) < 0 - )[0, 0] + 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) @@ -407,12 +402,9 @@ class Seg3dLossless(nn.Module): title="conflicts", ) - conflicts_boundary = ( - ( - conflicts_coords.int() + - self.gird8_offsets.unsqueeze(1) * stride.int() - ).reshape(-1, 3).long().unique(dim=0) - ) + 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 @@ -466,10 +458,8 @@ class Seg3dLossless(nn.Module): with torch.no_grad(): # conflicts - conflicts = ( - (occupancys_interp - self.balance_value) * - (occupancys_topk - self.balance_value) < 0 - )[0, 0] + 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) diff --git a/lib/common/seg3d_utils.py b/lib/common/seg3d_utils.py index bee264615a54777bb948414c82f502c678664329..8f2c547d2e02b40c9e149071818116160df0be3a 100644 --- a/lib/common/seg3d_utils.py +++ b/lib/common/seg3d_utils.py @@ -14,10 +14,10 @@ # # Contact: ps-license@tuebingen.mpg.de +import matplotlib.pyplot as plt 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): @@ -140,9 +140,8 @@ class SmoothConv2D(nn.Module): 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) + 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): @@ -155,9 +154,8 @@ class SmoothConv3D(nn.Module): 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) + 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): @@ -185,9 +183,8 @@ def build_smooth_conv2D(in_channels=1, out_channels=1, kernel_size=3, padding=1) 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.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 diff --git a/lib/common/train_util.py b/lib/common/train_util.py index a39102a5849e25d056b9d96d2df9538790bec6ea..324547c05fd7281381262f16d6240d1b9f2240da 100644 --- a/lib/common/train_util.py +++ b/lib/common/train_util.py @@ -14,11 +14,12 @@ # # Contact: ps-license@tuebingen.mpg.de +import pytorch_lightning as pl import torch +from termcolor import colored + from ..dataset.mesh_util import * from ..net.geometry import orthogonal -from termcolor import colored -import pytorch_lightning as pl class Format: @@ -30,50 +31,23 @@ def init_loss(): losses = { # Cloth: chamfer distance - "cloth": { - "weight": 1e3, - "value": 0.0 - }, + "cloth": {"weight": 1e3, "value": 0.0}, # Stiffness: [RT]_v1 - [RT]_v2 (v1-edge-v2) - "stiff": { - "weight": 1e5, - "value": 0.0 - }, + "stiff": {"weight": 1e5, "value": 0.0}, # Cloth: det(R) = 1 - "rigid": { - "weight": 1e5, - "value": 0.0 - }, + "rigid": {"weight": 1e5, "value": 0.0}, # Cloth: edge length - "edge": { - "weight": 0, - "value": 0.0 - }, + "edge": {"weight": 0, "value": 0.0}, # Cloth: normal consistency - "nc": { - "weight": 0, - "value": 0.0 - }, + "nc": {"weight": 0, "value": 0.0}, # Cloth: laplacian smoonth - "lapla": { - "weight": 1e2, - "value": 0.0 - }, + "lapla": {"weight": 1e2, "value": 0.0}, # Body: Normal_pred - Normal_smpl - "normal": { - "weight": 1e0, - "value": 0.0 - }, + "normal": {"weight": 1e0, "value": 0.0}, # Body: Silhouette_pred - Silhouette_smpl - "silhouette": { - "weight": 1e0, - "value": 0.0 - }, + "silhouette": {"weight": 1e0, "value": 0.0}, # Joint: reprojected joints difference - "joint": { - "weight": 5e0, - "value": 0.0 - }, + "joint": {"weight": 5e0, "value": 0.0}, } return losses @@ -143,9 +117,9 @@ def query_func_IF(batch, netG, points): def batch_mean(res, key): - return torch.stack( - [x[key] if torch.is_tensor(x[key]) else torch.as_tensor(x[key]) for x in res] - ).mean() + return torch.stack([ + x[key] if torch.is_tensor(x[key]) else torch.as_tensor(x[key]) for x in res + ]).mean() def accumulate(outputs, rot_num, split): diff --git a/lib/common/voxelize.py b/lib/common/voxelize.py index f792189ccc185e9a7b596eae5a9230fe21482aef..44cbf2b4f43aeed217a56b543feddf6110336773 100644 --- a/lib/common/voxelize.py +++ b/lib/common/voxelize.py @@ -1,15 +1,14 @@ -import trimesh -import numpy as np import os import traceback -import torch import numpy as np +import torch import trimesh from scipy import ndimage from skimage.measure import block_reduce -from lib.common.libvoxelize.voxelize import voxelize_mesh_ + from lib.common.libmesh.inside_mesh import check_mesh_contains +from lib.common.libvoxelize.voxelize import voxelize_mesh_ # From Occupancy Networks, Mescheder et. al. CVPR'19 @@ -147,76 +146,63 @@ class VoxelGrid: f2_r_x, f2_r_y, f2_r_z = np.where(f2_r) f3_r_x, f3_r_y, f3_r_z = np.where(f3_r) - faces_1_l = np.stack( - [ - v_idx[f1_l_x, f1_l_y, f1_l_z], - v_idx[f1_l_x, f1_l_y, f1_l_z + 1], - v_idx[f1_l_x, f1_l_y + 1, f1_l_z + 1], - v_idx[f1_l_x, f1_l_y + 1, f1_l_z], - ], - axis=1 - ) - - faces_1_r = np.stack( - [ - v_idx[f1_r_x, f1_r_y, f1_r_z], - v_idx[f1_r_x, f1_r_y + 1, f1_r_z], - v_idx[f1_r_x, f1_r_y + 1, f1_r_z + 1], - v_idx[f1_r_x, f1_r_y, f1_r_z + 1], - ], - axis=1 - ) - - faces_2_l = np.stack( - [ - v_idx[f2_l_x, f2_l_y, f2_l_z], - v_idx[f2_l_x + 1, f2_l_y, f2_l_z], - v_idx[f2_l_x + 1, f2_l_y, f2_l_z + 1], - v_idx[f2_l_x, f2_l_y, f2_l_z + 1], - ], - axis=1 - ) - - faces_2_r = np.stack( - [ - v_idx[f2_r_x, f2_r_y, f2_r_z], - v_idx[f2_r_x, f2_r_y, f2_r_z + 1], - v_idx[f2_r_x + 1, f2_r_y, f2_r_z + 1], - v_idx[f2_r_x + 1, f2_r_y, f2_r_z], - ], - axis=1 - ) - - faces_3_l = np.stack( - [ - v_idx[f3_l_x, f3_l_y, f3_l_z], - v_idx[f3_l_x, f3_l_y + 1, f3_l_z], - v_idx[f3_l_x + 1, f3_l_y + 1, f3_l_z], - v_idx[f3_l_x + 1, f3_l_y, f3_l_z], - ], - axis=1 - ) - - faces_3_r = np.stack( - [ - v_idx[f3_r_x, f3_r_y, f3_r_z], - v_idx[f3_r_x + 1, f3_r_y, f3_r_z], - v_idx[f3_r_x + 1, f3_r_y + 1, f3_r_z], - v_idx[f3_r_x, f3_r_y + 1, f3_r_z], - ], - axis=1 - ) - - faces = np.concatenate( - [ - faces_1_l, - faces_1_r, - faces_2_l, - faces_2_r, - faces_3_l, - faces_3_r, - ], axis=0 - ) + faces_1_l = np.stack([ + v_idx[f1_l_x, f1_l_y, f1_l_z], + v_idx[f1_l_x, f1_l_y, f1_l_z + 1], + v_idx[f1_l_x, f1_l_y + 1, f1_l_z + 1], + v_idx[f1_l_x, f1_l_y + 1, f1_l_z], + ], + axis=1) + + faces_1_r = np.stack([ + v_idx[f1_r_x, f1_r_y, f1_r_z], + v_idx[f1_r_x, f1_r_y + 1, f1_r_z], + v_idx[f1_r_x, f1_r_y + 1, f1_r_z + 1], + v_idx[f1_r_x, f1_r_y, f1_r_z + 1], + ], + axis=1) + + faces_2_l = np.stack([ + v_idx[f2_l_x, f2_l_y, f2_l_z], + v_idx[f2_l_x + 1, f2_l_y, f2_l_z], + v_idx[f2_l_x + 1, f2_l_y, f2_l_z + 1], + v_idx[f2_l_x, f2_l_y, f2_l_z + 1], + ], + axis=1) + + faces_2_r = np.stack([ + v_idx[f2_r_x, f2_r_y, f2_r_z], + v_idx[f2_r_x, f2_r_y, f2_r_z + 1], + v_idx[f2_r_x + 1, f2_r_y, f2_r_z + 1], + v_idx[f2_r_x + 1, f2_r_y, f2_r_z], + ], + axis=1) + + faces_3_l = np.stack([ + v_idx[f3_l_x, f3_l_y, f3_l_z], + v_idx[f3_l_x, f3_l_y + 1, f3_l_z], + v_idx[f3_l_x + 1, f3_l_y + 1, f3_l_z], + v_idx[f3_l_x + 1, f3_l_y, f3_l_z], + ], + axis=1) + + faces_3_r = np.stack([ + v_idx[f3_r_x, f3_r_y, f3_r_z], + v_idx[f3_r_x + 1, f3_r_y, f3_r_z], + v_idx[f3_r_x + 1, f3_r_y + 1, f3_r_z], + v_idx[f3_r_x, f3_r_y + 1, f3_r_z], + ], + axis=1) + + faces = np.concatenate([ + faces_1_l, + faces_1_r, + faces_2_l, + faces_2_r, + faces_3_l, + faces_3_r, + ], + axis=0) vertices = self.loc + self.scale * vertices mesh = trimesh.Trimesh(vertices, faces, process=False) diff --git a/lib/dataset/EvalDataset.py b/lib/dataset/EvalDataset.py index d7aca595c2f3027ee3ba3fdfa9a99ea742afe2b7..9ef6426285b103b04db4f0c827c6755056bcc01b 100644 --- a/lib/dataset/EvalDataset.py +++ b/lib/dataset/EvalDataset.py @@ -14,22 +14,24 @@ # # Contact: ps-license@tuebingen.mpg.de -import torch.nn.functional as F -from lib.common.render import Render -from lib.dataset.mesh_util import (SMPLX, projection, rescale_smpl, HoppeMesh) -import os.path as osp -import numpy as np -from PIL import Image import os +import os.path as osp + import cv2 -import trimesh +import numpy as np import torch +import torch.nn.functional as F import torchvision.transforms as transforms +import trimesh +from PIL import Image + +from lib.common.render import Render +from lib.dataset.mesh_util import SMPLX, HoppeMesh, projection, rescale_smpl cape_gender = { "male": - ['00032', '00096', '00122', '00127', '00145', '00215', '02474', '03284', '03375', '03394'], - "female": ['00134', '00159', '03223', '03331', '03383'] + ['00032', '00096', '00122', '00127', '00145', '00215', '02474', '03284', '03375', + '03394'], "female": ['00134', '00159', '03223', '03331', '03383'] } @@ -74,30 +76,27 @@ class EvalDataset: "scale": self.scales[dataset_id], } - self.datasets_dict[dataset].update( - {"subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str)} - ) + self.datasets_dict[dataset].update({ + "subjects": + np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str) + }) self.subject_list = self.get_subject_list() 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)), - ] - ) + 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.mask_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.0, ), (1.0, )), + ]) self.device = device self.render = Render(size=512, device=self.device) @@ -154,27 +153,23 @@ class EvalDataset: } if dataset == "cape": - data_dict.update( - { - "mesh_path": - osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj"), - "smpl_path": - osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.obj"), - } - ) + data_dict.update({ + "mesh_path": + osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj"), + "smpl_path": + osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.obj"), + }) else: - data_dict.update( - { - "mesh_path": - osp.join( - self.datasets_dict[dataset]["mesh_dir"], - f"{subject}.obj", - ), - "smplx_path": - osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}.obj"), - } - ) + data_dict.update({ + "mesh_path": + osp.join( + self.datasets_dict[dataset]["mesh_dir"], + f"{subject}.obj", + ), + "smplx_path": + osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}.obj"), + }) # load training data data_dict.update(self.load_calib(data_dict)) @@ -183,18 +178,17 @@ class EvalDataset: 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") - } - ) + data_dict.update({ + f"{name}_path": + osp.join(self.root, render_folder, name, f"{rotation:03d}.png") + }) # tensor update if os.path.exists(data_dict[f"{name}_path"]): - data_dict.update( - {name: self.imagepath2tensor(data_dict[f"{name}_path"], channel, inv=False)} - ) + 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.load_smpl(data_dict)) diff --git a/lib/dataset/Evaluator.py b/lib/dataset/Evaluator.py index 066e183617fbe79532d7a7f291227d8dfaef9c7c..3d840033735f04447fc296888e2775a5f8686b5e 100644 --- a/lib/dataset/Evaluator.py +++ b/lib/dataset/Evaluator.py @@ -14,20 +14,21 @@ # # Contact: ps-license@tuebingen.mpg.de -from lib.dataset.mesh_util import projection -from lib.common.render import Render +from typing import Tuple + import numpy as np import torch -from torchvision.utils import make_grid +from PIL import Image from pytorch3d import _C +from pytorch3d.ops.mesh_face_areas_normals import mesh_face_areas_normals +from pytorch3d.ops.packed_to_padded import packed_to_padded +from pytorch3d.structures import Pointclouds from torch.autograd import Function from torch.autograd.function import once_differentiable -from pytorch3d.structures import Pointclouds -from PIL import Image +from torchvision.utils import make_grid -from typing import Tuple -from pytorch3d.ops.mesh_face_areas_normals import mesh_face_areas_normals -from pytorch3d.ops.packed_to_padded import packed_to_padded +from lib.common.render import Render +from lib.dataset.mesh_util import projection _DEFAULT_MIN_TRIANGLE_AREA: float = 5e-3 @@ -278,12 +279,10 @@ class Evaluator: # error_hf = ((((src_normal_arr - tgt_normal_arr) * sim_mask)**2).sum(dim=0).mean()) * 4.0 - normal_img = Image.fromarray( - ( - torch.cat([src_normal_arr, tgt_normal_arr], - dim=1).permute(1, 2, 0).detach().cpu().numpy() * 255.0 - ).astype(np.uint8) - ) + normal_img = Image.fromarray(( + torch.cat([src_normal_arr, tgt_normal_arr], + dim=1).permute(1, 2, 0).detach().cpu().numpy() * 255.0 + ).astype(np.uint8)) normal_img.save(normal_path) return error diff --git a/lib/dataset/NormalDataset.py b/lib/dataset/NormalDataset.py index 3567ac8cd5a83517a93c80c008bbb9b8d23616a7..59132030d61c27f4cf6883025a440613d3eecbce 100644 --- a/lib/dataset/NormalDataset.py +++ b/lib/dataset/NormalDataset.py @@ -14,12 +14,13 @@ # # Contact: ps-license@tuebingen.mpg.de -import kornia import os.path as osp + +import kornia import numpy as np +import torchvision.transforms as transforms from PIL import Image from termcolor import colored -import torchvision.transforms as transforms class NormalDataset: @@ -59,22 +60,18 @@ class NormalDataset: 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)), - ] - ) + 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.mask_to_tensor = transforms.Compose([ + transforms.Resize(self.input_size), + transforms.ToTensor(), + transforms.Normalize((0.0, ), (1.0, )), + ]) def get_subject_list(self, split): @@ -128,21 +125,15 @@ class NormalDataset: 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") - } - ) - - data_dict.update( - { - name: - self.imagepath2tensor( - data_dict[f"{name}_path"], channel, inv=False, erasing=False - ) - } - ) + data_dict.update({ + f"{name}_path": + osp.join(self.root, render_folder, name, f"{rotation:03d}.png") + }) + + data_dict.update({ + name: + self.imagepath2tensor(data_dict[f"{name}_path"], channel, inv=False, erasing=False) + }) path_keys = [key for key in data_dict.keys() if "_path" in key or "_dir" in key] diff --git a/lib/dataset/NormalModule.py b/lib/dataset/NormalModule.py index ff672b3c42f5951f4ebf6c8446014d1d277ab02c..16dd02ec26789d40715b24b67f371da45aff2f8f 100644 --- a/lib/dataset/NormalModule.py +++ b/lib/dataset/NormalModule.py @@ -14,11 +14,11 @@ # # Contact: ps-license@tuebingen.mpg.de -from torch.utils.data import DataLoader -from lib.dataset.NormalDataset import NormalDataset - # pytorch lightning related libs import pytorch_lightning as pl +from torch.utils.data import DataLoader + +from lib.dataset.NormalDataset import NormalDataset class NormalModule(pl.LightningDataModule): diff --git a/lib/dataset/PointFeat.py b/lib/dataset/PointFeat.py index 457b949e5ce712a1eace33b1306fd48613ba8887..f6b5a708cca885f99205662fe6b42e13bf432da7 100644 --- a/lib/dataset/PointFeat.py +++ b/lib/dataset/PointFeat.py @@ -1,5 +1,6 @@ -from pytorch3d.structures import Meshes, Pointclouds import torch +from pytorch3d.structures import Meshes, Pointclouds + from lib.common.render_utils import face_vertices from lib.dataset.Evaluator import point_mesh_distance from lib.dataset.mesh_util import SMPLX, barycentric_coordinates_of_projection diff --git a/lib/dataset/TestDataset.py b/lib/dataset/TestDataset.py index e99627ec8eeb6a369e10caf4d19c198910fc4a2e..b6eef20dd4a33654fa10b2f68bc5eb82715fa7db 100644 --- a/lib/dataset/TestDataset.py +++ b/lib/dataset/TestDataset.py @@ -14,37 +14,34 @@ # # Contact: ps-license@tuebingen.mpg.de -import warnings import logging +import warnings warnings.filterwarnings("ignore") logging.getLogger("lightning").setLevel(logging.ERROR) logging.getLogger("trimesh").setLevel(logging.ERROR) -from lib.pixielib.utils.config import cfg as pixie_cfg -from lib.pixielib.pixie import PIXIE -from lib.pixielib.models.SMPLX import SMPLX as PIXIE_SMPLX -from lib.common.imutils import process_image -from lib.common.train_util import Format -from lib.net.geometry import rotation_matrix_to_angle_axis, rot6d_to_rotmat - -from lib.pymafx.core import path_config -from lib.pymafx.models import pymaf_net +import glob +import os.path as osp -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 numpy as np +import torch import torch.nn.functional as F +from PIL import ImageFile +from termcolor import colored from torchvision import transforms from torchvision.models import detection -import os.path as osp -import torch -import glob -import numpy as np -from termcolor import colored -from PIL import ImageFile +from lib.common.config import cfg +from lib.common.imutils import process_image +from lib.common.render import Render +from lib.common.train_util import Format +from lib.dataset.mesh_util import SMPLX, get_visibility +from lib.pixielib.models.SMPLX import SMPLX as PIXIE_SMPLX +from lib.pixielib.pixie import PIXIE +from lib.pixielib.utils.config import cfg as pixie_cfg +from lib.pymafx.core import path_config +from lib.pymafx.models import pymaf_net ImageFile.LOAD_TRUNCATED_IMAGES = True @@ -66,9 +63,8 @@ class TestDataset: keep_lst = sorted(glob.glob(f"{self.image_dir}/*")) img_fmts = ["jpg", "png", "jpeg", "JPG", "bmp", "exr"] - self.subject_list = sorted( - [item for item in keep_lst if item.split(".")[-1] in img_fmts], reverse=False - ) + self.subject_list = sorted([item for item in keep_lst if item.split(".")[-1] in img_fmts], + reverse=False) # smpl related self.smpl_data = SMPLX() diff --git a/lib/dataset/body_model.py b/lib/dataset/body_model.py index cebb105591cab29d833f2965ec609c85fd522881..dff2e07ef1277930974dc066da1a2685bcd0cdf3 100644 --- a/lib/dataset/body_model.py +++ b/lib/dataset/body_model.py @@ -14,10 +14,11 @@ # # Contact: ps-license@tuebingen.mpg.de -import numpy as np +import os import pickle + +import numpy as np import torch -import os class SMPLModel: @@ -126,12 +127,10 @@ class SMPLModel: 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])), - ] - ) + np.hstack([ + self.R[i], + ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1])), + ]) ) ) # remove the transformation due to the rest pose @@ -163,19 +162,17 @@ class SMPLModel: 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]) + 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 @@ -357,12 +354,10 @@ class TetraSMPLModel: 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])), - ] - ) + np.hstack([ + self.R[i], + ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1])), + ]) ) ) # remove the transformation due to the rest pose @@ -398,19 +393,17 @@ class TetraSMPLModel: 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]) + 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 diff --git a/lib/dataset/mesh_util.py b/lib/dataset/mesh_util.py index 7dc46d51c3abd9a4e22c7bbba96c78ef86e03dbf..6f257765db55681016424c446db86e1f157e3866 100644 --- a/lib/dataset/mesh_util.py +++ b/lib/dataset/mesh_util.py @@ -14,25 +14,25 @@ # # Contact: ps-license@tuebingen.mpg.de +import json import os +import os.path as osp + +import _pickle as cPickle import numpy as np +import open3d as o3d import torch +import torch.nn.functional as F import torchvision import trimesh -import json -import open3d as o3d -import os.path as osp -import _pickle as cPickle -from termcolor import colored +from PIL import Image, ImageDraw, ImageFont +from pytorch3d.loss import mesh_laplacian_smoothing, mesh_normal_consistency +from pytorch3d.renderer.mesh import rasterize_meshes +from pytorch3d.structures import Meshes from scipy.spatial import cKDTree -from pytorch3d.structures import Meshes -import torch.nn.functional as F import lib.smplx as smplx -from lib.common.render_utils import Pytorch3dRasterizer -from pytorch3d.renderer.mesh import rasterize_meshes -from PIL import Image, ImageFont, ImageDraw -from pytorch3d.loss import mesh_laplacian_smoothing, mesh_normal_consistency +from lib.common.render_utils import Pytorch3dRasterizer, face_vertices class Format: @@ -74,19 +74,17 @@ class SMPLX: self.smplx_vertex_lmkid = np.load(self.smplx_vertex_lmkid_path) self.smpl_vert_seg = json.load(open(self.smpl_vert_seg_path)) - self.smpl_mano_vid = np.concatenate( - [ - self.smpl_vert_seg["rightHand"], self.smpl_vert_seg["rightHandIndex1"], - self.smpl_vert_seg["leftHand"], self.smpl_vert_seg["leftHandIndex1"] - ] - ) + self.smpl_mano_vid = np.concatenate([ + self.smpl_vert_seg["rightHand"], self.smpl_vert_seg["rightHandIndex1"], + self.smpl_vert_seg["leftHand"], self.smpl_vert_seg["leftHandIndex1"] + ]) self.smplx_eyeball_fid_mask = np.load(self.smplx_eyeball_fid_path) self.smplx_mouth_fid = np.load(self.smplx_fill_mouth_fid_path) self.smplx_mano_vid_dict = np.load(self.smplx_mano_vid_path, allow_pickle=True) - self.smplx_mano_vid = np.concatenate( - [self.smplx_mano_vid_dict["left_hand"], self.smplx_mano_vid_dict["right_hand"]] - ) + self.smplx_mano_vid = np.concatenate([ + self.smplx_mano_vid_dict["left_hand"], self.smplx_mano_vid_dict["right_hand"] + ]) self.smplx_flame_vid = np.load(self.smplx_flame_vid_path, allow_pickle=True) self.smplx_front_flame_vid = self.smplx_flame_vid[np.load(self.front_flame_path)] @@ -110,26 +108,22 @@ class SMPLX: self.model_dir = osp.join(self.current_dir, "models") - self.ghum_smpl_pairs = torch.tensor( - [ - (0, 24), (2, 26), (5, 25), (7, 28), (8, 27), (11, 16), (12, 17), (13, 18), (14, 19), - (15, 20), (16, 21), (17, 39), (18, 44), (19, 36), (20, 41), (21, 35), (22, 40), - (23, 1), (24, 2), (25, 4), (26, 5), (27, 7), (28, 8), (29, 31), (30, 34), (31, 29), - (32, 32) - ] - ).long() + self.ghum_smpl_pairs = torch.tensor([(0, 24), (2, 26), (5, 25), (7, 28), (8, 27), (11, 16), + (12, 17), (13, 18), (14, 19), (15, 20), (16, 21), + (17, 39), (18, 44), (19, 36), (20, 41), (21, 35), + (22, 40), (23, 1), (24, 2), (25, 4), (26, 5), (27, 7), + (28, 8), (29, 31), (30, 34), (31, 29), + (32, 32)]).long() # smpl-smplx correspondence self.smpl_joint_ids_24 = np.arange(22).tolist() + [68, 73] self.smpl_joint_ids_24_pixie = np.arange(22).tolist() + [61 + 68, 72 + 68] self.smpl_joint_ids_45 = np.arange(22).tolist() + [68, 73] + np.arange(55, 76).tolist() - self.extra_joint_ids = np.array( - [ - 61, 72, 66, 69, 58, 68, 57, 56, 64, 59, 67, 75, 70, 65, 60, 61, 63, 62, 76, 71, 72, - 74, 73 - ] - ) + self.extra_joint_ids = np.array([ + 61, 72, 66, 69, 58, 68, 57, 56, 64, 59, 67, 75, 70, 65, 60, 61, 63, 62, 76, 71, 72, 74, + 73 + ]) self.extra_joint_ids += 68 @@ -369,9 +363,9 @@ def mesh_edge_loss(meshes, target_length: float = 0.0): return loss_all -def remesh_laplacian(mesh, obj_path): +def remesh_laplacian(mesh, obj_path, face_count=50000): - mesh = mesh.simplify_quadratic_decimation(50000) + mesh = mesh.simplify_quadratic_decimation(face_count) mesh = trimesh.smoothing.filter_humphrey( mesh, alpha=0.1, beta=0.5, iterations=10, laplacian_operator=None ) @@ -380,7 +374,7 @@ def remesh_laplacian(mesh, obj_path): return mesh -def poisson(mesh, obj_path, depth=10, decimation=True): +def poisson(mesh, obj_path, depth=10, face_count=50000): pcd_path = obj_path[:-4] + "_soups.ply" assert (mesh.vertex_normals.shape[1] == 3) @@ -395,12 +389,9 @@ def poisson(mesh, obj_path, depth=10, decimation=True): largest_mesh = keep_largest(trimesh.Trimesh(np.array(mesh.vertices), np.array(mesh.triangles))) largest_mesh.export(obj_path) - if decimation: - # mesh decimation for faster rendering - low_res_mesh = largest_mesh.simplify_quadratic_decimation(50000) - return low_res_mesh - else: - return largest_mesh + # mesh decimation for faster rendering + low_res_mesh = largest_mesh.simplify_quadratic_decimation(face_count) + return low_res_mesh # Losses to smooth / regularize the mesh shape @@ -437,10 +428,9 @@ def read_smpl_constants(folder): smpl_tetras = (np.loadtxt(os.path.join(folder, "tetrahedrons.txt"), dtype=np.int32) - 1) return_dict = { - "smpl_vertex_code": torch.tensor(smpl_vertex_code), - "smpl_face_code": torch.tensor(smpl_face_code), - "smpl_faces": torch.tensor(smpl_faces), - "smpl_tetras": torch.tensor(smpl_tetras) + "smpl_vertex_code": torch.tensor(smpl_vertex_code), "smpl_face_code": + torch.tensor(smpl_face_code), "smpl_faces": torch.tensor(smpl_faces), "smpl_tetras": + torch.tensor(smpl_tetras) } return return_dict @@ -598,22 +588,6 @@ def compute_normal(vertices, faces): return vert_norms, face_norms -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()] - - def compute_normal_batch(vertices, faces): if faces.shape[0] != vertices.shape[0]: @@ -657,20 +631,18 @@ def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type="smpl"): 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", - ] - ): + 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"] - ): + 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( @@ -751,3 +723,61 @@ def get_joint_mesh(joints, radius=2.0): else: combined = sum([combined, ball_new]) return combined + + +def preprocess_point_cloud(pcd, voxel_size): + pcd_down = pcd + pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature( + pcd_down, o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5.0, max_nn=100) + ) + return (pcd_down, pcd_fpfh) + + +def o3d_ransac(src, dst): + + voxel_size = 0.01 + distance_threshold = 1.5 * voxel_size + + o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Error) + + # print('Downsampling inputs') + src_down, src_fpfh = preprocess_point_cloud(src, voxel_size) + dst_down, dst_fpfh = preprocess_point_cloud(dst, voxel_size) + + # print('Running RANSAC') + result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( + src_down, + dst_down, + src_fpfh, + dst_fpfh, + mutual_filter=False, + max_correspondence_distance=distance_threshold, + estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False), + ransac_n=3, + checkers=[ + o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), + o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold) + ], + criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(1000000, 0.999) + ) + + return result.transformation + + +def export_obj(v_np, f_np, vt, ft, path): + + # write mtl info into obj + new_line = f"mtllib material.mtl \n" + vt_lines = "\nusemtl mat0 \n" + v_lines = "" + f_lines = "" + + for _v in v_np: + v_lines += f"v {_v[0]} {_v[1]} {_v[2]}\n" + for fid, _f in enumerate(f_np): + f_lines += f"f {_f[0]+1}/{ft[fid][0]+1} {_f[1]+1}/{ft[fid][1]+1} {_f[2]+1}/{ft[fid][2]+1}\n" + for _vt in vt: + vt_lines += f"vt {_vt[0]} {_vt[1]}\n" + new_file_data = new_line + v_lines + vt_lines + f_lines + with open(path, 'w') as file: + file.write(new_file_data) diff --git a/lib/net/BasePIFuNet.py b/lib/net/BasePIFuNet.py index eb18dbb3245d57c9e030c18094322a58e874db93..3b5a77aec804f0fcbf72f6320b3ef30cdeb83ea1 100644 --- a/lib/net/BasePIFuNet.py +++ b/lib/net/BasePIFuNet.py @@ -14,8 +14,8 @@ # # Contact: ps-license@tuebingen.mpg.de -import torch.nn as nn import pytorch_lightning as pl +import torch.nn as nn from .geometry import index, orthogonal, perspective diff --git a/lib/net/Discriminator.py b/lib/net/Discriminator.py index c60acdde000d414c78af0705ba268af3117c6ec9..b47ef9fd05ef645950be61111d417638a57ae3c6 100644 --- a/lib/net/Discriminator.py +++ b/lib/net/Discriminator.py @@ -1,11 +1,16 @@ """ The code is based on https://github.com/apple/ml-gsn/ with adaption. """ import math + import torch import torch.nn as nn -import math import torch.nn.functional as F -from lib.torch_utils.ops.native_ops import FusedLeakyReLU, fused_leaky_relu, upfirdn2d + +from lib.torch_utils.ops.native_ops import ( + FusedLeakyReLU, + fused_leaky_relu, + upfirdn2d, +) class DiscriminatorHead(nn.Module): diff --git a/lib/net/FBNet.py b/lib/net/FBNet.py index f4797667d4d800019967d7ee2ed944ec8b8550fc..5e04d5a04551c186379847637cc0a8d4b813b3da 100644 --- a/lib/net/FBNet.py +++ b/lib/net/FBNet.py @@ -19,13 +19,14 @@ 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 -from torchvision import models +import torch +import torch.nn as nn import torch.nn.functional as F +from torchvision import models ############################################################################### @@ -313,34 +314,28 @@ class NLayerDiscriminator(nn.Module): kw = 4 padw = int(np.ceil((kw - 1.0) / 2)) - sequence = [ - [ - nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), - nn.LeakyReLU(0.2, True) - ] - ] + sequence = [[ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, True) + ]] nf = ndf for n in range(1, n_layers): nf_prev = nf nf = min(nf * 2, 512) - sequence += [ - [ - nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw), - norm_layer(nf), - nn.LeakyReLU(0.2, True) - ] - ] + sequence += [[ + nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw), + norm_layer(nf), + nn.LeakyReLU(0.2, True) + ]] nf_prev = nf nf = min(nf * 2, 512) - sequence += [ - [ - nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw), - norm_layer(nf), - nn.LeakyReLU(0.2, True) - ] - ] + sequence += [[ + nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw), + norm_layer(nf), + nn.LeakyReLU(0.2, True) + ]] sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]] @@ -632,18 +627,16 @@ class GANLoss(pl.LightningModule): def get_target_tensor(self, input, target_is_real): target_tensor = None if target_is_real: - create_label = ( - (self.real_label_var is None) or (self.real_label_var.numel() != input.numel()) - ) + create_label = ((self.real_label_var is None) or + (self.real_label_var.numel() != input.numel())) if create_label: real_tensor = self.tensor(input.size()).fill_(self.real_label) self.real_label_var = real_tensor self.real_label_var.requires_grad = False target_tensor = self.real_label_var else: - create_label = ( - (self.fake_label_var is None) or (self.fake_label_var.numel() != input.numel()) - ) + create_label = ((self.fake_label_var is None) or + (self.fake_label_var.numel() != input.numel())) if create_label: fake_tensor = self.tensor(input.size()).fill_(self.fake_label) self.fake_label_var = fake_tensor diff --git a/lib/net/GANLoss.py b/lib/net/GANLoss.py index 5d6711479980e89a3fc067b5ef579bb382eb29df..79c778a403eb249687ac4850cfb2fb84e5e1dcfb 100644 --- a/lib/net/GANLoss.py +++ b/lib/net/GANLoss.py @@ -2,8 +2,9 @@ import torch import torch.nn as nn -from torch import autograd import torch.nn.functional as F +from torch import autograd + from lib.net.Discriminator import StyleDiscriminator diff --git a/lib/net/IFGeoNet.py b/lib/net/IFGeoNet.py index e7fd92a9fe766642100beb7083203401adffa20e..06d3deb4b406f10d21fb91767d824e7569fe5928 100644 --- a/lib/net/IFGeoNet.py +++ b/lib/net/IFGeoNet.py @@ -1,7 +1,9 @@ from pickle import TRUE + import torch import torch.nn as nn import torch.nn.functional as F + from lib.net.geometry import orthogonal @@ -151,13 +153,11 @@ class IFGeoNet(nn.Module): # here every channel corresponse to one feature. - features = torch.cat( - ( - feature_0_partial, feature_1_fused, feature_2, feature_3, feature_4, feature_5, - feature_6 - ), - dim=1 - ) # (B, features, 1,7,sample_num) + features = torch.cat(( + feature_0_partial, feature_1_fused, feature_2, feature_3, feature_4, feature_5, + feature_6 + ), + dim=1) # (B, features, 1,7,sample_num) shape = features.shape features = torch.reshape( features, (shape[0], shape[1] * shape[3], shape[4]) diff --git a/lib/net/IFGeoNet_nobody.py b/lib/net/IFGeoNet_nobody.py index 56de86268dcfbbf4a0226c206dcbb992f906db98..1daedb9ec1bccdf1a76c1a8938a902e11dbed9dc 100644 --- a/lib/net/IFGeoNet_nobody.py +++ b/lib/net/IFGeoNet_nobody.py @@ -1,7 +1,9 @@ from pickle import TRUE + import torch import torch.nn as nn import torch.nn.functional as F + from lib.net.geometry import orthogonal @@ -136,13 +138,11 @@ class IFGeoNet(nn.Module): # here every channel corresponse to one feature. - features = torch.cat( - ( - feature_0_partial, feature_1_fused, feature_2, feature_3, feature_4, feature_5, - feature_6 - ), - dim=1 - ) # (B, features, 1,7,sample_num) + features = torch.cat(( + feature_0_partial, feature_1_fused, feature_2, feature_3, feature_4, feature_5, + feature_6 + ), + dim=1) # (B, features, 1,7,sample_num) shape = features.shape features = torch.reshape( features, (shape[0], shape[1] * shape[3], shape[4]) diff --git a/lib/net/NormalNet.py b/lib/net/NormalNet.py index a065840ed859137e72ba1e37a40c636da7c32e6f..d4f0fe8e9a9ab935ffc8a49ae2d22d23bc44d4cc 100644 --- a/lib/net/NormalNet.py +++ b/lib/net/NormalNet.py @@ -14,14 +14,14 @@ # # Contact: ps-license@tuebingen.mpg.de -from lib.net.FBNet import define_G, define_D, VGGLoss, GANLoss, IDMRFLoss -from lib.net.net_util import init_net -from lib.net.BasePIFuNet import BasePIFuNet - import torch import torch.nn as nn import torch.nn.functional as F +from lib.net.BasePIFuNet import BasePIFuNet +from lib.net.FBNet import GANLoss, IDMRFLoss, VGGLoss, define_D, define_G +from lib.net.net_util import init_net + class NormalNet(BasePIFuNet): """ @@ -63,12 +63,12 @@ class NormalNet(BasePIFuNet): 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.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") diff --git a/lib/net/geometry.py b/lib/net/geometry.py index 6d7d82d2cb6b760596d1bbf70804e542999f802e..4c98530b316a3164ad82d742ac55da9e2fc8c212 100644 --- a/lib/net/geometry.py +++ b/lib/net/geometry.py @@ -14,11 +14,13 @@ # # Contact: ps-license@tuebingen.mpg.de -import torch -import numpy as np import numbers -from torch.nn import functional as F + +import numpy as np +import torch from einops.einops import rearrange +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 @@ -42,13 +44,11 @@ def quaternion_to_rotation_matrix(quat): 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) + 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 @@ -508,12 +508,10 @@ def estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000, img_si 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 + 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 @@ -580,13 +578,11 @@ def Rot_y(angle, category="torch", prepend_dim=True, device=None): prepend_dim: prepend an extra dimension Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) """ - m = np.array( - [ - [np.cos(angle), 0.0, np.sin(angle)], - [0.0, 1.0, 0.0], - [-np.sin(angle), 0.0, np.cos(angle)], - ] - ) + m = np.array([ + [np.cos(angle), 0.0, np.sin(angle)], + [0.0, 1.0, 0.0], + [-np.sin(angle), 0.0, np.cos(angle)], + ]) if category == "torch": if prepend_dim: return torch.tensor(m, dtype=torch.float, device=device).unsqueeze(0) @@ -608,13 +604,11 @@ def Rot_x(angle, category="torch", prepend_dim=True, device=None): prepend_dim: prepend an extra dimension Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) """ - m = np.array( - [ - [1.0, 0.0, 0.0], - [0.0, np.cos(angle), -np.sin(angle)], - [0.0, np.sin(angle), np.cos(angle)], - ] - ) + m = np.array([ + [1.0, 0.0, 0.0], + [0.0, np.cos(angle), -np.sin(angle)], + [0.0, np.sin(angle), np.cos(angle)], + ]) if category == "torch": if prepend_dim: return torch.tensor(m, dtype=torch.float, device=device).unsqueeze(0) @@ -636,13 +630,11 @@ def Rot_z(angle, category="torch", prepend_dim=True, device=None): 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.0], - [np.sin(angle), np.cos(angle), 0.0], - [0.0, 0.0, 1.0], - ] - ) + m = np.array([ + [np.cos(angle), -np.sin(angle), 0.0], + [np.sin(angle), np.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ]) if category == "torch": if prepend_dim: return torch.tensor(m, dtype=torch.float, device=device).unsqueeze(0) diff --git a/lib/net/net_util.py b/lib/net/net_util.py index d89fcff5670909cd41c2e917e87b3bdb25870d8a..fa3c9491596688de0425b4471318ff5c23a9a909 100644 --- a/lib/net/net_util.py +++ b/lib/net/net_util.py @@ -14,12 +14,13 @@ # # Contact: ps-license@tuebingen.mpg.de +import functools + 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 +from torch.nn import init def gradient(inputs, outputs): diff --git a/lib/net/voxelize.py b/lib/net/voxelize.py index 394b40e6eeeb158bb691c1e518b6b1f7a889b8d8..8525ef6cf389c40d8e4a82e1a442c24915a7acd8 100644 --- a/lib/net/voxelize.py +++ b/lib/net/voxelize.py @@ -1,11 +1,11 @@ from __future__ import division, print_function + +import numpy as np 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 +from torch.autograd import Function class VoxelizationFunction(Function): diff --git a/lib/pixielib/models/FLAME.py b/lib/pixielib/models/FLAME.py index b62b1069b6083685e8ff1511e57c48ccf79bc927..fb9ca09c4890f17206905546f4373ab186a5e6d1 100755 --- a/lib/pixielib/models/FLAME.py +++ b/lib/pixielib/models/FLAME.py @@ -13,10 +13,11 @@ # For comments or questions, please email us at pixie@tue.mpg.de # For commercial licensing contact, please contact ps-license@tuebingen.mpg.de +import pickle + +import numpy as np import torch import torch.nn as nn -import numpy as np -import pickle import torch.nn.functional as F diff --git a/lib/pixielib/models/SMPLX.py b/lib/pixielib/models/SMPLX.py index beb672facebc7fa9e61eee7f8e7f3f185ac6cdad..9f07f5740100133c94ba9e5f2f9767ba7ea4b42c 100644 --- a/lib/pixielib/models/SMPLX.py +++ b/lib/pixielib/models/SMPLX.py @@ -3,19 +3,20 @@ original from https://github.com/vchoutas/smplx modified by Vassilis and Yao """ +import pickle + +import numpy as np import torch import torch.nn as nn -import numpy as np -import pickle from .lbs import ( + JointsFromVerticesSelector, Struct, - to_tensor, - to_np, + find_dynamic_lmk_idx_and_bcoords, lbs, + to_np, + to_tensor, vertices2landmarks, - JointsFromVerticesSelector, - find_dynamic_lmk_idx_and_bcoords, ) # SMPLX @@ -209,468 +210,452 @@ extra_names = [ SMPLX_names += extra_names part_indices = {} -part_indices["body"] = np.array( - [ - 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, - 123, - 124, - 125, - 126, - 127, - 132, - 134, - 135, - 136, - 137, - 138, - 143, - ] -) -part_indices["torso"] = np.array( - [ - 0, - 1, - 2, - 3, - 6, - 9, - 12, - 13, - 14, - 15, - 16, - 17, - 18, - 19, - 22, - 23, - 24, - 55, - 56, - 57, - 58, - 59, - 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, - ] -) -part_indices["head"] = np.array( - [ - 12, - 15, - 22, - 23, - 24, - 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, - 125, - 126, - 134, - 136, - 137, - ] -) -part_indices["face"] = np.array( - [ - 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, - ] -) -part_indices["upper"] = np.array( - [ - 12, - 13, - 14, - 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, - ] -) -part_indices["hand"] = np.array( - [ - 20, - 21, - 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, - 128, - 129, - 130, - 131, - 133, - 139, - 140, - 141, - 142, - 144, - ] -) -part_indices["left_hand"] = np.array( - [ - 20, - 25, - 26, - 27, - 28, - 29, - 30, - 31, - 32, - 33, - 34, - 35, - 36, - 37, - 38, - 39, - 128, - 129, - 130, - 131, - 133, - ] -) -part_indices["right_hand"] = np.array( - [ - 21, - 40, - 41, - 42, - 43, - 44, - 45, - 46, - 47, - 48, - 49, - 50, - 51, - 52, - 53, - 54, - 139, - 140, - 141, - 142, - 144, - ] -) +part_indices["body"] = np.array([ + 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, + 123, + 124, + 125, + 126, + 127, + 132, + 134, + 135, + 136, + 137, + 138, + 143, +]) +part_indices["torso"] = np.array([ + 0, + 1, + 2, + 3, + 6, + 9, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 22, + 23, + 24, + 55, + 56, + 57, + 58, + 59, + 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, +]) +part_indices["head"] = np.array([ + 12, + 15, + 22, + 23, + 24, + 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, + 125, + 126, + 134, + 136, + 137, +]) +part_indices["face"] = np.array([ + 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, +]) +part_indices["upper"] = np.array([ + 12, + 13, + 14, + 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, +]) +part_indices["hand"] = np.array([ + 20, + 21, + 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, + 128, + 129, + 130, + 131, + 133, + 139, + 140, + 141, + 142, + 144, +]) +part_indices["left_hand"] = np.array([ + 20, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 128, + 129, + 130, + 131, + 133, +]) +part_indices["right_hand"] = np.array([ + 21, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 139, + 140, + 141, + 142, + 144, +]) # kinematic tree head_kin_chain = [15, 12, 9, 6, 3, 0] diff --git a/lib/pixielib/models/encoders.py b/lib/pixielib/models/encoders.py index 0783c9265ab442a259fd693a55039026cc7608db..44f979a2063fe62e3de451bebb267a8852e85955 100755 --- a/lib/pixielib/models/encoders.py +++ b/lib/pixielib/models/encoders.py @@ -1,6 +1,6 @@ import numpy as np -import torch.nn as nn import torch +import torch.nn as nn import torch.nn.functional as F diff --git a/lib/pixielib/models/hrnet.py b/lib/pixielib/models/hrnet.py index c1fd871abf8ae79dd87f96e30d14d726c913db05..665b96efa29fb273b2e28773e5ea35391d99b90e 100644 --- a/lib/pixielib/models/hrnet.py +++ b/lib/pixielib/models/hrnet.py @@ -3,10 +3,10 @@ borrowed from https://github.com/vchoutas/expose/blob/master/expose/models/backb """ import os.path as osp + import torch import torch.nn as nn - -from torchvision.models.resnet import Bottleneck, BasicBlock +from torchvision.models.resnet import BasicBlock, Bottleneck BN_MOMENTUM = 0.1 @@ -15,42 +15,38 @@ def load_HRNet(pretrained=False): hr_net_cfg_dict = { "use_old_impl": False, "pretrained_layers": ["*"], - "stage1": - { - "num_modules": 1, - "num_branches": 1, - "num_blocks": [4], - "num_channels": [64], - "block": "BOTTLENECK", - "fuse_method": "SUM", - }, - "stage2": - { - "num_modules": 1, - "num_branches": 2, - "num_blocks": [4, 4], - "num_channels": [48, 96], - "block": "BASIC", - "fuse_method": "SUM", - }, - "stage3": - { - "num_modules": 4, - "num_branches": 3, - "num_blocks": [4, 4, 4], - "num_channels": [48, 96, 192], - "block": "BASIC", - "fuse_method": "SUM", - }, - "stage4": - { - "num_modules": 3, - "num_branches": 4, - "num_blocks": [4, 4, 4, 4], - "num_channels": [48, 96, 192, 384], - "block": "BASIC", - "fuse_method": "SUM", - }, + "stage1": { + "num_modules": 1, + "num_branches": 1, + "num_blocks": [4], + "num_channels": [64], + "block": "BOTTLENECK", + "fuse_method": "SUM", + }, + "stage2": { + "num_modules": 1, + "num_branches": 2, + "num_blocks": [4, 4], + "num_channels": [48, 96], + "block": "BASIC", + "fuse_method": "SUM", + }, + "stage3": { + "num_modules": 4, + "num_branches": 3, + "num_blocks": [4, 4, 4], + "num_channels": [48, 96, 192], + "block": "BASIC", + "fuse_method": "SUM", + }, + "stage4": { + "num_modules": 3, + "num_branches": 4, + "num_blocks": [4, 4, 4, 4], + "num_channels": [48, 96, 192, 384], + "block": "BASIC", + "fuse_method": "SUM", + }, } hr_net_cfg = hr_net_cfg_dict model = HighResolutionNet(hr_net_cfg) diff --git a/lib/pixielib/models/lbs.py b/lib/pixielib/models/lbs.py index a2252a9a81c7e9ca3633a02cc08f3fafd5bd22cc..7b490bd9bc79a0e252aec2df99bead814edf4195 100755 --- a/lib/pixielib/models/lbs.py +++ b/lib/pixielib/models/lbs.py @@ -14,15 +14,14 @@ # # Contact: ps-license@tuebingen.mpg.de -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division +from __future__ import absolute_import, division, print_function -import numpy as np import os -import yaml + +import numpy as np import torch import torch.nn.functional as F +import yaml from torch import nn diff --git a/lib/pixielib/models/moderators.py b/lib/pixielib/models/moderators.py index 3ab139ac2ad3e0cbd99c8e40dbf6136a37e53cb5..205777192a5601c4f37c75a22981fcda8e0416e0 100644 --- a/lib/pixielib/models/moderators.py +++ b/lib/pixielib/models/moderators.py @@ -3,8 +3,8 @@ # output: fused feature, weight """ import numpy as np -import torch.nn as nn import torch +import torch.nn as nn import torch.nn.functional as F # MLP + temperature softmax diff --git a/lib/pixielib/models/resnet.py b/lib/pixielib/models/resnet.py index 162bc655bff1bd3ca2058334de2e15660de8f5f5..9732daf972dd6d7adebb96927bbc80242a3233dd 100755 --- a/lib/pixielib/models/resnet.py +++ b/lib/pixielib/models/resnet.py @@ -11,13 +11,14 @@ Loads different resnet models mark: copied from pytorch source code """ +import math + +import numpy as np +import torch import torch.nn as nn import torch.nn.functional as F -import torch -from torch.nn.parameter import Parameter -import numpy as np -import math import torchvision +from torch.nn.parameter import Parameter from torchvision import models diff --git a/lib/pixielib/pixie.py b/lib/pixielib/pixie.py index 545bc46f92b73aff4037da7ff3c6ebeba2b4c361..6cadc83bd8b67f32d5819e0d5218de194e63692c 100644 --- a/lib/pixielib/pixie.py +++ b/lib/pixielib/pixie.py @@ -14,21 +14,20 @@ # For commercial licensing contact, please contact ps-license@tuebingen.mpg.de import os -import torch -import torchvision -import torch.nn.functional as F -import torch.nn as nn +import cv2 import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision from skimage.io import imread -import cv2 -from .models.encoders import ResnetEncoder, MLP, HRNEncoder +from .models.encoders import MLP, HRNEncoder, ResnetEncoder from .models.moderators import TempSoftmaxFusion from .models.SMPLX import SMPLX -from .utils import util from .utils import rotation_converter as converter -from .utils import tensor_cropper +from .utils import tensor_cropper, util from .utils.config import cfg @@ -55,9 +54,7 @@ class PIXIE(object): # encode + decode param_dict = self.encode( - {"body": { - "image": data - }}, + {"body": {"image": data}}, threthold=True, keep_local=True, copy_and_paste=False, @@ -559,9 +556,10 @@ class PIXIE(object): } # change the order of face keypoints, to be the same as "standard" 68 keypoints - prediction["face_kpt"] = torch.cat( - [prediction["face_kpt"][:, -17:], prediction["face_kpt"][:, :-17]], dim=1 - ) + prediction["face_kpt"] = torch.cat([ + prediction["face_kpt"][:, -17:], prediction["face_kpt"][:, :-17] + ], + dim=1) prediction.update(param_dict) diff --git a/lib/pixielib/utils/array_cropper.py b/lib/pixielib/utils/array_cropper.py index fbee84b6a6f0f3dcad7fcd6b33bf03faf56be625..d18e15b504fd7d894dcb3ed72830374a9837e42b 100644 --- a/lib/pixielib/utils/array_cropper.py +++ b/lib/pixielib/utils/array_cropper.py @@ -8,7 +8,7 @@ only support crop to squared images """ import numpy as np -from skimage.transform import estimate_transform, warp, resize, rescale +from skimage.transform import estimate_transform, rescale, resize, warp def points2bbox(points, points_scale=None): @@ -47,13 +47,11 @@ def crop_array(image, center, bboxsize, crop_size): tform: 3x3 affine matrix """ # points: top-left, top-right, bottom-right - src_pts = np.array( - [ - [center[0] - bboxsize / 2, center[1] - bboxsize / 2], - [center[0] + bboxsize / 2, center[1] - bboxsize / 2], - [center[0] + bboxsize / 2, center[1] + bboxsize / 2], - ] - ) + src_pts = np.array([ + [center[0] - bboxsize / 2, center[1] - bboxsize / 2], + [center[0] + bboxsize / 2, center[1] - bboxsize / 2], + [center[0] + bboxsize / 2, center[1] + bboxsize / 2], + ]) DST_PTS = np.array([[0, 0], [crop_size - 1, 0], [crop_size - 1, crop_size - 1]]) # estimate transformation between points diff --git a/lib/pixielib/utils/config.py b/lib/pixielib/utils/config.py index 115a38e9c52b7cf025defa4a3d37d9490fc71833..cd8c87dc43a3adb4ae7c3f3a69595b0fb35d2d07 100644 --- a/lib/pixielib/utils/config.py +++ b/lib/pixielib/utils/config.py @@ -1,11 +1,12 @@ """ Default config for PIXIE """ -from yacs.config import CfgNode as CN import argparse -import yaml import os +import yaml +from yacs.config import CfgNode as CN + cfg = CN() abs_pixie_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) diff --git a/lib/pixielib/utils/renderer.py b/lib/pixielib/utils/renderer.py index eb2dc795e01b3e5c78a4ce848777d6cbc5558401..efcc560d7856959f70005858432171776b1bffc7 100755 --- a/lib/pixielib/utils/renderer.py +++ b/lib/pixielib/utils/renderer.py @@ -3,12 +3,12 @@ Author: Yao Feng Copyright (c) 2020, Yao Feng All rights reserved. """ +import imageio import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from skimage.io import imread -import imageio from . import util @@ -16,17 +16,18 @@ from . import util def set_rasterizer(type="pytorch3d"): if type == "pytorch3d": global Meshes, load_obj, rasterize_meshes - from pytorch3d.structures import Meshes from pytorch3d.io import load_obj from pytorch3d.renderer.mesh import rasterize_meshes + from pytorch3d.structures import Meshes elif type == "standard": global standard_rasterize, load_obj import os - from .util import load_obj # Use JIT Compiling Extensions # ref: https://pytorch.org/tutorials/advanced/cpp_extension.html - from torch.utils.cpp_extension import load, CUDA_HOME + from torch.utils.cpp_extension import CUDA_HOME, load + + from .util import load_obj curr_dir = os.path.dirname(__file__) standard_rasterize_cuda = load( @@ -207,19 +208,17 @@ class SRenderY(nn.Module): # SH factors for lighting pi = np.pi - constant_factor = torch.tensor( - [ - 1 / np.sqrt(4 * pi), - ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), - ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), - ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), - (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), - (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), - (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), - (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))), - (pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi))), - ] - ).float() + constant_factor = torch.tensor([ + 1 / np.sqrt(4 * pi), + ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), + ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), + ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), + (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), + (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), + (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), + (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))), + (pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi))), + ]).float() self.register_buffer("constant_factor", constant_factor) def forward( @@ -310,17 +309,17 @@ class SRenderY(nn.Module): normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]), lights, ) - shading_images = shading.reshape( - [batch_size, albedo_images.shape[2], albedo_images.shape[3], 3] - ).permute(0, 3, 1, 2) + shading_images = shading.reshape([ + batch_size, albedo_images.shape[2], albedo_images.shape[3], 3 + ]).permute(0, 3, 1, 2) else: shading = self.add_directionlight( normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]), lights, ) - shading_images = shading.reshape( - [batch_size, albedo_images.shape[2], albedo_images.shape[3], 3] - ).permute(0, 3, 1, 2) + shading_images = shading.reshape([ + batch_size, albedo_images.shape[2], albedo_images.shape[3], 3 + ]).permute(0, 3, 1, 2) images = albedo_images * shading_images else: images = albedo_images @@ -402,9 +401,8 @@ class SRenderY(nn.Module): ) # normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.) # normals_dot_lights = (normals[:,None,:,:]*directions_to_lights).sum(dim=3) - normals_dot_lights = torch.clamp( - (normals[:, None, :, :] * directions_to_lights).sum(dim=3), 0.0, 1.0 - ) + normals_dot_lights = torch.clamp((normals[:, None, :, :] * directions_to_lights).sum(dim=3), + 0.0, 1.0) shading = normals_dot_lights[:, :, :, None] * light_intensities[:, :, None, :] return shading.mean(1) diff --git a/lib/pixielib/utils/rotation_converter.py b/lib/pixielib/utils/rotation_converter.py index f8057cab4e0f84d035a0b8f964823bd61e91dae4..e66386d6427db73e50f791a6fd479f085230da51 100644 --- a/lib/pixielib/utils/rotation_converter.py +++ b/lib/pixielib/utils/rotation_converter.py @@ -1,6 +1,7 @@ +import numpy as np import torch import torch.nn.functional as F -import numpy as np + """ Rotation Converter This function is borrowed from https://github.com/kornia/kornia diff --git a/lib/pixielib/utils/tensor_cropper.py b/lib/pixielib/utils/tensor_cropper.py index c486f7709ad9216080102ee275f7165d276eb0ce..6863ff044a71d054b460f78557f6d09d11f20a30 100644 --- a/lib/pixielib/utils/tensor_cropper.py +++ b/lib/pixielib/utils/tensor_cropper.py @@ -8,9 +8,9 @@ only support crop to squared images """ import torch from kornia.geometry.transform.imgwarp import ( - warp_perspective, get_perspective_transform, warp_affine, + warp_perspective, ) diff --git a/lib/pixielib/utils/util.py b/lib/pixielib/utils/util.py index 566eda3a6e6ddf7f236bf4e20bf7220b39981ce3..5affbd8314fe175cd673ad07a64aa6964e581cbf 100755 --- a/lib/pixielib/utils/util.py +++ b/lib/pixielib/utils/util.py @@ -1,10 +1,11 @@ +import os +import pickle +from collections import OrderedDict + +import cv2 import numpy as np import torch import torch.nn.functional as F -from collections import OrderedDict -import os -import cv2 -import pickle # ---------------------------- process/generate vertices, normals, faces diff --git a/lib/pymafx/core/cfgs.py b/lib/pymafx/core/cfgs.py index c970c6c0caafe7a4c2f3abbb311adcd0cef42b94..17abd247de8d335131d8facc866d95e485ea9a7a 100644 --- a/lib/pymafx/core/cfgs.py +++ b/lib/pymafx/core/cfgs.py @@ -14,12 +14,13 @@ # # Contact: ps-license@tuebingen.mpg.de -import os +import argparse import json +import os import random import string -import argparse from datetime import datetime + from yacs.config import CfgNode as CN # Configuration variables diff --git a/lib/pymafx/core/constants.py b/lib/pymafx/core/constants.py index 5354a289f892a764a16221b469fc49794ff54127..24077a0c7b89215315b39dcbdf9335193ee6ce50 100644 --- a/lib/pymafx/core/constants.py +++ b/lib/pymafx/core/constants.py @@ -79,55 +79,16 @@ 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 + '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 @@ -163,30 +124,11 @@ SMPL_J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, SMPLX2SMPL_J45 = [i for i in range(22)] + [30, 45] + [i for i in range(55, 55 + 21)] SMPL_PART_ID = { - 'rightHand': 1, - 'rightUpLeg': 2, - 'leftArm': 3, - 'leftLeg': 4, - 'leftToeBase': 5, - 'leftFoot': 6, - 'spine1': 7, - 'spine2': 8, - 'leftShoulder': 9, - 'rightShoulder': 10, - 'rightFoot': 11, - 'head': 12, - 'rightArm': 13, - 'leftHandIndex1': 14, - 'rightLeg': 15, - 'rightHandIndex1': 16, - 'leftForeArm': 17, - 'rightForeArm': 18, - 'neck': 19, - 'rightToeBase': 20, - 'spine': 21, - 'leftUpLeg': 22, - 'leftHand': 23, - 'hips': 24 + 'rightHand': 1, 'rightUpLeg': 2, 'leftArm': 3, 'leftLeg': 4, 'leftToeBase': 5, 'leftFoot': 6, + 'spine1': 7, 'spine2': 8, 'leftShoulder': 9, 'rightShoulder': 10, 'rightFoot': 11, 'head': 12, + 'rightArm': 13, 'leftHandIndex1': 14, 'rightLeg': 15, 'rightHandIndex1': 16, 'leftForeArm': 17, + 'rightForeArm': 18, 'neck': 19, 'rightToeBase': 20, 'spine': 21, 'leftUpLeg': 22, 'leftHand': + 23, 'hips': 24 } # MANO_NAMES = [ diff --git a/lib/pymafx/models/attention.py b/lib/pymafx/models/attention.py index b0f7d3c5c63ba1471ff15ee1a3cf0d8c94a17699..b87008b0e68541970943eb3e04027223e85c68e8 100644 --- a/lib/pymafx/models/attention.py +++ b/lib/pymafx/models/attention.py @@ -4,15 +4,29 @@ Licensed under the MIT license. """ -from __future__ import absolute_import, division, print_function, unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) +import code import logging import math import os -import code + import torch from torch import nn -from .transformers.bert.modeling_bert import BertPreTrainedModel, BertEmbeddings, BertPooler, BertIntermediate, BertOutput, BertSelfOutput + +from .transformers.bert.modeling_bert import ( + BertEmbeddings, + BertIntermediate, + BertOutput, + BertPooler, + BertPreTrainedModel, + BertSelfOutput, +) # import src.modeling.data.config as cfg # from src.modeling._gcnn import GraphConvolution, GraphResBlock from .transformers.bert.modeling_utils import prune_linear_layer diff --git a/lib/pymafx/models/hmr.py b/lib/pymafx/models/hmr.py index da5459d355d3a3f00c53638a376ab3143b23c01e..e9ba5759d7a59cb2c5b9ce0964aaf899c27a1e8a 100755 --- a/lib/pymafx/models/hmr.py +++ b/lib/pymafx/models/hmr.py @@ -1,13 +1,14 @@ # This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/hmr.py +import logging +import math + +import numpy as np import torch import torch.nn as nn import torchvision.models.resnet as resnet -import numpy as np -import math -from lib.net.geometry import rot6d_to_rotmat -import logging +from lib.net.geometry import rot6d_to_rotmat logger = logging.getLogger(__name__) diff --git a/lib/pymafx/models/hr_module.py b/lib/pymafx/models/hr_module.py index 7396f1ea59860235db8fdd24434114381c4a7083..ad6243a463a45733a0c518e34c0dbcb115d39bcc 100644 --- a/lib/pymafx/models/hr_module.py +++ b/lib/pymafx/models/hr_module.py @@ -1,13 +1,14 @@ +import logging import os + import torch -import torch.nn as nn import torch._utils +import torch.nn as nn import torch.nn.functional as F + # from core.cfgs import cfg from .res_module import BasicBlock, Bottleneck -import logging - logger = logging.getLogger(__name__) BN_MOMENTUM = 0.1 diff --git a/lib/pymafx/models/maf_extractor.py b/lib/pymafx/models/maf_extractor.py index 34237bc55663dcbcbd67beb4c5d0b6e693aae266..ffe4e73427e30848798df2f57e835a8b10ae2934 100644 --- a/lib/pymafx/models/maf_extractor.py +++ b/lib/pymafx/models/maf_extractor.py @@ -1,23 +1,23 @@ # This script is borrowed and extended from https://github.com/shunsukesaito/PIFu/blob/master/lib/model/SurfaceClassifier.py -import torch -import scipy +import logging + import numpy as np +import scipy +import torch import torch.nn as nn import torch.nn.functional as F from lib.pymafx.core import path_config from lib.pymafx.utils.geometry import projection -import logging - logger = logging.getLogger(__name__) +from lib.pymafx.utils.imutils import j2d_processing + from .transformers.net_utils import PosEnSine from .transformers.transformer_basics import OurMultiheadAttention -from lib.pymafx.utils.imutils import j2d_processing - class TransformerDecoderUnit(nn.Module): def __init__( diff --git a/lib/pymafx/models/pose_resnet.py b/lib/pymafx/models/pose_resnet.py index d97b6609cf02fd2a94d2951f82f71de2be2356c0..16b22e815f715d2ae8e5f217431055ee2ba57ddf 100644 --- a/lib/pymafx/models/pose_resnet.py +++ b/lib/pymafx/models/pose_resnet.py @@ -4,12 +4,10 @@ # Written by Bin Xiao (Bin.Xiao@microsoft.com) # ------------------------------------------------------------------------------ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function +from __future__ import absolute_import, division, print_function -import os import logging +import os import torch import torch.nn as nn @@ -277,11 +275,8 @@ class PoseResNet(nn.Module): 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]) + 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]) } diff --git a/lib/pymafx/models/pymaf_net.py b/lib/pymafx/models/pymaf_net.py index ca57e4b1c8ce971d76ce53d02827f441016a19ab..1f46c10f4bb0384277a4951135ce6edaabe5ac7d 100644 --- a/lib/pymafx/models/pymaf_net.py +++ b/lib/pymafx/models/pymaf_net.py @@ -1,21 +1,34 @@ +import logging + +import numpy as np import torch import torch.nn as nn -import numpy as np -from lib.pymafx.core import constants from lib.common.config import cfg -from lib.pymafx.utils.geometry import rot6d_to_rotmat, rotmat_to_rot6d, projection, rotation_matrix_to_angle_axis, compute_twist_rotation -from .maf_extractor import MAF_Extractor, Mesh_Sampler -from .smpl import SMPL, SMPL_MODEL_DIR, SMPL_MEAN_PARAMS, get_partial_smpl, SMPL_Family +from lib.pymafx.core import constants +from lib.pymafx.utils.cam_params import homo_vector +from lib.pymafx.utils.geometry import ( + compute_twist_rotation, + projection, + rot6d_to_rotmat, + rotation_matrix_to_angle_axis, + rotmat_to_rot6d, +) +from lib.pymafx.utils.imutils import j2d_processing from lib.smplx.lbs import batch_rodrigues -from .res_module import IUV_predict_layer + +from .attention import get_att_block from .hr_module import get_hrnet_encoder +from .maf_extractor import MAF_Extractor, Mesh_Sampler from .pose_resnet import get_resnet_encoder -from lib.pymafx.utils.imutils import j2d_processing -from lib.pymafx.utils.cam_params import homo_vector -from .attention import get_att_block - -import logging +from .res_module import IUV_predict_layer +from .smpl import ( + SMPL, + SMPL_MEAN_PARAMS, + SMPL_MODEL_DIR, + SMPL_Family, + get_partial_smpl, +) logger = logging.getLogger(__name__) @@ -169,14 +182,14 @@ class Regressor(nn.Module): if not self.smpl_mode: lhand_mean_rot6d = rotmat_to_rot6d( - batch_rodrigues(self.smpl.model.model_neutral.left_hand_mean.view(-1, 3)).view( - [-1, 3, 3] - ) + batch_rodrigues(self.smpl.model.model_neutral.left_hand_mean.view(-1, 3)).view([ + -1, 3, 3 + ]) ) rhand_mean_rot6d = rotmat_to_rot6d( - batch_rodrigues(self.smpl.model.model_neutral.right_hand_mean.view(-1, 3)).view( - [-1, 3, 3] - ) + batch_rodrigues(self.smpl.model.model_neutral.right_hand_mean.view(-1, 3)).view([ + -1, 3, 3 + ]) ) init_lhand = lhand_mean_rot6d.reshape(-1).unsqueeze(0) init_rhand = rhand_mean_rot6d.reshape(-1).unsqueeze(0) @@ -300,9 +313,9 @@ class Regressor(nn.Module): else: vfov = rw_cam['vfov'][:, None] crop_ratio = rw_cam['crop_ratio'][:, None] - crop_center = rw_cam['bbox_center'] / torch.cat( - [rw_cam['img_w'][:, None], rw_cam['img_h'][:, None]], 1 - ) + crop_center = rw_cam['bbox_center'] / torch.cat([ + rw_cam['img_w'][:, None], rw_cam['img_h'][:, None] + ], 1) xc = torch.cat([xc, vfov, crop_ratio, crop_center], 1) xc = self.fc1(xc) @@ -338,9 +351,9 @@ class Regressor(nn.Module): pred_lhand = self.decrhand(xc_lhand) + pred_lhand if cfg.MODEL.PyMAF.OPT_WRIST: - xc_lhand = torch.cat( - [xc_lhand, pred_shape_lh, pred_orient_lh, pred_cam_lh], 1 - ) + xc_lhand = torch.cat([ + xc_lhand, pred_shape_lh, pred_orient_lh, pred_cam_lh + ], 1) xc_lhand = self.drop3_hand(self.fc3_hand(xc_lhand)) pred_shape_lh = self.decshape_rhand(xc_lhand) + pred_shape_lh @@ -353,9 +366,9 @@ class Regressor(nn.Module): pred_rhand = self.decrhand(xc_rhand) + pred_rhand if cfg.MODEL.MESH_MODEL == 'mano' or cfg.MODEL.PyMAF.OPT_WRIST: - xc_rhand = torch.cat( - [xc_rhand, pred_shape_rh, pred_orient_rh, pred_cam_rh], 1 - ) + xc_rhand = torch.cat([ + xc_rhand, pred_shape_rh, pred_orient_rh, pred_cam_rh + ], 1) xc_rhand = self.drop3_hand(self.fc3_hand(xc_rhand)) pred_shape_rh = self.decshape_rhand(xc_rhand) + pred_shape_rh @@ -363,9 +376,10 @@ class Regressor(nn.Module): pred_cam_rh = self.deccam_rhand(xc_rhand) + pred_cam_rh if cfg.MODEL.MESH_MODEL == 'mano': - pred_cam = torch.cat( - [pred_cam_rh[:, 0:1] * 10., pred_cam_rh[:, 1:] / 10.], dim=1 - ) + pred_cam = torch.cat([ + pred_cam_rh[:, 0:1] * 10., pred_cam_rh[:, 1:] / 10. + ], + dim=1) if 'face' in self.part_names: xc_face = self.drop1_face(self.fc1_face(xc_face)) @@ -374,9 +388,9 @@ class Regressor(nn.Module): pred_exp = self.decexp(xc_face) + pred_exp if cfg.MODEL.MESH_MODEL == 'flame': - xc_face = torch.cat( - [xc_face, pred_shape_fa, pred_orient_fa, pred_cam_fa], 1 - ) + xc_face = torch.cat([ + xc_face, pred_shape_fa, pred_orient_fa, pred_cam_fa + ], 1) xc_face = self.drop3_face(self.fc3_face(xc_face)) pred_shape_fa = self.decshape_face(xc_face) + pred_shape_fa @@ -384,9 +398,10 @@ class Regressor(nn.Module): pred_cam_fa = self.deccam_face(xc_face) + pred_cam_fa if cfg.MODEL.MESH_MODEL == 'flame': - pred_cam = torch.cat( - [pred_cam_fa[:, 0:1] * 10., pred_cam_fa[:, 1:] / 10.], dim=1 - ) + pred_cam = torch.cat([ + pred_cam_fa[:, 0:1] * 10., pred_cam_fa[:, 1:] / 10. + ], + dim=1) if self.full_body_mode or self.body_hand_mode: if cfg.MODEL.PyMAF.PRED_VIS_H: @@ -500,15 +515,13 @@ class Regressor(nn.Module): opt_lelbow = torch.stack(opt_lelbow_filtered) opt_relbow = torch.stack(opt_relbow_filtered) - pred_rotmat_body = torch.cat( - [ - pred_rotmat_body[:, :18], - opt_lelbow.unsqueeze(1), - opt_relbow.unsqueeze(1), - opt_lwrist.unsqueeze(1), - opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:] - ], 1 - ) + pred_rotmat_body = torch.cat([ + pred_rotmat_body[:, :18], + opt_lelbow.unsqueeze(1), + opt_relbow.unsqueeze(1), + opt_lwrist.unsqueeze(1), + opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:] + ], 1) else: if cfg.MODEL.PyMAF.PRED_VIS_H and global_iter == ( cfg.MODEL.PyMAF.N_ITER - 1 @@ -527,13 +540,11 @@ class Regressor(nn.Module): opt_lwrist = torch.stack(opt_lwrist_filtered) opt_rwrist = torch.stack(opt_rwrist_filtered) - pred_rotmat_body = torch.cat( - [ - pred_rotmat_body[:, :20], - opt_lwrist.unsqueeze(1), - opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:] - ], 1 - ) + pred_rotmat_body = torch.cat([ + pred_rotmat_body[:, :20], + opt_lwrist.unsqueeze(1), + opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:] + ], 1) if self.hand_only_mode: pred_rotmat_rh = rot6d_to_rotmat( @@ -630,19 +641,15 @@ class Regressor(nn.Module): elif self.face_only_mode: pred_joints_full = pred_output.face_joints elif self.smplx_mode: - pred_joints_full = torch.cat( - [ - pred_joints, pred_output.lhand_joints, pred_output.rhand_joints, - pred_output.face_joints, pred_output.lfoot_joints, pred_output.rfoot_joints - ], - dim=1 - ) + pred_joints_full = torch.cat([ + pred_joints, pred_output.lhand_joints, pred_output.rhand_joints, + pred_output.face_joints, pred_output.lfoot_joints, pred_output.rfoot_joints + ], + dim=1) else: pred_joints_full = pred_joints pred_keypoints_2d = projection( - pred_joints_full, { - **rw_cam, 'cam_sxy': pred_cam - }, iwp_mode=cfg.MODEL.USE_IWP_CAM + pred_joints_full, {**rw_cam, 'cam_sxy': pred_cam}, iwp_mode=cfg.MODEL.USE_IWP_CAM ) if cfg.MODEL.USE_IWP_CAM: # Normalize keypoints to [-1,1] @@ -661,119 +668,109 @@ class Regressor(nn.Module): else: kp_3d = pred_joints pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape(-1, 72) - output.update( - { - 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), - 'verts': pred_vertices, - 'kp_2d': pred_keypoints_2d[:, :len_b_kp], - 'kp_3d': kp_3d, - 'pred_joints': pred_joints, - 'smpl_kp_3d': pred_output.smpl_joints, - 'rotmat': pred_rotmat, - 'pred_cam': pred_cam, - 'pred_shape': pred_shape, - 'pred_pose': pred_pose, - } - ) + output.update({ + 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), + 'verts': pred_vertices, + 'kp_2d': pred_keypoints_2d[:, :len_b_kp], + 'kp_3d': kp_3d, + 'pred_joints': pred_joints, + 'smpl_kp_3d': pred_output.smpl_joints, + 'rotmat': pred_rotmat, + 'pred_cam': pred_cam, + 'pred_shape': pred_shape, + 'pred_pose': pred_pose, + }) # if self.full_body_mode: if self.smplx_mode: # assert pred_keypoints_2d.shape[1] == 144 len_h_kp = len(constants.HAND_NAMES) len_f_kp = len(constants.FACIAL_LANDMARKS) len_feet_kp = 2 * len(constants.FOOT_NAMES) - output.update( - { - 'smplx_verts': - pred_output.smplx_vertices if cfg.MODEL.EVAL_MODE else None, - 'pred_lhand': - pred_lhand, - 'pred_rhand': - pred_rhand, - 'pred_face': - pred_face, - 'pred_exp': - pred_exp, - 'verts_lh': - pred_output.lhand_vertices, - 'verts_rh': - pred_output.rhand_vertices, + output.update({ + 'smplx_verts': + pred_output.smplx_vertices if cfg.MODEL.EVAL_MODE else None, + 'pred_lhand': + pred_lhand, + 'pred_rhand': + pred_rhand, + 'pred_face': + pred_face, + 'pred_exp': + pred_exp, + 'verts_lh': + pred_output.lhand_vertices, + 'verts_rh': + pred_output.rhand_vertices, # 'pred_arm_rotmat': pred_arm_rotmat, # 'pred_hfrotmat': pred_hfrotmat, - 'pred_lhand_rotmat': - pred_lhand_rotmat, - 'pred_rhand_rotmat': - pred_rhand_rotmat, - 'pred_face_rotmat': - pred_face_rotmat, - 'pred_lhand_kp3d': - pred_output.lhand_joints, - 'pred_rhand_kp3d': - pred_output.rhand_joints, - 'pred_face_kp3d': - pred_output.face_joints, - 'pred_lhand_kp2d': - pred_keypoints_2d[:, len_b_kp:len_b_kp + len_h_kp], - 'pred_rhand_kp2d': - pred_keypoints_2d[:, len_b_kp + len_h_kp:len_b_kp + len_h_kp * 2], - 'pred_face_kp2d': - pred_keypoints_2d[:, len_b_kp + len_h_kp * 2:len_b_kp + len_h_kp * 2 + - len_f_kp], - 'pred_feet_kp2d': - pred_keypoints_2d[:, len_b_kp + len_h_kp * 2 + len_f_kp:len_b_kp + - len_h_kp * 2 + len_f_kp + len_feet_kp], - } - ) + 'pred_lhand_rotmat': + pred_lhand_rotmat, + 'pred_rhand_rotmat': + pred_rhand_rotmat, + 'pred_face_rotmat': + pred_face_rotmat, + 'pred_lhand_kp3d': + pred_output.lhand_joints, + 'pred_rhand_kp3d': + pred_output.rhand_joints, + 'pred_face_kp3d': + pred_output.face_joints, + 'pred_lhand_kp2d': + pred_keypoints_2d[:, len_b_kp:len_b_kp + len_h_kp], + 'pred_rhand_kp2d': + pred_keypoints_2d[:, len_b_kp + len_h_kp:len_b_kp + len_h_kp * 2], + 'pred_face_kp2d': + pred_keypoints_2d[:, + len_b_kp + len_h_kp * 2:len_b_kp + len_h_kp * 2 + len_f_kp], + 'pred_feet_kp2d': + pred_keypoints_2d[:, len_b_kp + len_h_kp * 2 + len_f_kp:len_b_kp + + len_h_kp * 2 + len_f_kp + len_feet_kp], + }) if cfg.MODEL.PyMAF.OPT_WRIST: - output.update( - { - 'pred_orient_lh': pred_orient_lh, - 'pred_shape_lh': pred_shape_lh, - 'pred_orient_rh': pred_orient_rh, - 'pred_shape_rh': pred_shape_rh, - 'pred_cam_fa': pred_cam_fa, - 'pred_cam_lh': pred_cam_lh, - 'pred_cam_rh': pred_cam_rh, - } - ) + output.update({ + 'pred_orient_lh': pred_orient_lh, + 'pred_shape_lh': pred_shape_lh, + 'pred_orient_rh': pred_orient_rh, + 'pred_shape_rh': pred_shape_rh, + 'pred_cam_fa': pred_cam_fa, + 'pred_cam_lh': pred_cam_lh, + 'pred_cam_rh': pred_cam_rh, + }) if cfg.MODEL.PyMAF.PRED_VIS_H: output.update({'pred_vis_hands': pred_vis_hands}) elif self.hand_only_mode: # hand mesh out assert pred_keypoints_2d.shape[1] == 21 - output.update( - { - 'theta': pred_cam, - 'pred_cam': pred_cam, - 'pred_rhand': pred_rhand, - 'pred_rhand_rotmat': pred_rotmat_rh[:, 1:], - 'pred_orient_rh': pred_orient_rh, - 'pred_orient_rh_rotmat': pred_rotmat_rh[:, 0], - 'verts_rh': pred_output.rhand_vertices, - 'pred_cam_rh': pred_cam_rh, - 'pred_shape_rh': pred_shape_rh, - 'pred_rhand_kp3d': pred_output.rhand_joints, - 'pred_rhand_kp2d': pred_keypoints_2d, - } - ) + output.update({ + 'theta': pred_cam, + 'pred_cam': pred_cam, + 'pred_rhand': pred_rhand, + 'pred_rhand_rotmat': pred_rotmat_rh[:, 1:], + 'pred_orient_rh': pred_orient_rh, + 'pred_orient_rh_rotmat': pred_rotmat_rh[:, 0], + 'verts_rh': pred_output.rhand_vertices, + 'pred_cam_rh': pred_cam_rh, + 'pred_shape_rh': pred_shape_rh, + 'pred_rhand_kp3d': pred_output.rhand_joints, + 'pred_rhand_kp2d': pred_keypoints_2d, + }) elif self.face_only_mode: # face mesh out assert pred_keypoints_2d.shape[1] == 68 - output.update( - { - 'theta': pred_cam, - 'pred_cam': pred_cam, - 'pred_face': pred_face, - 'pred_exp': pred_exp, - 'pred_face_rotmat': pred_rotmat_fa[:, 1:], - 'pred_orient_fa': pred_orient_fa, - 'pred_orient_fa_rotmat': pred_rotmat_fa[:, 0], - 'verts_fa': pred_output.flame_vertices, - 'pred_cam_fa': pred_cam_fa, - 'pred_shape_fa': pred_shape_fa, - 'pred_face_kp3d': pred_output.face_joints, - 'pred_face_kp2d': pred_keypoints_2d, - } - ) + output.update({ + 'theta': pred_cam, + 'pred_cam': pred_cam, + 'pred_face': pred_face, + 'pred_exp': pred_exp, + 'pred_face_rotmat': pred_rotmat_fa[:, 1:], + 'pred_orient_fa': pred_orient_fa, + 'pred_orient_fa_rotmat': pred_rotmat_fa[:, 0], + 'verts_fa': pred_output.flame_vertices, + 'pred_cam_fa': pred_cam_fa, + 'pred_shape_fa': pred_shape_fa, + 'pred_face_kp3d': pred_output.face_joints, + 'pred_face_kp2d': pred_keypoints_2d, + }) return output @@ -946,16 +943,16 @@ class PyMAF(nn.Module): self.dp_head_hf['hand'] = IUV_predict_layer( feat_dim=hf_sfeat_dim[-1], mode='pncc' ) - self.part_module_names['hand'].update( - {'dp_head_hf.hand': self.dp_head_hf['hand']} - ) + self.part_module_names['hand'].update({ + 'dp_head_hf.hand': self.dp_head_hf['hand'] + }) if 'face' in bhf_names: self.dp_head_hf['face'] = IUV_predict_layer( feat_dim=hf_sfeat_dim[-1], mode='pncc' ) - self.part_module_names['face'].update( - {'dp_head_hf.face': self.dp_head_hf['face']} - ) + self.part_module_names['face'].update({ + 'dp_head_hf.face': self.dp_head_hf['face'] + }) smpl2limb_vert_faces = get_partial_smpl() @@ -964,10 +961,10 @@ class PyMAF(nn.Module): # grid points for grid feature extraction grid_size = 21 - xv, yv = torch.meshgrid( - [torch.linspace(-1, 1, grid_size), - torch.linspace(-1, 1, grid_size)] - ) + xv, yv = torch.meshgrid([ + torch.linspace(-1, 1, grid_size), + torch.linspace(-1, 1, grid_size) + ]) grid_points = torch.stack([xv.reshape(-1), yv.reshape(-1)]).unsqueeze(0) self.register_buffer('grid_points', grid_points) grid_feat_dim = grid_size * grid_size * cfg.MODEL.PyMAF.MLP_DIM[-1] @@ -995,9 +992,10 @@ class PyMAF(nn.Module): bhf_att_feat_dim.update({'hand': 1024}) if 'face' in self.bhf_names: - bhf_ma_feat_dim.update( - {'face': len(constants.FACIAL_LANDMARKS) * cfg.MODEL.PyMAF.HF_MLP_DIM[-1]} - ) + bhf_ma_feat_dim.update({ + 'face': + len(constants.FACIAL_LANDMARKS) * cfg.MODEL.PyMAF.HF_MLP_DIM[-1] + }) if self.fuse_grid_align: bhf_att_feat_dim.update({'face': 1024}) @@ -1022,9 +1020,10 @@ class PyMAF(nn.Module): ) for part in bhf_names: - self.part_module_names[part].update( - {f'align_attention.{part}': self.align_attention[part]} - ) + self.part_module_names[part].update({ + f'align_attention.{part}': + self.align_attention[part] + }) if self.fuse_grid_align: self.att_feat_reduce = get_fusion_modules( @@ -1035,9 +1034,10 @@ class PyMAF(nn.Module): out_feat_len=bhf_att_feat_dim ) for part in bhf_names: - self.part_module_names[part].update( - {f'att_feat_reduce.{part}': self.att_feat_reduce[part]} - ) + self.part_module_names[part].update({ + f'att_feat_reduce.{part}': + self.att_feat_reduce[part] + }) # build regressor for parameter prediction self.regressor = nn.ModuleList() @@ -1109,21 +1109,25 @@ class PyMAF(nn.Module): # assign sub-regressor to each part for dec_name, dec_module in self.regressor[-1].named_children(): if 'hand' in dec_name: - self.part_module_names['hand'].update( - {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module} - ) + self.part_module_names['hand'].update({ + 'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): + dec_module + }) elif 'face' in dec_name or 'head' in dec_name or 'exp' in dec_name: - self.part_module_names['face'].update( - {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module} - ) + self.part_module_names['face'].update({ + 'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): + dec_module + }) elif 'res' in dec_name or 'vis' in dec_name: - self.part_module_names['link'].update( - {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module} - ) + self.part_module_names['link'].update({ + 'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): + dec_module + }) elif 'body' in self.part_module_names: - self.part_module_names['body'].update( - {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module} - ) + self.part_module_names['body'].update({ + 'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): + dec_module + }) # mesh-aligned feature extractor self.maf_extractor = nn.ModuleDict() @@ -1373,17 +1377,15 @@ class PyMAF(nn.Module): pred_cam = mesh_output['pred_cam'].detach() pred_rhand_v = self.mano_sampler(mesh_output['verts_rh']) pred_rhand_proj = projection( - pred_rhand_v, { - **rw_cam, 'cam_sxy': pred_cam - }, iwp_mode=cfg.MODEL.USE_IWP_CAM + pred_rhand_v, {**rw_cam, 'cam_sxy': pred_cam}, iwp_mode=cfg.MODEL.USE_IWP_CAM ) if cfg.MODEL.USE_IWP_CAM: pred_rhand_proj = pred_rhand_proj / (224. / 2.) else: pred_rhand_proj = j2d_processing(pred_rhand_proj, rw_cam['kps_transf']) proj_hf_center = { - 'rhand': - mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1) + 'rhand': mesh_output['pred_rhand_kp2d'][:, + self.hf_root_idx['rhand']].unsqueeze(1) } proj_hf_pts = { 'rhand': torch.cat([proj_hf_center['rhand'], pred_rhand_proj], dim=1) @@ -1392,9 +1394,7 @@ class PyMAF(nn.Module): pred_cam = mesh_output['pred_cam'].detach() pred_face_v = mesh_output['pred_face_kp3d'] pred_face_proj = projection( - pred_face_v, { - **rw_cam, 'cam_sxy': pred_cam - }, iwp_mode=cfg.MODEL.USE_IWP_CAM + pred_face_v, {**rw_cam, 'cam_sxy': pred_cam}, iwp_mode=cfg.MODEL.USE_IWP_CAM ) if cfg.MODEL.USE_IWP_CAM: pred_face_proj = pred_face_proj / (224. / 2.) @@ -1409,9 +1409,7 @@ class PyMAF(nn.Module): pred_rhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2rhand]) pred_hand_v = torch.cat([pred_lhand_v, pred_rhand_v], dim=1) pred_hand_proj = projection( - pred_hand_v, { - **rw_cam, 'cam_sxy': pred_cam - }, iwp_mode=cfg.MODEL.USE_IWP_CAM + pred_hand_v, {**rw_cam, 'cam_sxy': pred_cam}, iwp_mode=cfg.MODEL.USE_IWP_CAM ) if cfg.MODEL.USE_IWP_CAM: pred_hand_proj = pred_hand_proj / (224. / 2.) @@ -1419,29 +1417,25 @@ class PyMAF(nn.Module): pred_hand_proj = j2d_processing(pred_hand_proj, rw_cam['kps_transf']) proj_hf_center = { - 'lhand': - mesh_output['pred_lhand_kp2d'][:, self.hf_root_idx['lhand']].unsqueeze(1), - 'rhand': - mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1), + 'lhand': mesh_output['pred_lhand_kp2d'][:, + self.hf_root_idx['lhand']].unsqueeze(1), + 'rhand': mesh_output['pred_rhand_kp2d'][:, + self.hf_root_idx['rhand']].unsqueeze(1), } proj_hf_pts = { 'lhand': - torch.cat( - [proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], dim=1 - ), + torch.cat([proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], + dim=1), 'rhand': - torch.cat( - [proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], dim=1 - ), + torch.cat([proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], + dim=1), } elif self.full_body_mode: pred_lhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2lhand]) pred_rhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2rhand]) pred_hand_v = torch.cat([pred_lhand_v, pred_rhand_v], dim=1) pred_hand_proj = projection( - pred_hand_v, { - **rw_cam, 'cam_sxy': pred_cam - }, iwp_mode=cfg.MODEL.USE_IWP_CAM + pred_hand_v, {**rw_cam, 'cam_sxy': pred_cam}, iwp_mode=cfg.MODEL.USE_IWP_CAM ) if cfg.MODEL.USE_IWP_CAM: pred_hand_proj = pred_hand_proj / (224. / 2.) @@ -1449,24 +1443,19 @@ class PyMAF(nn.Module): pred_hand_proj = j2d_processing(pred_hand_proj, rw_cam['kps_transf']) proj_hf_center = { - 'lhand': - mesh_output['pred_lhand_kp2d'][:, self.hf_root_idx['lhand']].unsqueeze(1), - 'rhand': - mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1), - 'face': - mesh_output['pred_face_kp2d'][:, self.hf_root_idx['face']].unsqueeze(1) + 'lhand': mesh_output['pred_lhand_kp2d'][:, + self.hf_root_idx['lhand']].unsqueeze(1), + 'rhand': mesh_output['pred_rhand_kp2d'][:, + self.hf_root_idx['rhand']].unsqueeze(1), + 'face': mesh_output['pred_face_kp2d'][:, self.hf_root_idx['face']].unsqueeze(1) } proj_hf_pts = { 'lhand': - torch.cat( - [proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], dim=1 - ), - 'rhand': - torch.cat( - [proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], dim=1 - ), - 'face': - torch.cat([proj_hf_center['face'], mesh_output['pred_face_kp2d']], dim=1) + torch.cat([proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], + dim=1), 'rhand': + torch.cat([proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], + dim=1), 'face': + torch.cat([proj_hf_center['face'], mesh_output['pred_face_kp2d']], dim=1) } # extract mesh-aligned features for the hand / face part @@ -1542,9 +1531,10 @@ class PyMAF(nn.Module): limb_grid_feature_ctd = self.maf_extractor[hf_key][limb_rf_i].sampling( grid_points, im_feat=limb_feat_i, reduce_dim=limb_reduce_dim ) - limb_grid_ref_feat_ctd = torch.cat( - [limb_grid_feature_ctd, limb_ref_feat_ctd], dim=-1 - ).permute(0, 2, 1) + limb_grid_ref_feat_ctd = torch.cat([ + limb_grid_feature_ctd, limb_ref_feat_ctd + ], + dim=-1).permute(0, 2, 1) if cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT: att_ref_feat_ctd = self.align_attention[hf_key][ @@ -1581,9 +1571,7 @@ class PyMAF(nn.Module): ref_feature = self.maf_extractor['body'][rf_i]( pred_smpl_verts_ds, im_feat=s_feat_i, - cam={ - **rw_cam, 'cam_sxy': pred_cam - }, + cam={**rw_cam, 'cam_sxy': pred_cam}, add_att=True, reduce_dim=reduce_dim ) # [B, 431 * n_feat] diff --git a/lib/pymafx/models/res_module.py b/lib/pymafx/models/res_module.py index 94de7ecaa2ba3ead51c5f960e0ae08b806d9cd80..69ca53960d7ce7159eb25deb8684f1605eb2c830 100644 --- a/lib/pymafx/models/res_module.py +++ b/lib/pymafx/models/res_module.py @@ -1,18 +1,18 @@ # 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 +from __future__ import absolute_import, division, print_function +import logging import os +from collections import OrderedDict + import torch import torch.nn as nn import torch.nn.functional as F -from collections import OrderedDict + from lib.pymafx.core.cfgs import cfg -# from .transformers.tokenlearner import TokenLearner -import logging +# from .transformers.tokenlearner import TokenLearner logger = logging.getLogger(__name__) @@ -119,11 +119,8 @@ class Bottleneck(nn.Module): 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]) + 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]) } diff --git a/lib/pymafx/models/smpl.py b/lib/pymafx/models/smpl.py index 0a69eaf24e518545542cdd1eb55c819a549e8d55..6dcb6127886e9671fde6a4036d0889ab39ff2b66 100644 --- a/lib/pymafx/models/smpl.py +++ b/lib/pymafx/models/smpl.py @@ -1,20 +1,25 @@ # This script is extended based on https://github.com/nkolot/SPIN/blob/master/models/smpl.py -from typing import Optional +import json +import os +import pickle from dataclasses import dataclass +from typing import Optional -import os +import numpy as np import torch import torch.nn as nn -import numpy as np -import pickle + +from lib.pymafx.core import constants, path_config from lib.smplx import SMPL as _SMPL -from lib.smplx import SMPLXLayer, MANOLayer, FLAMELayer -from lib.smplx.lbs import batch_rodrigues, transform_mat, vertices2joints, blend_shapes +from lib.smplx import FLAMELayer, MANOLayer, SMPLXLayer from lib.smplx.body_models import SMPLXOutput -import json - -from lib.pymafx.core import path_config, constants +from lib.smplx.lbs import ( + batch_rodrigues, + blend_shapes, + transform_mat, + vertices2joints, +) SMPL_MEAN_PARAMS = path_config.SMPL_MEAN_PARAMS SMPL_MODEL_DIR = path_config.SMPL_MODEL_DIR @@ -134,11 +139,11 @@ class SMPL(_SMPL): ).contiguous() # Concatenate all pose vectors - full_pose = torch.cat( - [global_orient.reshape(-1, 1, 3, 3), - body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3)], - dim=1 - ) + full_pose = torch.cat([ + global_orient.reshape(-1, 1, 3, 3), + body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3) + ], + dim=1) rot_mats = full_pose.view(batch_size, -1, 3, 3) @@ -279,18 +284,16 @@ class SMPLX(SMPLXLayer): -1).contiguous() # Concatenate all pose vectors - full_pose = torch.cat( - [ - global_orient.reshape(-1, 1, 3, 3), - body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3), - jaw_pose.reshape(-1, 1, 3, 3), - leye_pose.reshape(-1, 1, 3, 3), - reye_pose.reshape(-1, 1, 3, 3), - left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3), - right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3) - ], - dim=1 - ) + full_pose = torch.cat([ + global_orient.reshape(-1, 1, 3, 3), + body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3), + jaw_pose.reshape(-1, 1, 3, 3), + leye_pose.reshape(-1, 1, 3, 3), + reye_pose.reshape(-1, 1, 3, 3), + left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3), + right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3) + ], + dim=1) rot_mats = full_pose.view(batch_size, -1, 3, 3) @@ -339,22 +342,20 @@ class SMPLX_ALL(nn.Module): self.genders = ['neutral'] for gender in self.genders: assert gender in ['male', 'female', 'neutral'] - self.model_dict = nn.ModuleDict( - { - gender: SMPLX( - path_config.SMPL_MODEL_DIR, - gender=gender, - ext='npz', - num_betas=numBetas, - use_pca=False, - batch_size=batch_size, - use_face_contour=use_face_contour, - num_pca_comps=45, - **kwargs - ) - for gender in self.genders - } - ) + self.model_dict = nn.ModuleDict({ + gender: SMPLX( + path_config.SMPL_MODEL_DIR, + gender=gender, + ext='npz', + num_betas=numBetas, + use_pca=False, + batch_size=batch_size, + use_face_contour=use_face_contour, + num_pca_comps=45, + **kwargs + ) + for gender in self.genders + }) self.model_neutral = self.model_dict['neutral'] joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES] J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA) @@ -426,9 +427,9 @@ class SMPLX_ALL(nn.Module): # kwargs[key] += self.model_neutral.left_hand_mean # elif key == 'right_hand_pose': # kwargs[key] += self.model_neutral.right_hand_mean - kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view( - [batch_size, -1, 3, 3] - ) + kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([ + batch_size, -1, 3, 3 + ]) if kwargs['body_pose'].shape[1] == 23: # remove hand pose in the body_pose kwargs['body_pose'] = kwargs['body_pose'][:, :21] @@ -570,9 +571,9 @@ class MANO(MANOLayer): if kwargs['pose2rot']: for key in pose_keys: if key in kwargs: - kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view( - [batch_size, -1, 3, 3] - ) + kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([ + batch_size, -1, 3, 3 + ]) kwargs['hand_pose'] = kwargs.pop('right_hand_pose') mano_output = super().forward(*args, **kwargs) th_verts = mano_output.vertices @@ -605,9 +606,9 @@ class FLAME(FLAMELayer): if kwargs['pose2rot']: for key in pose_keys: if key in kwargs: - kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view( - [batch_size, -1, 3, 3] - ) + kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([ + batch_size, -1, 3, 3 + ]) flame_output = super().forward(*args, **kwargs) output = ModelOutput( flame_vertices=flame_output.vertices, @@ -745,9 +746,8 @@ def get_part_joints(smpl_joints): # 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) - ] + 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) diff --git a/lib/pymafx/models/transformers/bert/__init__.py b/lib/pymafx/models/transformers/bert/__init__.py index 0432a1e92856c438e5fd2f550dc5029a78fa354c..d66c20fc9ec44b8fb3ae68a611b784bc624c2616 100644 --- a/lib/pymafx/models/transformers/bert/__init__.py +++ b/lib/pymafx/models/transformers/bert/__init__.py @@ -1,19 +1,23 @@ __version__ = "1.0.0" +from .file_utils import PYTORCH_PRETRAINED_BERT_CACHE, cached_path from .modeling_bert import ( - BertConfig, BertModel, load_tf_weights_in_bert, BERT_PRETRAINED_MODEL_ARCHIVE_MAP, - BERT_PRETRAINED_CONFIG_ARCHIVE_MAP + BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, + BERT_PRETRAINED_MODEL_ARCHIVE_MAP, + BertConfig, + BertModel, + load_tf_weights_in_bert, ) - from .modeling_graphormer import Graphormer - -# from .e2e_body_network import Graphormer_Body_Network - -# from .e2e_hand_network import Graphormer_Hand_Network - from .modeling_utils import ( - WEIGHTS_NAME, CONFIG_NAME, TF_WEIGHTS_NAME, PretrainedConfig, PreTrainedModel, prune_layer, - Conv1D + CONFIG_NAME, + TF_WEIGHTS_NAME, + WEIGHTS_NAME, + Conv1D, + PretrainedConfig, + PreTrainedModel, + prune_layer, ) -from .file_utils import (PYTORCH_PRETRAINED_BERT_CACHE, cached_path) +# from .e2e_body_network import Graphormer_Body_Network +# from .e2e_hand_network import Graphormer_Hand_Network diff --git a/lib/pymafx/models/transformers/bert/e2e_body_network.py b/lib/pymafx/models/transformers/bert/e2e_body_network.py index 9d1c75e276aa18fa1e8f2d865cbef7a275f71b8c..6e00cf8716db068ecdca3d790ecacbdb9e518edc 100644 --- a/lib/pymafx/models/transformers/bert/e2e_body_network.py +++ b/lib/pymafx/models/transformers/bert/e2e_body_network.py @@ -4,8 +4,8 @@ Licensed under the MIT license. """ -import torch import src.modeling.data.config as cfg +import torch class Graphormer_Body_Network(torch.nn.Module): diff --git a/lib/pymafx/models/transformers/bert/e2e_hand_network.py b/lib/pymafx/models/transformers/bert/e2e_hand_network.py index 410968c4abc63e1ae8281b2e0297c8eef4e7bbcf..be88fc5f32099daf9ca56eaa8d362efd5d8915b6 100644 --- a/lib/pymafx/models/transformers/bert/e2e_hand_network.py +++ b/lib/pymafx/models/transformers/bert/e2e_hand_network.py @@ -4,8 +4,8 @@ Licensed under the MIT license. """ -import torch import src.modeling.data.config as cfg +import torch class Graphormer_Hand_Network(torch.nn.Module): diff --git a/lib/pymafx/models/transformers/bert/file_utils.py b/lib/pymafx/models/transformers/bert/file_utils.py index ee58bed427f90be254caee9a0733d81ae92c8711..2159d8ea87e216039104054f70a159f8c534179b 100644 --- a/lib/pymafx/models/transformers/bert/file_utils.py +++ b/lib/pymafx/models/transformers/bert/file_utils.py @@ -3,15 +3,20 @@ Utilities for working with the local dataset cache. This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp Copyright by the AllenNLP authors. """ -from __future__ import (absolute_import, division, print_function, unicode_literals) +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) -import sys +import fnmatch import json import logging import os import shutil +import sys import tempfile -import fnmatch from functools import wraps from hashlib import sha256 from io import open diff --git a/lib/pymafx/models/transformers/bert/modeling_bert.py b/lib/pymafx/models/transformers/bert/modeling_bert.py index c4a7f27f1bc0e69d87ac3747b8d8acfafb03b4b8..c922f5229d9b85f4eeecade118264e18e6e56753 100644 --- a/lib/pymafx/models/transformers/bert/modeling_bert.py +++ b/lib/pymafx/models/transformers/bert/modeling_bert.py @@ -15,7 +15,12 @@ # limitations under the License. """PyTorch BERT model. """ -from __future__ import absolute_import, division, print_function, unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) import json import logging @@ -29,68 +34,72 @@ from torch import nn from torch.nn import CrossEntropyLoss, MSELoss from .modeling_utils import ( - WEIGHTS_NAME, CONFIG_NAME, PretrainedConfig, PreTrainedModel, prune_linear_layer, - add_start_docstrings + CONFIG_NAME, + WEIGHTS_NAME, + PretrainedConfig, + PreTrainedModel, + add_start_docstrings, + prune_linear_layer, ) logger = logging.getLogger(__name__) BERT_PRETRAINED_MODEL_ARCHIVE_MAP = { 'bert-base-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin", 'bert-large-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch_model.bin", 'bert-base-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch_model.bin", 'bert-large-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch_model.bin", 'bert-base-multilingual-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch_model.bin", 'bert-base-multilingual-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch_model.bin", 'bert-base-chinese': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch_model.bin", 'bert-base-german-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-pytorch_model.bin", 'bert-large-uncased-whole-word-masking': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-pytorch_model.bin", 'bert-large-cased-whole-word-masking': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-pytorch_model.bin", 'bert-large-uncased-whole-word-masking-finetuned-squad': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-pytorch_model.bin", 'bert-large-cased-whole-word-masking-finetuned-squad': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-pytorch_model.bin", 'bert-base-cased-finetuned-mrpc': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-pytorch_model.bin", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-pytorch_model.bin", } BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = { 'bert-base-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json", 'bert-large-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json", 'bert-base-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json", 'bert-large-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json", 'bert-base-multilingual-uncased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json", 'bert-base-multilingual-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json", 'bert-base-chinese': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json", 'bert-base-german-cased': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-config.json", 'bert-large-uncased-whole-word-masking': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json", 'bert-large-cased-whole-word-masking': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json", 'bert-large-uncased-whole-word-masking-finetuned-squad': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json", 'bert-large-cased-whole-word-masking-finetuned-squad': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json", 'bert-base-cased-finetuned-mrpc': - "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json", + "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json", } @@ -99,6 +108,7 @@ def load_tf_weights_in_bert(model, config, tf_checkpoint_path): """ try: import re + import numpy as np import tensorflow as tf except ImportError: diff --git a/lib/pymafx/models/transformers/bert/modeling_graphormer.py b/lib/pymafx/models/transformers/bert/modeling_graphormer.py index e318af8a45d34148e0db68f42181f692afbf8754..315e2e6fe0b3f4b4a35add9dd9fe3d4e2b2a53c4 100644 --- a/lib/pymafx/models/transformers/bert/modeling_graphormer.py +++ b/lib/pymafx/models/transformers/bert/modeling_graphormer.py @@ -4,15 +4,29 @@ Licensed under the MIT license. """ -from __future__ import absolute_import, division, print_function, unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) +import code import logging import math import os -import code + import torch from torch import nn -from .modeling_bert import BertPreTrainedModel, BertEmbeddings, BertPooler, BertIntermediate, BertOutput, BertSelfOutput + +from .modeling_bert import ( + BertEmbeddings, + BertIntermediate, + BertOutput, + BertPooler, + BertPreTrainedModel, + BertSelfOutput, +) # import src.modeling.data.config as cfg # from src.modeling._gcnn import GraphConvolution, GraphResBlock from .modeling_utils import prune_linear_layer @@ -180,9 +194,9 @@ class GraphormerEncoder(nn.Module): super(GraphormerEncoder, self).__init__() self.output_attentions = config.output_attentions self.output_hidden_states = config.output_hidden_states - self.layer = nn.ModuleList( - [GraphormerLayer(config) for _ in range(config.num_hidden_layers)] - ) + self.layer = nn.ModuleList([ + GraphormerLayer(config) for _ in range(config.num_hidden_layers) + ]) def forward(self, hidden_states, attention_mask, head_mask=None, encoder_history_states=None): all_hidden_states = () diff --git a/lib/pymafx/models/transformers/bert/modeling_utils.py b/lib/pymafx/models/transformers/bert/modeling_utils.py index 40a0915822c8e736de8ac2466c075e6cc5ef7e83..56cb54703a5ca9b068933a92440fce1c278ea01e 100644 --- a/lib/pymafx/models/transformers/bert/modeling_utils.py +++ b/lib/pymafx/models/transformers/bert/modeling_utils.py @@ -15,7 +15,12 @@ # limitations under the License. """PyTorch BERT model.""" -from __future__ import (absolute_import, division, print_function, unicode_literals) +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) import copy import json @@ -552,9 +557,8 @@ class PreTrainedModel(nn.Module): if output_loading_info: loading_info = { - "missing_keys": missing_keys, - "unexpected_keys": unexpected_keys, - "error_msgs": error_msgs + "missing_keys": missing_keys, "unexpected_keys": unexpected_keys, "error_msgs": + error_msgs } return model, loading_info @@ -893,9 +897,8 @@ class SequenceSummary(nn.Module): ) else: token_ids = token_ids.unsqueeze(-1).unsqueeze(-1) - token_ids = token_ids.expand( - (-1, ) * (token_ids.dim() - 1) + (hidden_states.size(-1), ) - ) + token_ids = token_ids.expand((-1, ) * (token_ids.dim() - 1) + + (hidden_states.size(-1), )) # shape of token_ids: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states output = hidden_states.gather(-2, token_ids).squeeze(-2) # shape (bsz, XX, hidden_size) diff --git a/lib/pymafx/models/transformers/net_utils.py b/lib/pymafx/models/transformers/net_utils.py index 52782911e276705ec0dd908ce9676430c0a58d72..ff59e7ebdf0756306352ccd1ebc2827b4e032c90 100644 --- a/lib/pymafx/models/transformers/net_utils.py +++ b/lib/pymafx/models/transformers/net_utils.py @@ -1,6 +1,7 @@ -import torch.nn as nn -import torch import math + +import torch +import torch.nn as nn import torch.nn.functional as F diff --git a/lib/pymafx/models/transformers/texformer.py b/lib/pymafx/models/transformers/texformer.py index 4266b24ed6839f91ce5ca819cc3750143387d48f..aada36cd9fad64a8ff47743da40155ed6b012b38 100644 --- a/lib/pymafx/models/transformers/texformer.py +++ b/lib/pymafx/models/transformers/texformer.py @@ -1,5 +1,12 @@ import torch.nn as nn -from .net_utils import single_conv, double_conv, double_conv_down, double_conv_up, PosEnSine + +from .net_utils import ( + PosEnSine, + double_conv, + double_conv_down, + double_conv_up, + single_conv, +) from .transformer_basics import OurMultiheadAttention @@ -86,14 +93,12 @@ class Texformer(nn.Module): self.unet_k = Unet(src_ch, self.feat_dim, self.feat_dim) self.unet_v = Unet(v_ch, self.feat_dim, self.feat_dim) - self.trans_dec = nn.ModuleList( - [ - None, None, None, - TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'softmax'), - TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'dotproduct'), - TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'dotproduct') - ] - ) + self.trans_dec = nn.ModuleList([ + None, None, None, + TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'softmax'), + TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'dotproduct'), + TransformerDecoderUnit(self.feat_dim, opts.nhead, True, 'dotproduct') + ]) self.conv0 = double_conv(self.feat_dim, self.feat_dim) self.conv1 = double_conv_down(self.feat_dim, self.feat_dim) diff --git a/lib/pymafx/models/transformers/transformer_basics.py b/lib/pymafx/models/transformers/transformer_basics.py index 144ccd76b7e2f73189634ab551691c4262781b9d..f2c9d9533926b88d0308c18560b9cf327e9d317a 100644 --- a/lib/pymafx/models/transformers/transformer_basics.py +++ b/lib/pymafx/models/transformers/transformer_basics.py @@ -1,6 +1,13 @@ import torch.nn as nn -from .net_utils import PosEnSine, softmax_attention, dotproduct_attention, long_range_attention, \ - short_range_attention, patch_attention + +from .net_utils import ( + PosEnSine, + dotproduct_attention, + long_range_attention, + patch_attention, + short_range_attention, + softmax_attention, +) class OurMultiheadAttention(nn.Module): diff --git a/lib/pymafx/utils/blob.py b/lib/pymafx/utils/blob.py index 00123338e18a3fa74a6c3cb730cac9fb41b59ac5..11814bbec48887f622d11a786ab25271f98d5450 100644 --- a/lib/pymafx/utils/blob.py +++ b/lib/pymafx/utils/blob.py @@ -22,16 +22,17 @@ # -------------------------------------------------------- """blob helper functions.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) -from six.moves import cPickle as pickle -import numpy as np import cv2 - +import numpy as np from models.core.config import cfg +from six.moves import cPickle as pickle def get_image_blob(im, target_scale, target_max_size): diff --git a/lib/pymafx/utils/cam_params.py b/lib/pymafx/utils/cam_params.py index 1f6c1a8d89b2c80d72c90c841d02425df77aa4a5..f8138bede445a95571e0b179f4f8515a7a2cb672 100644 --- a/lib/pymafx/utils/cam_params.py +++ b/lib/pymafx/utils/cam_params.py @@ -15,10 +15,11 @@ # Contact: ps-license@tuebingen.mpg.de import os -from numpy.testing._private.utils import print_assert_equal -import torch -import numpy as np + import joblib +import numpy as np +import torch +from numpy.testing._private.utils import print_assert_equal from .geometry import batch_euler2matrix diff --git a/lib/pymafx/utils/collections.py b/lib/pymafx/utils/collections.py index edd20a8c89d5d2221dc9d35948eda12c6304ba29..b7288875cd49a06d8305f95a104e0309907d9ec0 100644 --- a/lib/pymafx/utils/collections.py +++ b/lib/pymafx/utils/collections.py @@ -14,10 +14,12 @@ ############################################################################## """A simple attribute dictionary used for representing configuration options.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) class AttrDict(dict): diff --git a/lib/pymafx/utils/colormap.py b/lib/pymafx/utils/colormap.py index 44ef28c050021a6f03d088e9437de0c4adeb5ee5..1b275cf20b8c307564a2469d77426fe3ac5d2996 100644 --- a/lib/pymafx/utils/colormap.py +++ b/lib/pymafx/utils/colormap.py @@ -14,39 +14,38 @@ ############################################################################## """An awesome colormap for really neat visualizations.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) import numpy as np def colormap(rgb=False): - color_list = np.array( - [ - 0.000, 0.447, 0.741, 0.850, 0.325, 0.098, 0.929, 0.694, 0.125, 0.494, 0.184, 0.556, - 0.466, 0.674, 0.188, 0.301, 0.745, 0.933, 0.635, 0.078, 0.184, 0.300, 0.300, 0.300, - 0.600, 0.600, 0.600, 1.000, 0.000, 0.000, 1.000, 0.500, 0.000, 0.749, 0.749, 0.000, - 0.000, 1.000, 0.000, 0.000, 0.000, 1.000, 0.667, 0.000, 1.000, 0.333, 0.333, 0.000, - 0.333, 0.667, 0.000, 0.333, 1.000, 0.000, 0.667, 0.333, 0.000, 0.667, 0.667, 0.000, - 0.667, 1.000, 0.000, 1.000, 0.333, 0.000, 1.000, 0.667, 0.000, 1.000, 1.000, 0.000, - 0.000, 0.333, 0.500, 0.000, 0.667, 0.500, 0.000, 1.000, 0.500, 0.333, 0.000, 0.500, - 0.333, 0.333, 0.500, 0.333, 0.667, 0.500, 0.333, 1.000, 0.500, 0.667, 0.000, 0.500, - 0.667, 0.333, 0.500, 0.667, 0.667, 0.500, 0.667, 1.000, 0.500, 1.000, 0.000, 0.500, - 1.000, 0.333, 0.500, 1.000, 0.667, 0.500, 1.000, 1.000, 0.500, 0.000, 0.333, 1.000, - 0.000, 0.667, 1.000, 0.000, 1.000, 1.000, 0.333, 0.000, 1.000, 0.333, 0.333, 1.000, - 0.333, 0.667, 1.000, 0.333, 1.000, 1.000, 0.667, 0.000, 1.000, 0.667, 0.333, 1.000, - 0.667, 0.667, 1.000, 0.667, 1.000, 1.000, 1.000, 0.000, 1.000, 1.000, 0.333, 1.000, - 1.000, 0.667, 1.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, - 0.667, 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, 0.167, 0.000, - 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 0.000, 0.000, 0.833, 0.000, - 0.000, 1.000, 0.000, 0.000, 0.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, - 0.000, 0.000, 0.667, 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, - 0.143, 0.143, 0.143, 0.286, 0.286, 0.286, 0.429, 0.429, 0.429, 0.571, 0.571, 0.571, - 0.714, 0.714, 0.714, 0.857, 0.857, 0.857, 1.000, 1.000, 1.000 - ] - ).astype(np.float32) + color_list = np.array([ + 0.000, 0.447, 0.741, 0.850, 0.325, 0.098, 0.929, 0.694, 0.125, 0.494, 0.184, 0.556, 0.466, + 0.674, 0.188, 0.301, 0.745, 0.933, 0.635, 0.078, 0.184, 0.300, 0.300, 0.300, 0.600, 0.600, + 0.600, 1.000, 0.000, 0.000, 1.000, 0.500, 0.000, 0.749, 0.749, 0.000, 0.000, 1.000, 0.000, + 0.000, 0.000, 1.000, 0.667, 0.000, 1.000, 0.333, 0.333, 0.000, 0.333, 0.667, 0.000, 0.333, + 1.000, 0.000, 0.667, 0.333, 0.000, 0.667, 0.667, 0.000, 0.667, 1.000, 0.000, 1.000, 0.333, + 0.000, 1.000, 0.667, 0.000, 1.000, 1.000, 0.000, 0.000, 0.333, 0.500, 0.000, 0.667, 0.500, + 0.000, 1.000, 0.500, 0.333, 0.000, 0.500, 0.333, 0.333, 0.500, 0.333, 0.667, 0.500, 0.333, + 1.000, 0.500, 0.667, 0.000, 0.500, 0.667, 0.333, 0.500, 0.667, 0.667, 0.500, 0.667, 1.000, + 0.500, 1.000, 0.000, 0.500, 1.000, 0.333, 0.500, 1.000, 0.667, 0.500, 1.000, 1.000, 0.500, + 0.000, 0.333, 1.000, 0.000, 0.667, 1.000, 0.000, 1.000, 1.000, 0.333, 0.000, 1.000, 0.333, + 0.333, 1.000, 0.333, 0.667, 1.000, 0.333, 1.000, 1.000, 0.667, 0.000, 1.000, 0.667, 0.333, + 1.000, 0.667, 0.667, 1.000, 0.667, 1.000, 1.000, 1.000, 0.000, 1.000, 1.000, 0.333, 1.000, + 1.000, 0.667, 1.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, + 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, 0.167, 0.000, 0.000, 0.333, + 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 0.000, 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, + 0.000, 0.000, 0.167, 0.000, 0.000, 0.333, 0.000, 0.000, 0.500, 0.000, 0.000, 0.667, 0.000, + 0.000, 0.833, 0.000, 0.000, 1.000, 0.000, 0.000, 0.000, 0.143, 0.143, 0.143, 0.286, 0.286, + 0.286, 0.429, 0.429, 0.429, 0.571, 0.571, 0.571, 0.714, 0.714, 0.714, 0.857, 0.857, 0.857, + 1.000, 1.000, 1.000 + ]).astype(np.float32) color_list = color_list.reshape((-1, 3)) * 255 if not rgb: color_list = color_list[:, ::-1] diff --git a/lib/pymafx/utils/common.py b/lib/pymafx/utils/common.py index f3330ea18c4783ccacb21657808b8b8ce2301f86..bdcb5d0e3b18ab60f88bcab94de75922513e6195 100755 --- a/lib/pymafx/utils/common.py +++ b/lib/pymafx/utils/common.py @@ -1,7 +1,9 @@ -import torch -import numpy as np import logging from copy import deepcopy + +import numpy as np +import torch + from .utils.libkdtree import KDTree logger_py = logging.getLogger(__name__) @@ -170,14 +172,12 @@ def check_ray_intersection_with_unit_cube(ray0, ray_direction, padding=0.1, eps= ray_direction.unsqueeze(-2) # Calculate mask where points intersect unit cube - p_mask_inside_cube = ( - (p_intersect[:, :, :, 0] <= p_distance + eps) & - (p_intersect[:, :, :, 1] <= p_distance + eps) & - (p_intersect[:, :, :, 2] <= p_distance + eps) & - (p_intersect[:, :, :, 0] >= -(p_distance + eps)) & - (p_intersect[:, :, :, 1] >= -(p_distance + eps)) & - (p_intersect[:, :, :, 2] >= -(p_distance + eps)) - ).cpu() + p_mask_inside_cube = ((p_intersect[:, :, :, 0] <= p_distance + eps) & + (p_intersect[:, :, :, 1] <= p_distance + eps) & + (p_intersect[:, :, :, 2] <= p_distance + eps) & + (p_intersect[:, :, :, 0] >= -(p_distance + eps)) & + (p_intersect[:, :, :, 1] >= -(p_distance + eps)) & + (p_intersect[:, :, :, 2] >= -(p_distance + eps))).cpu() # Correct rays are these which intersect exactly 2 times mask_inside_cube = p_mask_inside_cube.sum(-1) == 2 @@ -190,13 +190,11 @@ def check_ray_intersection_with_unit_cube(ray0, ray_direction, padding=0.1, eps= # Calculate ray lengths for the interval points d_intervals_batch = torch.zeros(batch_size, n_pts, 2).to(device) norm_ray = torch.norm(ray_direction[mask_inside_cube], dim=-1) - d_intervals_batch[mask_inside_cube] = torch.stack( - [ - torch.norm(p_intervals[:, 0] - ray0[mask_inside_cube], dim=-1) / norm_ray, - torch.norm(p_intervals[:, 1] - ray0[mask_inside_cube], dim=-1) / norm_ray, - ], - dim=-1 - ) + d_intervals_batch[mask_inside_cube] = torch.stack([ + torch.norm(p_intervals[:, 0] - ray0[mask_inside_cube], dim=-1) / norm_ray, + torch.norm(p_intervals[:, 1] - ray0[mask_inside_cube], dim=-1) / norm_ray, + ], + dim=-1) # Sort the ray lengths d_intervals_batch, indices_sort = d_intervals_batch.sort() diff --git a/lib/pymafx/utils/data_loader.py b/lib/pymafx/utils/data_loader.py index cc92ad223836e9de322bc80bbab887bb9ec3f17b..3d109f82b3473242a9fb9442037c47471fd0f7d2 100644 --- a/lib/pymafx/utils/data_loader.py +++ b/lib/pymafx/utils/data_loader.py @@ -1,4 +1,5 @@ from __future__ import division + import torch from torch.utils.data import DataLoader from torch.utils.data.sampler import Sampler diff --git a/lib/pymafx/utils/demo_utils.py b/lib/pymafx/utils/demo_utils.py index b1ad8da91c7a7f6f67d4770c9866a02a78aa5275..437a2b42cae65f51dc357ffffef0cc63d6aa58b9 100644 --- a/lib/pymafx/utils/demo_utils.py +++ b/lib/pymafx/utils/demo_utils.py @@ -14,20 +14,20 @@ # # Contact: ps-license@tuebingen.mpg.de -import os -import cv2 -import time import json -import torch -import subprocess -import numpy as np +import os import os.path as osp +import subprocess +import time # from pytube import YouTube from collections import OrderedDict -from utils.smooth_bbox import get_smooth_bbox_params, get_all_bbox_params +import cv2 +import numpy as np +import torch from datasets.data_utils.img_utils import get_single_image_crop_demo from utils.geometry import rotation_matrix_to_angle_axis +from utils.smooth_bbox import get_all_bbox_params, get_smooth_bbox_params def preprocess_video(video, joints2d, bboxes, frames, scale=1.0, crop_size=224): @@ -112,9 +112,10 @@ def smplify_runner( pred_pose = pred_rotmat # Calculate camera parameters for smplify - pred_cam_t = torch.stack( - [pred_cam[:, 1], pred_cam[:, 2], 2 * 5000 / (224 * pred_cam[:, 0] + 1e-9)], dim=-1 - ) + pred_cam_t = torch.stack([ + pred_cam[:, 1], pred_cam[:, 2], 2 * 5000 / (224 * pred_cam[:, 0] + 1e-9) + ], + dim=-1) gt_keypoints_2d_orig = j2d # Before running compute reprojection error of the network @@ -285,14 +286,11 @@ def prepare_rendering_results(results_dict, nframes): for person_id, person_data in results_dict.items(): for idx, frame_id in enumerate(person_data['frame_ids']): frame_results[frame_id][person_id] = { - 'verts': - person_data['verts'][idx], + 'verts': person_data['verts'][idx], 'smplx_verts': - person_data['smplx_verts'][idx] if 'smplx_verts' in person_data else None, - 'cam': - person_data['orig_cam'][idx], - 'cam_t': - person_data['orig_cam_t'][idx] if 'orig_cam_t' in person_data else None, + person_data['smplx_verts'][idx] if 'smplx_verts' in person_data else None, + 'cam': person_data['orig_cam'][idx], + 'cam_t': person_data['orig_cam_t'][idx] if 'orig_cam_t' in person_data else None, # 'cam': person_data['pred_cam'][idx], } @@ -300,9 +298,9 @@ def prepare_rendering_results(results_dict, nframes): for frame_id, frame_data in enumerate(frame_results): # sort based on y-scale of the cam in original image coords sort_idx = np.argsort([v['cam'][1] for k, v in frame_data.items()]) - frame_results[frame_id] = OrderedDict( - {list(frame_data.keys())[i]: frame_data[list(frame_data.keys())[i]] - for i in sort_idx} - ) + frame_results[frame_id] = OrderedDict({ + list(frame_data.keys())[i]: frame_data[list(frame_data.keys())[i]] + for i in sort_idx + }) return frame_results diff --git a/lib/pymafx/utils/densepose_methods.py b/lib/pymafx/utils/densepose_methods.py index 93fdf66a6651dcfe05f6e95c55379eaa00c52cb0..7e02eafd746a4a152cbb3d3569ed3fdaf673f7fe 100644 --- a/lib/pymafx/utils/densepose_methods.py +++ b/lib/pymafx/utils/densepose_methods.py @@ -4,12 +4,13 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. -import numpy as np import copy +import os + import cv2 -from scipy.io import loadmat +import numpy as np import scipy.spatial.distance -import os +from scipy.io import loadmat class DensePoseMethods: @@ -102,24 +103,18 @@ class DensePoseMethods: FaceIndicesNow = np.where(self.FaceIndices == I_point) FacesNow = self.FacesDensePose[FaceIndicesNow] # - P_0 = np.vstack( - ( - self.U_norm[FacesNow][:, 0], self.V_norm[FacesNow][:, 0], - np.zeros(self.U_norm[FacesNow][:, 0].shape) - ) - ).transpose() - P_1 = np.vstack( - ( - self.U_norm[FacesNow][:, 1], self.V_norm[FacesNow][:, 1], - np.zeros(self.U_norm[FacesNow][:, 1].shape) - ) - ).transpose() - P_2 = np.vstack( - ( - self.U_norm[FacesNow][:, 2], self.V_norm[FacesNow][:, 2], - np.zeros(self.U_norm[FacesNow][:, 2].shape) - ) - ).transpose() + P_0 = np.vstack(( + self.U_norm[FacesNow][:, 0], self.V_norm[FacesNow][:, 0], + np.zeros(self.U_norm[FacesNow][:, 0].shape) + )).transpose() + P_1 = np.vstack(( + self.U_norm[FacesNow][:, 1], self.V_norm[FacesNow][:, 1], + np.zeros(self.U_norm[FacesNow][:, 1].shape) + )).transpose() + P_2 = np.vstack(( + self.U_norm[FacesNow][:, 2], self.V_norm[FacesNow][:, 2], + np.zeros(self.U_norm[FacesNow][:, 2].shape) + )).transpose() # for i, [P0, P1, P2] in enumerate(zip(P_0, P_1, P_2)): diff --git a/lib/pymafx/utils/geometry.py b/lib/pymafx/utils/geometry.py index 608288fc4d73a4918ab95938a7bf5dbe98ce606f..5c2d2c1f8b22147344afc549ae74630bea56d3ef 100644 --- a/lib/pymafx/utils/geometry.py +++ b/lib/pymafx/utils/geometry.py @@ -1,8 +1,10 @@ -import torch -from torch.nn import functional as F -import numpy as np import numbers + +import numpy as np +import torch from einops.einops import rearrange +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 @@ -43,13 +45,11 @@ def quat_to_rotmat(quat): 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) + 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 @@ -225,39 +225,31 @@ def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6): 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 - ) + 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 - ) + 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 - ) + 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 - ) + 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 @@ -321,13 +313,11 @@ def quaternion_to_rotation_matrix(quat): 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) + 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 @@ -405,9 +395,10 @@ def projection(pred_joints, pred_camera, retain_z=False, iwp_mode=True): batch_size = pred_joints.shape[0] if iwp_mode: cam_sxy = pred_camera['cam_sxy'] - pred_cam_t = torch.stack( - [cam_sxy[:, 1], cam_sxy[:, 2], 2 * 5000. / (224. * cam_sxy[:, 0] + 1e-9)], dim=-1 - ) + pred_cam_t = torch.stack([ + cam_sxy[:, 1], cam_sxy[:, 2], 2 * 5000. / (224. * cam_sxy[:, 0] + 1e-9) + ], + dim=-1) camera_center = torch.zeros(batch_size, 2) pred_keypoints_2d = perspective_projection( @@ -537,12 +528,10 @@ def estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000, img_si 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 + 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 @@ -609,10 +598,8 @@ def Rot_y(angle, category='torch', prepend_dim=True, device=None): 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)]] - ) + 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) @@ -634,10 +621,8 @@ def Rot_x(angle, category='torch', prepend_dim=True, device=None): 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)]] - ) + 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) @@ -659,9 +644,8 @@ def Rot_z(angle, category='torch', prepend_dim=True, device=None): 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.]] - ) + 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) diff --git a/lib/pymafx/utils/imutils.py b/lib/pymafx/utils/imutils.py index b3522fee118cf47c5101bfd8e16991e5c30f58ad..f3df8bf7ccf43a2f30074f52ccbc025017af66cf 100644 --- a/lib/pymafx/utils/imutils.py +++ b/lib/pymafx/utils/imutils.py @@ -1,10 +1,10 @@ """ This file contains functions that are used to perform data augmentation. """ -import torch -import numpy as np import cv2 +import numpy as np import skimage.transform +import torch from PIL import Image from lib.pymafx.core import constants @@ -129,12 +129,8 @@ def uncrop(img, center, scale, orig_shape, rot=0, is_rgb=True): 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] - ] - ) + 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 @@ -246,9 +242,9 @@ def generate_heatmap(joints, heatmap_size, sigma=1, joints_vis=None): 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 - ) + target = torch.zeros((num_joints, heatmap_size[1], heatmap_size[0]), + dtype=torch.float32, + device=cur_device) tmp_size = sigma * 3 diff --git a/lib/pymafx/utils/io.py b/lib/pymafx/utils/io.py index 0926624ddeb1eccf2e9c6393595acfd34a62e84d..67d5b50542ef8f831ac59d46947631ae8f2cc78e 100644 --- a/lib/pymafx/utils/io.py +++ b/lib/pymafx/utils/io.py @@ -14,17 +14,21 @@ ############################################################################## """IO utilities.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) -from six.moves import cPickle as pickle import hashlib import logging import os import re import sys + +from six.moves import cPickle as pickle + try: from urllib.request import urlopen except ImportError: #python2 diff --git a/lib/pymafx/utils/iuvmap.py b/lib/pymafx/utils/iuvmap.py index 7f7c25398e04e30b2b244d44badc83415d583852..8c1914dac03f553f7651fac74ca3bd1b45b65645 100644 --- a/lib/pymafx/utils/iuvmap.py +++ b/lib/pymafx/utils/iuvmap.py @@ -115,10 +115,8 @@ def iuv_img2map(uvimages, uv_rois=None, new_size=None, n_part=24): batch_size = uvimages.size(0) uvimg_size = uvimages.size(-1) - Index2mask = [ - [0], [1, 2], [3], [4], [5], [6], [7, 9], [8, 10], [11, 13], [12, 14], [15, 17], [16, 18], - [19, 21], [20, 22], [23, 24] - ] + Index2mask = [[0], [1, 2], [3], [4], [5], [6], [7, 9], [8, 10], [11, 13], [12, 14], [15, 17], + [16, 18], [19, 21], [20, 22], [23, 24]] part_ind = torch.round(uvimages[:, 0, :, :] * n_part) part_u = uvimages[:, 1, :, :] diff --git a/lib/pymafx/utils/keypoints.py b/lib/pymafx/utils/keypoints.py index 2ab223c2bef79518adc523da1606cfc331ef8251..a4e52d6a51dd24624be07a469fef62e1d0130995 100644 --- a/lib/pymafx/utils/keypoints.py +++ b/lib/pymafx/utils/keypoints.py @@ -14,16 +14,18 @@ ############################################################################## """Keypoint utilities (somewhat specific to COCO keypoints).""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) import cv2 import numpy as np import torch -import torch.nn.functional as F import torch.cuda.comm +import torch.nn.functional as F # from core.config import cfg # import utils.blob as blob_utils @@ -39,14 +41,9 @@ def get_keypoints(): 'left_knee', 'right_knee', 'left_ankle', 'right_ankle' ] keypoint_flip_map = { - 'left_eye': 'right_eye', - 'left_ear': 'right_ear', - 'left_shoulder': 'right_shoulder', - 'left_elbow': 'right_elbow', - 'left_wrist': 'right_wrist', - 'left_hip': 'right_hip', - 'left_knee': 'right_knee', - 'left_ankle': 'right_ankle' + 'left_eye': 'right_eye', 'left_ear': 'right_ear', 'left_shoulder': 'right_shoulder', + 'left_elbow': 'right_elbow', 'left_wrist': 'right_wrist', 'left_hip': 'right_hip', + 'left_knee': 'right_knee', 'left_ankle': 'right_ankle' } return keypoints, keypoint_flip_map @@ -232,9 +229,9 @@ def compute_oks(src_keypoints, src_roi, dst_keypoints, dst_roi): dst_roi: Nx4 """ - sigmas = np.array( - [.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89] - ) / 10.0 + sigmas = np.array([ + .26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89 + ]) / 10.0 vars = (sigmas * 2)**2 # area diff --git a/lib/pymafx/utils/mesh_generation.py b/lib/pymafx/utils/mesh_generation.py index 2876209e7678d2906a84850208f6c288103d07c5..3442d16a79e7594a17c1a3ca7ae195ec5c463b4c 100644 --- a/lib/pymafx/utils/mesh_generation.py +++ b/lib/pymafx/utils/mesh_generation.py @@ -1,16 +1,16 @@ import time -import torch -import trimesh + import numpy as np +import torch import torch.optim as optim +import trimesh from torch import autograd -from torch.utils.data import TensorDataset, DataLoader +from torch.utils.data import DataLoader, TensorDataset -from .common import make_3d_grid +from .common import make_3d_grid, transform_pointcloud from .utils import libmcubes from .utils.libmise import MISE from .utils.libsimplify import simplify_mesh -from .common import transform_pointcloud class Generator3D(object): @@ -286,9 +286,8 @@ class Generator3D(object): colors = np.concatenate(colors, axis=0) colors = np.clip(colors, 0, 1) colors = (colors * 255).astype(np.uint8) - colors = np.concatenate( - [colors, np.full((colors.shape[0], 1), 255, dtype=np.uint8)], axis=1 - ) + colors = np.concatenate([colors, np.full((colors.shape[0], 1), 255, dtype=np.uint8)], + axis=1) return colors def estimate_normals(self, vertices, c=None): @@ -375,13 +374,11 @@ class Generator3D(object): face_normal = face_normal / \ (face_normal.norm(dim=1, keepdim=True) + 1e-10) - face_value = torch.cat( - [ - torch.sigmoid(self.model.decode(p_split, c).logits) - for p_split in torch.split(face_point.unsqueeze(0), 20000, dim=1) - ], - dim=1 - ) + face_value = torch.cat([ + torch.sigmoid(self.model.decode(p_split, c).logits) + for p_split in torch.split(face_point.unsqueeze(0), 20000, dim=1) + ], + dim=1) normal_target = -autograd.grad([face_value.sum()], [face_point], create_graph=True)[0] diff --git a/lib/pymafx/utils/part_utils.py b/lib/pymafx/utils/part_utils.py index 12f0de443fa11e90674761816a644cf82a48a786..5b17763404418e9b97d738af4c497be2d1918cde 100644 --- a/lib/pymafx/utils/part_utils.py +++ b/lib/pymafx/utils/part_utils.py @@ -1,8 +1,7 @@ -import torch -import numpy as np import neural_renderer as nr +import numpy as np +import torch from core import path_config - from models import SMPL @@ -42,13 +41,11 @@ class PartRenderer(): def __call__(self, vertices, camera): """Wrapper function for rendering process.""" # Estimate camera parameters given a fixed focal length - cam_t = torch.stack( - [ - camera[:, 1], camera[:, 2], 2 * self.focal_length / - (self.render_res * camera[:, 0] + 1e-9) - ], - dim=-1 - ) + cam_t = torch.stack([ + camera[:, 1], camera[:, 2], 2 * self.focal_length / + (self.render_res * camera[:, 0] + 1e-9) + ], + dim=-1) batch_size = vertices.shape[0] K = torch.eye(3, device=vertices.device) K[0, 0] = self.focal_length diff --git a/lib/pymafx/utils/pose_tracker.py b/lib/pymafx/utils/pose_tracker.py index 92c383cdb3dba6053a0595b9f03305c02e9fc277..be6008b31580d600533921ed722a89568c3a36fd 100644 --- a/lib/pymafx/utils/pose_tracker.py +++ b/lib/pymafx/utils/pose_tracker.py @@ -14,12 +14,13 @@ # # Contact: ps-license@tuebingen.mpg.de -import os import json +import os +import os.path as osp import shutil import subprocess + import numpy as np -import os.path as osp def run_openpose( diff --git a/lib/pymafx/utils/pose_utils.py b/lib/pymafx/utils/pose_utils.py index 55eb1d771376da71c864a715d1dd6b5d66e9894e..966eb51850b6fa793b1e11d4dcad1eed2c6698da 100644 --- a/lib/pymafx/utils/pose_utils.py +++ b/lib/pymafx/utils/pose_utils.py @@ -1,9 +1,8 @@ """ Parts of the code are adapted from https://github.com/akanazawa/hmr """ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function +from __future__ import absolute_import, division, print_function + import numpy as np import torch diff --git a/lib/pymafx/utils/renderer.py b/lib/pymafx/utils/renderer.py index 9fb19568680b839f93c00a5288c94a5a52025242..5b201badc2c6da232d0d753c9303b7c62361ca7f 100644 --- a/lib/pymafx/utils/renderer.py +++ b/lib/pymafx/utils/renderer.py @@ -1,45 +1,59 @@ import imp +import json import os from pickle import NONE + +import numpy as np # os.environ['PYOPENGL_PLATFORM'] = 'osmesa' import torch +import torch.nn.functional as F import trimesh -import numpy as np +from core import constants, path_config +from models.smpl import get_model_faces, get_model_tpose, get_smpl_faces # import neural_renderer as nr from skimage.transform import resize from torchvision.utils import make_grid -import torch.nn.functional as F - -from models.smpl import get_smpl_faces, get_model_faces, get_model_tpose from utils.densepose_methods import DensePoseMethods -from core import constants, path_config -import json -from .geometry import convert_to_full_img_cam from utils.imutils import crop +from .geometry import convert_to_full_img_cam + try: import math + import pyrender from pyrender.constants import RenderFlags except: pass try: - from opendr.renderer import ColoredRenderer - from opendr.lighting import LambertianPointLight, SphericalHarmonics from opendr.camera import ProjectPoints + from opendr.lighting import LambertianPointLight, SphericalHarmonics + from opendr.renderer import ColoredRenderer except: pass -from pytorch3d.structures.meshes import Meshes -# from pytorch3d.renderer.mesh.renderer import MeshRendererWithFragments +import logging from pytorch3d.renderer import ( - look_at_view_transform, FoVPerspectiveCameras, PerspectiveCameras, AmbientLights, PointLights, - RasterizationSettings, BlendParams, MeshRenderer, MeshRasterizer, SoftPhongShader, - SoftSilhouetteShader, HardPhongShader, HardGouraudShader, HardFlatShader, TexturesVertex + AmbientLights, + BlendParams, + FoVPerspectiveCameras, + HardFlatShader, + HardGouraudShader, + HardPhongShader, + MeshRasterizer, + MeshRenderer, + PerspectiveCameras, + PointLights, + RasterizationSettings, + SoftPhongShader, + SoftSilhouetteShader, + TexturesVertex, + look_at_view_transform, ) +from pytorch3d.structures.meshes import Meshes -import logging +# from pytorch3d.renderer.mesh.renderer import MeshRendererWithFragments logger = logging.getLogger(__name__) @@ -172,15 +186,15 @@ class PyRenderer: if len(cam) == 4: sx, sy, tx, ty = cam # sy = sx - camera_translation = np.array( - [tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9)] - ) + camera_translation = np.array([ + tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9) + ]) elif len(cam) == 3: sx, tx, ty = cam sy = sx - camera_translation = np.array( - [-tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9)] - ) + camera_translation = np.array([ + -tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9) + ]) render_res = resolution self.renderer.viewport_width = render_res[1] self.renderer.viewport_height = render_res[0] @@ -283,12 +297,8 @@ class OpenDRenderer: self.resolution = (resolution[0] * ratio, resolution[1] * ratio) self.ratio = ratio self.focal_length = 5000. - self.K = np.array( - [ - [self.focal_length, 0., self.resolution[1] / 2.], - [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.] - ] - ) + self.K = np.array([[self.focal_length, 0., self.resolution[1] / 2.], + [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.]]) self.colors_dict = { 'red': np.array([0.5, 0.2, 0.2]), 'pink': np.array([0.7, 0.5, 0.5]), @@ -303,12 +313,8 @@ class OpenDRenderer: def reset_res(self, resolution): self.resolution = (resolution[0] * self.ratio, resolution[1] * self.ratio) - self.K = np.array( - [ - [self.focal_length, 0., self.resolution[1] / 2.], - [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.] - ] - ) + self.K = np.array([[self.focal_length, 0., self.resolution[1] / 2.], + [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.]]) def __call__( self, @@ -446,27 +452,22 @@ class OpenDRenderer: # https://github.com/classner/up/blob/master/up_tools/camera.py def rotateY(points, angle): """Rotate all points in a 2D array around the y axis.""" - ry = np.array( - [[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.], [-np.sin(angle), 0., - np.cos(angle)]] - ) + ry = np.array([[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.], + [-np.sin(angle), 0., np.cos(angle)]]) return np.dot(points, ry) def rotateX(points, angle): """Rotate all points in a 2D array around the x axis.""" - rx = np.array( - [[1., 0., 0.], [0., np.cos(angle), -np.sin(angle)], [0., np.sin(angle), - np.cos(angle)]] - ) + rx = np.array([[1., 0., 0.], [0., np.cos(angle), -np.sin(angle)], + [0., np.sin(angle), np.cos(angle)]]) return np.dot(points, rx) def rotateZ(points, angle): """Rotate all points in a 2D array around the z axis.""" - rz = np.array( - [[np.cos(angle), -np.sin(angle), 0.], [np.sin(angle), np.cos(angle), 0.], [0., 0., 1.]] - ) + rz = np.array([[np.cos(angle), -np.sin(angle), 0.], [np.sin(angle), + np.cos(angle), 0.], [0., 0., 1.]]) return np.dot(points, rz) @@ -514,12 +515,8 @@ class IUV_Renderer(object): break np.save(dp_vert_pid_fname, np.array(dp_vert_pid)) - textures_vts = np.array( - [ - (dp_vert_pid[i] / num_part, DP.U_norm[i], DP.V_norm[i]) - for i in range(len(vert_mapping)) - ] - ) + textures_vts = np.array([(dp_vert_pid[i] / num_part, DP.U_norm[i], DP.V_norm[i]) + for i in range(len(vert_mapping))]) self.textures_vts = torch.from_numpy( textures_vts[None].astype(np.float32) ) # (1, 7829, 3) @@ -569,12 +566,8 @@ class IUV_Renderer(object): # range(n_verts)]) self.textures_vts = torch.from_numpy(textures_vts[None].astype(np.float32)) - K = np.array( - [ - [self.focal_length, 0., self.orig_size / 2.], - [0., self.focal_length, self.orig_size / 2.], [0., 0., 1.] - ] - ) + K = np.array([[self.focal_length, 0., self.orig_size / 2.], + [0., self.focal_length, self.orig_size / 2.], [0., 0., 1.]]) R = np.array([[-1., 0., 0.], [0., -1., 0.], [0., 0., 1.]]) @@ -620,10 +613,10 @@ class IUV_Renderer(object): K = self.K.repeat(batch_size, 1, 1) R = self.R.repeat(batch_size, 1, 1) - t = torch.stack( - [-cam[:, 1], -cam[:, 2], 2 * self.focal_length / (self.orig_size * cam[:, 0] + 1e-9)], - dim=-1 - ) + t = torch.stack([ + -cam[:, 1], -cam[:, 2], 2 * self.focal_length / (self.orig_size * cam[:, 0] + 1e-9) + ], + dim=-1) if cam.is_cuda: # device_id = cam.get_device() diff --git a/lib/pymafx/utils/sample_mesh.py b/lib/pymafx/utils/sample_mesh.py index 2599bee12d2577b6826ea8bfad8c937f2bcc2db2..1a696cf9580d7cd0e27ac0ade1db3c9e6a9c6e84 100644 --- a/lib/pymafx/utils/sample_mesh.py +++ b/lib/pymafx/utils/sample_mesh.py @@ -1,6 +1,8 @@ import os -import trimesh + import numpy as np +import trimesh + from .utils.libmesh import check_mesh_contains diff --git a/lib/pymafx/utils/saver.py b/lib/pymafx/utils/saver.py index 6a6bd3a184cc658dbc666ad2dcf3bc15d8cc427b..faed475e6bc4a8f1e2e3cd16d81b267d9bbb8496 100644 --- a/lib/pymafx/utils/saver.py +++ b/lib/pymafx/utils/saver.py @@ -1,8 +1,10 @@ from __future__ import division -import os -import torch + import datetime import logging +import os + +import torch logger = logging.getLogger(__name__) @@ -102,10 +104,8 @@ class CheckpointSaver(): if optimizer in checkpoint: optimizers[optimizer].load_state_dict(checkpoint[optimizer]) return { - 'epoch': checkpoint['epoch'], - 'batch_idx': checkpoint['batch_idx'], - 'batch_size': checkpoint['batch_size'], - 'total_step_count': checkpoint['total_step_count'] + 'epoch': checkpoint['epoch'], 'batch_idx': checkpoint['batch_idx'], 'batch_size': + checkpoint['batch_size'], 'total_step_count': checkpoint['total_step_count'] } def get_latest_checkpoint(self): diff --git a/lib/pymafx/utils/segms.py b/lib/pymafx/utils/segms.py index 44c617529d67323a8664c3e00872e5db091b8be6..c8fbf7e2c49422449cf4a8c4a38e1f320a0b15c0 100644 --- a/lib/pymafx/utils/segms.py +++ b/lib/pymafx/utils/segms.py @@ -21,13 +21,14 @@ The following terms are used in this module RLE: COCO's run length encoding format """ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) import numpy as np - import pycocotools.mask as mask_util diff --git a/lib/pymafx/utils/smooth_bbox.py b/lib/pymafx/utils/smooth_bbox.py index 4393320e7f50128d6838d99c76b5d0f8f45f6efc..a5eb8e431840f974b2cd896d80e8368755f9a682 100644 --- a/lib/pymafx/utils/smooth_bbox.py +++ b/lib/pymafx/utils/smooth_bbox.py @@ -93,12 +93,10 @@ def get_all_bbox_params(kps, vis_thresh=2): # Linearly interpolate each param. previous = bbox_params[-1] # This will be 3x(n+2) - interpolated = np.array( - [ - np.linspace(prev, curr, num_to_interpolate + 2) - for prev, curr in zip(previous, bbox_param) - ] - ) + interpolated = np.array([ + np.linspace(prev, curr, num_to_interpolate + 2) + for prev, curr in zip(previous, bbox_param) + ]) bbox_params = np.vstack((bbox_params, interpolated.T[1:-1])) num_to_interpolate = 0 bbox_params = np.vstack((bbox_params, bbox_param)) diff --git a/lib/pymafx/utils/transforms.py b/lib/pymafx/utils/transforms.py index 25534674631d40b8b263b242d05339443b169dcb..5f4189ee0e2da45e565b322d207b011ae3ed70f5 100644 --- a/lib/pymafx/utils/transforms.py +++ b/lib/pymafx/utils/transforms.py @@ -4,12 +4,10 @@ # Written by Bin Xiao (Bin.Xiao@microsoft.com) # ------------------------------------------------------------------------------ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function +from __future__ import absolute_import, division, print_function -import numpy as np import cv2 +import numpy as np def flip_back(output_flipped, matched_parts): diff --git a/lib/pymafx/utils/uv_vis.py b/lib/pymafx/utils/uv_vis.py index 86fdd33ddee774c2bbe02478b2d74f53f8522256..2bac9e86cdeb7d0091382914188755e7ecbf541a 100644 --- a/lib/pymafx/utils/uv_vis.py +++ b/lib/pymafx/utils/uv_vis.py @@ -1,10 +1,11 @@ import os -import torch + +# Use a non-interactive backend +import matplotlib import numpy as np +import torch import torch.nn.functional as F from skimage.transform import resize -# Use a non-interactive backend -import matplotlib matplotlib.use('Agg') @@ -100,9 +101,8 @@ def vis_smpl_iuv( for draw_i in range(len(cam_pred)): err_val = '{:06d}_'.format(int(10 * vert_errors_batch[draw_i])) draw_name = err_val + image_name[draw_i] - K = np.array( - [[focal_length, 0., orig_size / 2.], [0., focal_length, orig_size / 2.], [0., 0., 1.]] - ) + K = np.array([[focal_length, 0., orig_size / 2.], [0., focal_length, orig_size / 2.], + [0., 0., 1.]]) # img_orig, img_resized, img_smpl, render_smpl_rgba = dr_render( # image[draw_i], diff --git a/lib/pymafx/utils/vis.py b/lib/pymafx/utils/vis.py index 5273707c05f66275150e7cb2d86f44dcf4c92223..f64a490e60ded94ee8cffa51dae26f2ffa3a6ead 100644 --- a/lib/pymafx/utils/vis.py +++ b/lib/pymafx/utils/vis.py @@ -17,24 +17,26 @@ # limitations under the License. ############################################################################## -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import ( + absolute_import, + division, + print_function, + unicode_literals, +) + +import math +import os import cv2 +# Use a non-interactive backend +import matplotlib import numpy as np -import os import pycocotools.mask as mask_util -import math import torchvision from .colormap import colormap -from .keypoints import get_keypoints from .imutils import normalize_2d_kp - -# Use a non-interactive backend -import matplotlib +from .keypoints import get_keypoints matplotlib.use('Agg') import matplotlib.pyplot as plt @@ -191,15 +193,13 @@ def vis_one_image( print(dataset.classes[classes[i]], score) # show box (off by default, box_alpha=0.0) ax.add_patch( - plt.Rectangle( - (bbox[0], bbox[1]), - bbox[2] - bbox[0], - bbox[3] - bbox[1], - fill=False, - edgecolor='g', - linewidth=0.5, - alpha=box_alpha - ) + plt.Rectangle((bbox[0], bbox[1]), + bbox[2] - bbox[0], + bbox[3] - bbox[1], + fill=False, + edgecolor='g', + linewidth=0.5, + alpha=box_alpha) ) if show_class: diff --git a/lib/smplx/__init__.py b/lib/smplx/__init__.py index 886949df670691d1ef5995737cafa285224826c4..f29a068748ad6f0cd5e246a7678aeda15991c1ea 100644 --- a/lib/smplx/__init__.py +++ b/lib/smplx/__init__.py @@ -15,16 +15,16 @@ # Contact: ps-license@tuebingen.mpg.de from .body_models import ( - create, + FLAME, + MANO, SMPL, SMPLH, SMPLX, - MANO, - FLAME, - build_layer, - SMPLLayer, + FLAMELayer, + MANOLayer, SMPLHLayer, + SMPLLayer, SMPLXLayer, - MANOLayer, - FLAMELayer, + build_layer, + create, ) diff --git a/lib/smplx/body_models.py b/lib/smplx/body_models.py index b98adb635a3f9296c102e2bb6ca93bcdb14ab57d..ad55315dc5aa23422352bcca1d17b5f177d55e5d 100644 --- a/lib/smplx/body_models.py +++ b/lib/smplx/body_models.py @@ -14,36 +14,34 @@ # # Contact: ps-license@tuebingen.mpg.de -from typing import Optional, Dict, Union +import logging import os import os.path as osp import pickle +from collections import namedtuple +from typing import Dict, Optional, Union import numpy as np import torch import torch.nn as nn -from collections import namedtuple - -import logging logging.getLogger("smplx").setLevel(logging.ERROR) -from .lbs import lbs, vertices2landmarks, find_dynamic_lmk_idx_and_bcoords - -from .vertex_ids import vertex_ids as VERTEX_IDS +from .lbs import find_dynamic_lmk_idx_and_bcoords, lbs, vertices2landmarks from .utils import ( - Struct, - to_np, - to_tensor, - Tensor, Array, - SMPLOutput, + FLAMEOutput, + MANOOutput, SMPLHOutput, + SMPLOutput, SMPLXOutput, - MANOOutput, - FLAMEOutput, + Struct, + Tensor, find_joint_kin_chain, + to_np, + to_tensor, ) +from .vertex_ids import vertex_ids as VERTEX_IDS from .vertex_joint_selector import VertexJointSelector ModelOutput = namedtuple( @@ -1110,9 +1108,8 @@ class SMPLX(SMPLH): if create_expression: if expression is None: - default_expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype - ) + default_expression = torch.zeros([batch_size, self.num_expression_coeffs], + dtype=dtype) else: default_expression = torch.tensor(expression, dtype=dtype) expression_param = nn.Parameter(default_expression, requires_grad=True) @@ -1337,9 +1334,9 @@ class SMPLX(SMPLH): dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1 - ) + lmk_bary_coords = torch.cat([ + lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords + ], 1) landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords) @@ -1513,9 +1510,9 @@ class SMPLXLayer(SMPLX): dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() ) if expression is None: - expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype, device=device - ) + expression = torch.zeros([batch_size, self.num_expression_coeffs], + dtype=dtype, + device=device) if betas is None: betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device) if transl is None: @@ -1564,9 +1561,9 @@ class SMPLXLayer(SMPLX): dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1 - ) + lmk_bary_coords = torch.cat([ + lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords + ], 1) landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords) @@ -2044,9 +2041,8 @@ class FLAME(SMPL): if create_expression: if expression is None: - default_expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype - ) + default_expression = torch.zeros([batch_size, self.num_expression_coeffs], + dtype=dtype) else: default_expression = torch.tensor(expression, dtype=dtype) expression_param = nn.Parameter(default_expression, requires_grad=True) @@ -2202,9 +2198,9 @@ class FLAME(SMPL): ) dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1 - ) + lmk_bary_coords = torch.cat([ + lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords + ], 1) landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords) @@ -2331,9 +2327,9 @@ class FLAMELayer(FLAME): if betas is None: betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device) if expression is None: - expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype, device=device - ) + expression = torch.zeros([batch_size, self.num_expression_coeffs], + dtype=dtype, + device=device) if transl is None: transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) @@ -2367,9 +2363,9 @@ class FLAMELayer(FLAME): ) dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1 - ) + lmk_bary_coords = torch.cat([ + lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords + ], 1) landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords) diff --git a/lib/smplx/lbs.py b/lib/smplx/lbs.py index ac64f4b41be569331d632bfeb50fef9c50dc3d71..13862e837723d71ed70ecf7b68dd41e84ebd772c 100644 --- a/lib/smplx/lbs.py +++ b/lib/smplx/lbs.py @@ -14,17 +14,15 @@ # # Contact: ps-license@tuebingen.mpg.de -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division +from __future__ import absolute_import, division, print_function -from typing import Tuple, List, Optional -import numpy as np +from typing import List, Optional, Tuple +import numpy as np import torch import torch.nn.functional as F -from .utils import rot_mat_to_euler, Tensor +from .utils import Tensor, rot_mat_to_euler def find_dynamic_lmk_idx_and_bcoords( diff --git a/lib/smplx/utils.py b/lib/smplx/utils.py index d43a25217573f4c327adbf0411a76d1081632a69..6e2a2ff45f6597c2f8a0acf0820a3eb7dac9cfb2 100644 --- a/lib/smplx/utils.py +++ b/lib/smplx/utils.py @@ -14,8 +14,9 @@ # # Contact: ps-license@tuebingen.mpg.de -from typing import NewType, Union, Optional -from dataclasses import dataclass, asdict, fields +from dataclasses import asdict, dataclass, fields +from typing import NewType, Optional, Union + import numpy as np import torch diff --git a/lib/smplx/vertex_ids.py b/lib/smplx/vertex_ids.py index 31ed146ed4b3529bfbe0c92450bd3b02559f338b..060ed8ed60117a33358944abb85891abb3de8e30 100644 --- a/lib/smplx/vertex_ids.py +++ b/lib/smplx/vertex_ids.py @@ -14,61 +14,57 @@ # # Contact: ps-license@tuebingen.mpg.de -from __future__ import print_function -from __future__ import absolute_import -from __future__ import division +from __future__ import absolute_import, division, print_function # Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to # MSCOCO and OpenPose joints vertex_ids = { - "smplh": - { - "nose": 332, - "reye": 6260, - "leye": 2800, - "rear": 4071, - "lear": 583, - "rthumb": 6191, - "rindex": 5782, - "rmiddle": 5905, - "rring": 6016, - "rpinky": 6133, - "lthumb": 2746, - "lindex": 2319, - "lmiddle": 2445, - "lring": 2556, - "lpinky": 2673, - "LBigToe": 3216, - "LSmallToe": 3226, - "LHeel": 3387, - "RBigToe": 6617, - "RSmallToe": 6624, - "RHeel": 6787, - }, - "smplx": - { - "nose": 9120, - "reye": 9929, - "leye": 9448, - "rear": 616, - "lear": 6, - "rthumb": 8079, - "rindex": 7669, - "rmiddle": 7794, - "rring": 7905, - "rpinky": 8022, - "lthumb": 5361, - "lindex": 4933, - "lmiddle": 5058, - "lring": 5169, - "lpinky": 5286, - "LBigToe": 5770, - "LSmallToe": 5780, - "LHeel": 8846, - "RBigToe": 8463, - "RSmallToe": 8474, - "RHeel": 8635, - }, + "smplh": { + "nose": 332, + "reye": 6260, + "leye": 2800, + "rear": 4071, + "lear": 583, + "rthumb": 6191, + "rindex": 5782, + "rmiddle": 5905, + "rring": 6016, + "rpinky": 6133, + "lthumb": 2746, + "lindex": 2319, + "lmiddle": 2445, + "lring": 2556, + "lpinky": 2673, + "LBigToe": 3216, + "LSmallToe": 3226, + "LHeel": 3387, + "RBigToe": 6617, + "RSmallToe": 6624, + "RHeel": 6787, + }, + "smplx": { + "nose": 9120, + "reye": 9929, + "leye": 9448, + "rear": 616, + "lear": 6, + "rthumb": 8079, + "rindex": 7669, + "rmiddle": 7794, + "rring": 7905, + "rpinky": 8022, + "lthumb": 5361, + "lindex": 4933, + "lmiddle": 5058, + "lring": 5169, + "lpinky": 5286, + "LBigToe": 5770, + "LSmallToe": 5780, + "LHeel": 8846, + "RBigToe": 8463, + "RSmallToe": 8474, + "RHeel": 8635, + }, "mano": { "thumb": 744, "index": 320, diff --git a/lib/smplx/vertex_joint_selector.py b/lib/smplx/vertex_joint_selector.py index 1680e07acb03402a54fc0621ab36ec1d4de2c78e..dfee954435df6a3e674eb0c553476caf5d2a019a 100644 --- a/lib/smplx/vertex_joint_selector.py +++ b/lib/smplx/vertex_joint_selector.py @@ -14,12 +14,9 @@ # # Contact: ps-license@tuebingen.mpg.de -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division +from __future__ import absolute_import, division, print_function import numpy as np - import torch import torch.nn as nn diff --git a/lib/torch_utils/custom_ops.py b/lib/torch_utils/custom_ops.py index 2170f4732aba52f614b7cec09ac62465275ad90b..c76fc0e6a9c41c9e5b8e861079865eae41189226 100644 --- a/lib/torch_utils/custom_ops.py +++ b/lib/torch_utils/custom_ops.py @@ -6,15 +6,15 @@ # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. -import os import glob -import torch -import torch.utils.cpp_extension -import importlib import hashlib +import importlib +import os import shutil from pathlib import Path +import torch +import torch.utils.cpp_extension from torch.utils.file_baton import FileBaton #---------------------------------------------------------------------------- diff --git a/lib/torch_utils/misc.py b/lib/torch_utils/misc.py index 61c266a84d83e9a486df52e725af1c51488951e4..4946f0cc5fd29bb20bec9db27d0285c35878ec43 100644 --- a/lib/torch_utils/misc.py +++ b/lib/torch_utils/misc.py @@ -6,12 +6,13 @@ # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. -import re import contextlib -import numpy as np -import torch +import re import warnings + import dnnlib +import numpy as np +import torch #---------------------------------------------------------------------------- # Cached construction of constant tensors. Avoids CPU=>GPU copy when the @@ -272,15 +273,13 @@ def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True): buffer_size = sum(t.numel() for t in e.unique_buffers) output_shapes = [str(list(e.outputs[0].shape)) for t in e.outputs] output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs] - rows += [ - [ - name + (':0' if len(e.outputs) >= 2 else ''), - str(param_size) if param_size else '-', - str(buffer_size) if buffer_size else '-', - (output_shapes + ['-'])[0], - (output_dtypes + ['-'])[0], - ] - ] + rows += [[ + name + (':0' if len(e.outputs) >= 2 else ''), + str(param_size) if param_size else '-', + str(buffer_size) if buffer_size else '-', + (output_shapes + ['-'])[0], + (output_dtypes + ['-'])[0], + ]] for idx in range(1, len(e.outputs)): rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]] param_total += param_size diff --git a/lib/torch_utils/ops/bias_act.py b/lib/torch_utils/ops/bias_act.py index d8cfdb65d25ed077827862bc70e860c450fe929a..81d07ac029a2c36c12bcc6caff59a73476bdaf1e 100644 --- a/lib/torch_utils/ops/bias_act.py +++ b/lib/torch_utils/ops/bias_act.py @@ -8,94 +8,94 @@ """Custom PyTorch ops for efficient bias and activation.""" import os +import traceback import warnings + +import dnnlib import numpy as np import torch -import dnnlib -import traceback -from .. import custom_ops -from .. import misc +from .. import custom_ops, misc #---------------------------------------------------------------------------- activation_funcs = { 'linear': - dnnlib.EasyDict( - func=lambda x, **_: x, def_alpha=0, def_gain=1, cuda_idx=1, ref='', has_2nd_grad=False - ), + dnnlib.EasyDict( + func=lambda x, **_: x, def_alpha=0, def_gain=1, cuda_idx=1, ref='', has_2nd_grad=False + ), 'relu': - dnnlib.EasyDict( - func=lambda x, **_: torch.nn.functional.relu(x), - def_alpha=0, - def_gain=np.sqrt(2), - cuda_idx=2, - ref='y', - has_2nd_grad=False - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.nn.functional.relu(x), + def_alpha=0, + def_gain=np.sqrt(2), + cuda_idx=2, + ref='y', + has_2nd_grad=False + ), 'lrelu': - dnnlib.EasyDict( - func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha), - def_alpha=0.2, - def_gain=np.sqrt(2), - cuda_idx=3, - ref='y', - has_2nd_grad=False - ), + dnnlib.EasyDict( + func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha), + def_alpha=0.2, + def_gain=np.sqrt(2), + cuda_idx=3, + ref='y', + has_2nd_grad=False + ), 'tanh': - dnnlib.EasyDict( - func=lambda x, **_: torch.tanh(x), - def_alpha=0, - def_gain=1, - cuda_idx=4, - ref='y', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.tanh(x), + def_alpha=0, + def_gain=1, + cuda_idx=4, + ref='y', + has_2nd_grad=True + ), 'sigmoid': - dnnlib.EasyDict( - func=lambda x, **_: torch.sigmoid(x), - def_alpha=0, - def_gain=1, - cuda_idx=5, - ref='y', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.sigmoid(x), + def_alpha=0, + def_gain=1, + cuda_idx=5, + ref='y', + has_2nd_grad=True + ), 'elu': - dnnlib.EasyDict( - func=lambda x, **_: torch.nn.functional.elu(x), - def_alpha=0, - def_gain=1, - cuda_idx=6, - ref='y', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.nn.functional.elu(x), + def_alpha=0, + def_gain=1, + cuda_idx=6, + ref='y', + has_2nd_grad=True + ), 'selu': - dnnlib.EasyDict( - func=lambda x, **_: torch.nn.functional.selu(x), - def_alpha=0, - def_gain=1, - cuda_idx=7, - ref='y', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.nn.functional.selu(x), + def_alpha=0, + def_gain=1, + cuda_idx=7, + ref='y', + has_2nd_grad=True + ), 'softplus': - dnnlib.EasyDict( - func=lambda x, **_: torch.nn.functional.softplus(x), - def_alpha=0, - def_gain=1, - cuda_idx=8, - ref='y', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.nn.functional.softplus(x), + def_alpha=0, + def_gain=1, + cuda_idx=8, + ref='y', + has_2nd_grad=True + ), 'swish': - dnnlib.EasyDict( - func=lambda x, **_: torch.sigmoid(x) * x, - def_alpha=0, - def_gain=np.sqrt(2), - cuda_idx=9, - ref='x', - has_2nd_grad=True - ), + dnnlib.EasyDict( + func=lambda x, **_: torch.sigmoid(x) * x, + def_alpha=0, + def_gain=np.sqrt(2), + cuda_idx=9, + ref='x', + has_2nd_grad=True + ), } #---------------------------------------------------------------------------- diff --git a/lib/torch_utils/ops/conv2d_gradfix.py b/lib/torch_utils/ops/conv2d_gradfix.py index 29c3d8f5a8a1e2816e225af3157fc1bb99a4fd33..16bcdfdb229acf55b30a33c47ed419bd008bae7d 100644 --- a/lib/torch_utils/ops/conv2d_gradfix.py +++ b/lib/torch_utils/ops/conv2d_gradfix.py @@ -8,8 +8,9 @@ """Custom replacement for `torch.nn.functional.conv2d` that supports arbitrarily high order gradients with zero performance penalty.""" -import warnings import contextlib +import warnings + import torch # pylint: disable=redefined-builtin diff --git a/lib/torch_utils/ops/conv2d_resample.py b/lib/torch_utils/ops/conv2d_resample.py index 9f347c59165d1aceafee936b36281610b5a64e1b..2529947e6f5d34f9aa65bd21ede9e0fac87190ab 100644 --- a/lib/torch_utils/ops/conv2d_resample.py +++ b/lib/torch_utils/ops/conv2d_resample.py @@ -10,10 +10,8 @@ import torch from .. import misc -from . import conv2d_gradfix -from . import upfirdn2d -from .upfirdn2d import _parse_padding -from .upfirdn2d import _get_filter_size +from . import conv2d_gradfix, upfirdn2d +from .upfirdn2d import _get_filter_size, _parse_padding #---------------------------------------------------------------------------- diff --git a/lib/torch_utils/ops/grid_sample_gradfix.py b/lib/torch_utils/ops/grid_sample_gradfix.py index 850feacd5a6300b85493cd7f713bffab1af70536..b4049485344ac9cc09e5522b992962db05900dc7 100644 --- a/lib/torch_utils/ops/grid_sample_gradfix.py +++ b/lib/torch_utils/ops/grid_sample_gradfix.py @@ -11,6 +11,7 @@ Only works on 2D images and assumes `mode='bilinear'`, `padding_mode='zeros'`, `align_corners=False`.""" import warnings + import torch # pylint: disable=redefined-builtin diff --git a/lib/torch_utils/ops/upfirdn2d.py b/lib/torch_utils/ops/upfirdn2d.py index 86f6fb36eb83711db42aef6b05c003eceaeeaa69..0d5fa241d9d8d56a7a4e8c3cb9553c5916fbd4d2 100644 --- a/lib/torch_utils/ops/upfirdn2d.py +++ b/lib/torch_utils/ops/upfirdn2d.py @@ -8,13 +8,13 @@ """Custom PyTorch ops for efficient resampling of 2D images.""" import os +import traceback import warnings + import numpy as np import torch -import traceback -from .. import custom_ops -from .. import misc +from .. import custom_ops, misc from . import conv2d_gradfix #---------------------------------------------------------------------------- diff --git a/lib/torch_utils/persistence.py b/lib/torch_utils/persistence.py index c3263dc0690ac12d5d2e74a6d9d8d2af2fed0f5b..34f16d48f416fc07eb6954aa810acf852db5eed5 100644 --- a/lib/torch_utils/persistence.py +++ b/lib/torch_utils/persistence.py @@ -12,13 +12,14 @@ during unpickling. This way, any previously exported pickles will remain usable even if the original code is no longer available, or if the current version of the code is not consistent with what was originally pickled.""" -import sys -import pickle -import io -import inspect import copy -import uuid +import inspect +import io +import pickle +import sys import types +import uuid + import dnnlib #---------------------------------------------------------------------------- diff --git a/lib/torch_utils/training_stats.py b/lib/torch_utils/training_stats.py index 11658fdbf55450f5f0d4679e247ff65a4b37151e..a4fd0a4c3687aff712547b2688225ba1ec621f47 100644 --- a/lib/torch_utils/training_stats.py +++ b/lib/torch_utils/training_stats.py @@ -11,8 +11,10 @@ synchronization overhead as well as the amount of boilerplate in user code.""" import re + import numpy as np import torch + import lib.dnnlib from . import misc diff --git a/requirements.txt b/requirements.txt index ad16e028ac6e07019e69d3fa8407af0bd8eca9e4..44b6403e52c59039b093c764343d02700b0b42ef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,4 +14,5 @@ mediapipe einops boto3 open3d +xatlas git+https://github.com/YuliangXiu/rembg.git diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..718afe479d0ace65dabe450ac2d5b3d5060315f2 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,18 @@ +[yapf] +based_on_style = facebook +column_limit = 100 +indent_width = 4 +spaces_before_comment = 4 +split_all_comma_separated_values = false +split_all_top_level_comma_separated_values = false +dedent_closing_brackets = true +coalesce_brackets = true +split_before_dot = false +each_dict_entry_on_separate_line = false +indent_dictionary_value = false + +[isort] +multi_line_output = 3 +line_length = 80 +include_trailing_comma = true +skip=./log,./results,./data,./debug,./lib/common/libmesh/setup.py,./lib/common/libvoxelize/setup.py \ No newline at end of file