Support TEXTure

.gitignore

@@ -16,4 +16,5 @@ build
 dist
 *egg-info
 *.so
-run.sh
+run.sh
+*.log

README.md

@@ -25,6 +25,7 @@
     <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
     <a href="https://cupy.dev/"><img alt="cupy" src="https://img.shields.io/badge/-Cupy-46C02B?logo=numpy&logoColor=white"></a>
     <a href="https://twitter.com/yuliangxiu"><img alt='Twitter' src="https://img.shields.io/twitter/follow/yuliangxiu?label=%40yuliangxiu"></a>
+    <a href="https://discord.gg/Vqa7KBGRyk"><img alt="discord invitation link" src="https://dcbadge.vercel.app/api/server/Vqa7KBGRyk?style=flat"></a>
     <br></br>
     <a href='https://colab.research.google.com/drive/1YRgwoRCZIrSB2e7auEWFyG10Xzjbrbno?usp=sharing'><img src='https://colab.research.google.com/assets/colab-badge.svg' alt='Google Colab'></a>
     <a href='https://github.com/YuliangXiu/ECON/blob/master/docs/installation-docker.md'><img src='https://img.shields.io/badge/Docker-9cf.svg?logo=Docker' alt='Docker'></a>
@@ -35,7 +36,6 @@
     </a>
     <a href='https://xiuyuliang.cn/econ/'>
       <img src='https://img.shields.io/badge/ECON-Page-orange?style=for-the-badge&logo=Google%20chrome&logoColor=white&labelColor=D35400' alt='Project Page'></a>
-    <a href="https://discord.gg/Vqa7KBGRyk"><img src="https://img.shields.io/discord/940240966844035082?color=7289DA&labelColor=4a64bd&logo=discord&logoColor=white&style=for-the-badge"></a>
     <a href="https://youtu.be/j5hw4tsWpoY"><img alt="youtube views" title="Subscribe to my YouTube channel" src="https://img.shields.io/youtube/views/j5hw4tsWpoY?logo=youtube&labelColor=ce4630&style=for-the-badge"/></a>
   </p>
 </p>

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
 

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):

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")

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=[

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=[

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

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

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

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"

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

lib/common/blender_utils.py

@@ -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)

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
 

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

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),

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

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()])

lib/common/libmesh/triangle_hash.cpp

@@ -720,12 +720,12 @@ static CYTHON_INLINE float __PYX_NAN() {
 
     /* NumPy API declarations from "numpy/__init__.pxd" */
     
+#include <math.h>
 #include "ios"
 #include "new"
 #include "stdexcept"
 #include "typeinfo"
 #include <vector>
-#include <math.h>
 #include "pythread.h"
 #include <stdlib.h>
 #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] = <int> 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] = <int> 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] = <int> 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] = <int> 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);

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

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] = <int> (
  *             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] = <int> (             # <<<<<<<<<<<<<<
@@ -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] = <int> (
  *             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] = <int> (             # <<<<<<<<<<<<<<
@@ -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);

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],

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()

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()

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

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)

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
 

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):

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)

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))

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

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]
 

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):

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

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()

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

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)

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
 

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):

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

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
 
 

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])

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])

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")

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)

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):

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):

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
 
 

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]
 

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