remove MeshLab dependency with Open3D

README.md

@@ -103,20 +103,23 @@ python -m apps.avatarizer -n {filename}
 
 ### Some adjustable parameters in _config/econ.yaml_
 
-- `use_ifnet: True`
-  - True: use IF-Nets+ for mesh completion ( $\text{ECON}_\text{IF}$ - Better quality)
-  - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed)
+- `use_ifnet: False`
+  - True: use IF-Nets+ for mesh completion ( $\text{ECON}_\text{IF}$ - Better quality, **~2min / img**)
+  - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed, **~1.5min / img**)
 - `use_smpl: ["hand", "face"]`
   - [ ]: 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
 - `thickness: 2cm`
   - could be increased accordingly in case final reconstruction **xx_full.obj** looks flat
+- `k: 4`
+  - could be reduced accordingly in case the surface of **xx_full.obj** has discontinous artifacts
 - `hps_type: PIXIE`
   - "pixie": more accurate for face and hands
   - "pymafx": more robust for challenging poses
-- `k: 4`
-  - could be reduced accordingly in case the surface of **xx_full.obj** has discontinous artifacts
+- `texture_src: image`
+  - "image": direct mapping the aligned pixels to final mesh
+  - "SD": use Stable Diffusion to generate full texture (TODO)
 
 <br/>
 
@@ -160,7 +163,6 @@ Here are some great resources we benefit from:
 - [BiNI](https://github.com/hoshino042/bilateral_normal_integration) for Bilateral Normal Integration
 - [MonoPortDataset](https://github.com/Project-Splinter/MonoPortDataset) for Data Processing, [MonoPort](https://github.com/Project-Splinter/MonoPort) for fast implicit surface query
 - [rembg](https://github.com/danielgatis/rembg) for Human Segmentation
-- [pypoisson](https://github.com/mmolero/pypoisson) for poisson reconstruction
 - [MediaPipe](https://google.github.io/mediapipe/getting_started/python.html) for full-body landmark estimation
 - [PyTorch-NICP](https://github.com/wuhaozhe/pytorch-nicp) for non-rigid registration
 - [smplx](https://github.com/vchoutas/smplx), [PyMAF-X](https://www.liuyebin.com/pymaf-x/), [PIXIE](https://github.com/YadiraF/PIXIE) for Human Pose & Shape Estimation

apps/IFGeo.py

@@ -24,7 +24,6 @@ torch.backends.cudnn.benchmark = True
 
 
 class IFGeo(pl.LightningModule):
-
     def __init__(self, cfg):
         super(IFGeo, self).__init__()
 
@@ -44,14 +43,15 @@ class IFGeo(pl.LightningModule):
             from lib.net.IFGeoNet_nobody import IFGeoNet
             self.netG = IFGeoNet(cfg)
 
-
-        self.resolutions = (np.logspace(
-            start=5,
-            stop=np.log2(self.mcube_res),
-            base=2,
-            num=int(np.log2(self.mcube_res) - 4),
-            endpoint=True,
-        ) + 1.0)
+        self.resolutions = (
+            np.logspace(
+                start=5,
+                stop=np.log2(self.mcube_res),
+                base=2,
+                num=int(np.log2(self.mcube_res) - 4),
+                endpoint=True,
+            ) + 1.0
+        )
 
         self.resolutions = self.resolutions.astype(np.int16).tolist()
 
@@ -82,9 +82,9 @@ class IFGeo(pl.LightningModule):
 
         if self.cfg.optim == "Adadelta":
 
-            optimizer_G = torch.optim.Adadelta(optim_params_G,
-                                               lr=self.lr_G,
-                                               weight_decay=weight_decay)
+            optimizer_G = torch.optim.Adadelta(
+                optim_params_G, lr=self.lr_G, weight_decay=weight_decay
+            )
 
         elif self.cfg.optim == "Adam":
 
@@ -103,20 +103,14 @@ class IFGeo(pl.LightningModule):
             raise NotImplementedError
 
         # set scheduler
-        scheduler_G = torch.optim.lr_scheduler.MultiStepLR(optimizer_G,
-                                                           milestones=self.cfg.schedule,
-                                                           gamma=self.cfg.gamma)
+        scheduler_G = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_G, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
 
         return [optimizer_G], [scheduler_G]
 
     def training_step(self, batch, batch_idx):
 
-        # cfg log
-        if self.cfg.devices == 1:
-            if not self.cfg.fast_dev and self.global_step == 0:
-                export_cfg(self.logger, osp.join(self.cfg.results_path, self.cfg.name), self.cfg)
-                self.logger.experiment.config.update(convert_to_dict(self.cfg))
-
         self.netG.train()
 
         preds_G = self.netG(batch)
@@ -127,12 +121,9 @@ class IFGeo(pl.LightningModule):
             "loss": error_G,
         }
 
-        self.log_dict(metrics_log,
-                      prog_bar=True,
-                      logger=True,
-                      on_step=True,
-                      on_epoch=False,
-                      sync_dist=True)
+        self.log_dict(
+            metrics_log, prog_bar=True, logger=True, on_step=True, on_epoch=False, sync_dist=True
+        )
 
         return metrics_log
 
@@ -143,12 +134,14 @@ class IFGeo(pl.LightningModule):
             "train/avgloss": batch_mean(outputs, "loss"),
         }
 
-        self.log_dict(metrics_log,
-                      prog_bar=False,
-                      logger=True,
-                      on_step=False,
-                      on_epoch=True,
-                      rank_zero_only=True)
+        self.log_dict(
+            metrics_log,
+            prog_bar=False,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            rank_zero_only=True
+        )
 
     def validation_step(self, batch, batch_idx):
 
@@ -162,12 +155,9 @@ class IFGeo(pl.LightningModule):
             "val/loss": error_G,
         }
 
-        self.log_dict(metrics_log,
-                      prog_bar=True,
-                      logger=False,
-                      on_step=True,
-                      on_epoch=False,
-                      sync_dist=True)
+        self.log_dict(
+            metrics_log, prog_bar=True, logger=False, on_step=True, on_epoch=False, sync_dist=True
+        )
 
         return metrics_log
 
@@ -178,9 +168,11 @@ class IFGeo(pl.LightningModule):
             "val/avgloss": batch_mean(outputs, "val/loss"),
         }
 
-        self.log_dict(metrics_log,
-                      prog_bar=False,
-                      logger=True,
-                      on_step=False,
-                      on_epoch=True,
-                      rank_zero_only=True)
+        self.log_dict(
+            metrics_log,
+            prog_bar=False,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            rank_zero_only=True
+        )

apps/Normal.py

@@ -1,14 +1,12 @@
 from lib.net import NormalNet
-from lib.common.train_util import convert_to_dict, export_cfg, batch_mean
+from lib.common.train_util import batch_mean
 import torch
 import numpy as np
-import os.path as osp
 from skimage.transform import resize
 import pytorch_lightning as pl
 
 
 class Normal(pl.LightningModule):
-
     def __init__(self, cfg):
         super(Normal, self).__init__()
         self.cfg = cfg
@@ -44,19 +42,19 @@ class Normal(pl.LightningModule):
         optimizer_N_F = torch.optim.Adam(optim_params_N_F, lr=self.lr_F, betas=(0.5, 0.999))
         optimizer_N_B = torch.optim.Adam(optim_params_N_B, lr=self.lr_B, betas=(0.5, 0.999))
 
-        scheduler_N_F = torch.optim.lr_scheduler.MultiStepLR(optimizer_N_F,
-                                                             milestones=self.cfg.schedule,
-                                                             gamma=self.cfg.gamma)
+        scheduler_N_F = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_N_F, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
 
-        scheduler_N_B = torch.optim.lr_scheduler.MultiStepLR(optimizer_N_B,
-                                                             milestones=self.cfg.schedule,
-                                                             gamma=self.cfg.gamma)
+        scheduler_N_B = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer_N_B, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+        )
         if 'gan' in self.ALL_losses:
             optim_params_N_D = [{"params": self.netG.netD.parameters(), "lr": self.lr_D}]
             optimizer_N_D = torch.optim.Adam(optim_params_N_D, lr=self.lr_D, betas=(0.5, 0.999))
-            scheduler_N_D = torch.optim.lr_scheduler.MultiStepLR(optimizer_N_D,
-                                                                 milestones=self.cfg.schedule,
-                                                                 gamma=self.cfg.gamma)
+            scheduler_N_D = torch.optim.lr_scheduler.MultiStepLR(
+                optimizer_N_D, milestones=self.cfg.schedule, gamma=self.cfg.gamma
+            )
             self.schedulers = [scheduler_N_F, scheduler_N_B, scheduler_N_D]
             optims = [optimizer_N_F, optimizer_N_B, optimizer_N_D]
 
@@ -77,19 +75,16 @@ class Normal(pl.LightningModule):
                     ((render_tensor[name].cpu().numpy()[0] + 1.0) / 2.0).transpose(1, 2, 0),
                     (height, height),
                     anti_aliasing=True,
-                ))
+                )
+            )
 
-        self.logger.log_image(key=f"Normal/{dataset}/{idx if not self.overfit else 1}",
-                              images=[(np.concatenate(result_list, axis=1) * 255.0).astype(np.uint8)
-                                     ])
+        self.logger.log_image(
+            key=f"Normal/{dataset}/{idx if not self.overfit else 1}",
+            images=[(np.concatenate(result_list, axis=1) * 255.0).astype(np.uint8)]
+        )
 
     def training_step(self, batch, batch_idx):
 
-        # cfg log
-        if not self.cfg.fast_dev and self.global_step == 0 and self.cfg.devices == 1:
-            export_cfg(self.logger, osp.join(self.cfg.results_path, self.cfg.name), self.cfg)
-            self.logger.experiment.config.update(convert_to_dict(self.cfg))
-
         self.netG.train()
 
         # retrieve the data
@@ -125,7 +120,8 @@ class Normal(pl.LightningModule):
             opt_B.step()
 
         if batch_idx > 0 and batch_idx % int(
-                self.cfg.freq_show_train) == 0 and self.cfg.devices == 1:
+            self.cfg.freq_show_train
+        ) == 0 and self.cfg.devices == 1:
 
             self.netG.eval()
             with torch.no_grad():
@@ -142,12 +138,9 @@ class Normal(pl.LightningModule):
         for key in error_dict.keys():
             metrics_log["train/loss_" + key] = error_dict[key].item()
 
-        self.log_dict(metrics_log,
-                      prog_bar=True,
-                      logger=True,
-                      on_step=True,
-                      on_epoch=False,
-                      sync_dist=True)
+        self.log_dict(
+            metrics_log, prog_bar=True, logger=True, on_step=True, on_epoch=False, sync_dist=True
+        )
 
         return metrics_log
 
@@ -163,12 +156,14 @@ class Normal(pl.LightningModule):
                 loss_name = key
             metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key)
 
-        self.log_dict(metrics_log,
-                      prog_bar=False,
-                      logger=True,
-                      on_step=False,
-                      on_epoch=True,
-                      rank_zero_only=True)
+        self.log_dict(
+            metrics_log,
+            prog_bar=False,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            rank_zero_only=True
+        )
 
     def validation_step(self, batch, batch_idx):
 
@@ -212,9 +207,11 @@ class Normal(pl.LightningModule):
             [stage, loss_name] = key.split("/")
             metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key)
 
-        self.log_dict(metrics_log,
-                      prog_bar=False,
-                      logger=True,
-                      on_step=False,
-                      on_epoch=True,
-                      rank_zero_only=True)
+        self.log_dict(
+            metrics_log,
+            prog_bar=False,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            rank_zero_only=True
+        )

apps/avatarizer.py

@@ -44,7 +44,8 @@ smpl_model = smplx.create(
     use_pca=False,
     num_betas=200,
     num_expression_coeffs=50,
-    ext='pkl')
+    ext='pkl'
+)
 
 smpl_out_lst = []
 
@@ -62,7 +63,9 @@ for pose_type in ["t-pose", "da-pose", "pose"]:
             return_full_pose=True,
             return_joint_transformation=True,
             return_vertex_transformation=True,
-            pose_type=pose_type))
+            pose_type=pose_type
+        )
+    )
 
 smpl_verts = smpl_out_lst[2].vertices.detach()[0]
 smpl_tree = cKDTree(smpl_verts.cpu().numpy())
@@ -74,7 +77,8 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.
     econ_verts = torch.tensor(econ_obj.vertices).float()
     rot_mat_t = smpl_out_lst[2].vertex_transformation.detach()[0][idx[:, 0]]
     homo_coord = torch.ones_like(econ_verts)[..., :1]
-    econ_cano_verts = torch.inverse(rot_mat_t) @ torch.cat([econ_verts, homo_coord], dim=1).unsqueeze(-1)
+    econ_cano_verts = torch.inverse(rot_mat_t) @ torch.cat([econ_verts, homo_coord],
+                                                           dim=1).unsqueeze(-1)
     econ_cano_verts = econ_cano_verts[:, :3, 0].cpu()
     econ_cano = trimesh.Trimesh(econ_cano_verts, econ_obj.faces)
 
@@ -84,7 +88,9 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.
     econ_da = trimesh.Trimesh(econ_da_verts[:, :3, 0].cpu(), econ_obj.faces)
 
     # da-pose for SMPL-X
-    smpl_da = trimesh.Trimesh(smpl_out_lst[1].vertices.detach()[0], smpl_model.faces, maintain_orders=True, process=False)
+    smpl_da = trimesh.Trimesh(
+        smpl_out_lst[1].vertices.detach()[0], smpl_model.faces, maintain_orders=True, process=False
+    )
     smpl_da.export(f"{prefix}_smpl_da.obj")
 
     # remove hands from ECON for next registeration
@@ -97,7 +103,8 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.
     # remove SMPL-X hand and face
     register_mask = ~np.isin(
         np.arange(smpl_da.vertices.shape[0]),
-        np.concatenate([smplx_container.smplx_mano_vid, smplx_container.smplx_front_flame_vid]))
+        np.concatenate([smplx_container.smplx_mano_vid, smplx_container.smplx_front_flame_vid])
+    )
     register_mask *= ~smplx_container.eyeball_vertex_mask.bool().numpy()
     smpl_da_body = smpl_da.copy()
     smpl_da_body.update_faces(register_mask[smpl_da.faces].all(axis=1))
@@ -115,8 +122,13 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.
     # remove over-streched+hand faces from ECON
     econ_da_body = econ_da.copy()
     edge_before = np.sqrt(
-        ((econ_obj.vertices[econ_cano.edges[:, 0]] - econ_obj.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1))
-    edge_after = np.sqrt(((econ_da.vertices[econ_cano.edges[:, 0]] - econ_da.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1))
+        ((econ_obj.vertices[econ_cano.edges[:, 0]] -
+          econ_obj.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1)
+    )
+    edge_after = np.sqrt(
+        ((econ_da.vertices[econ_cano.edges[:, 0]] -
+          econ_da.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1)
+    )
     edge_diff = edge_after / edge_before.clip(1e-2)
     streched_mask = np.unique(econ_cano.edges[edge_diff > 6])
     mano_mask = ~np.isin(idx[:, 0], smplx_container.smplx_mano_vid)
@@ -148,8 +160,9 @@ econ_J_regressor = (smpl_model.J_regressor[:, idx] * knn_weights[None]).sum(axis
 econ_lbs_weights = (smpl_model.lbs_weights.T[:, idx] * knn_weights[None]).sum(axis=-1).T
 
 num_posedirs = smpl_model.posedirs.shape[0]
-econ_posedirs = (smpl_model.posedirs.view(num_posedirs, -1, 3)[:, idx, :] *
-                 knn_weights[None, ..., None]).sum(axis=-2).view(num_posedirs, -1).float()
+econ_posedirs = (
+    smpl_model.posedirs.view(num_posedirs, -1, 3)[:, idx, :] * knn_weights[None, ..., None]
+).sum(axis=-2).view(num_posedirs, -1).float()
 
 econ_J_regressor /= econ_J_regressor.sum(axis=1, keepdims=True)
 econ_lbs_weights /= econ_lbs_weights.sum(axis=1, keepdims=True)
@@ -157,8 +170,9 @@ econ_lbs_weights /= econ_lbs_weights.sum(axis=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()
 
 # ----------------------------------------------------
@@ -174,7 +188,8 @@ posed_econ_verts, _ = general_lbs(
     posedirs=econ_posedirs,
     J_regressor=econ_J_regressor,
     parents=smpl_model.parents,
-    lbs_weights=econ_lbs_weights)
+    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")
+econ_pose.export(f"{prefix}_econ_pose.obj")

apps/infer.py

@@ -34,7 +34,8 @@ 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, load_normal_networks, load_networks
+from lib.common.train_util import init_loss, Format
+from lib.common.imutils import blend_rgb_norm
 from lib.common.BNI import BNI
 from lib.common.BNI_utils import save_normal_tensor
 from lib.dataset.TestDataset import TestDataset
@@ -68,20 +69,25 @@ if __name__ == "__main__":
     device = torch.device(f"cuda:{args.gpu_device}")
 
     # setting for testing on in-the-wild images
-    cfg_show_list = ["test_gpus", [args.gpu_device], "mcube_res", 512, "clean_mesh", True, "test_mode", True, "batch_size", 1]
+    cfg_show_list = [
+        "test_gpus", [args.gpu_device], "mcube_res", 512, "clean_mesh", True, "test_mode", True,
+        "batch_size", 1
+    ]
 
     cfg.merge_from_list(cfg_show_list)
     cfg.freeze()
 
-    # load model
-    normal_model = Normal(cfg).to(device)
-    load_normal_networks(normal_model, cfg.normal_path)
-    normal_model.netG.eval()
-
-    # load IFGeo model
-    ifnet_model = IFGeo(cfg).to(device)
-    load_networks(ifnet_model, mlp_path=cfg.ifnet_path)
-    ifnet_model.netG.eval()
+    # load normal model
+    normal_net = Normal.load_from_checkpoint(
+        cfg=cfg, checkpoint_path=cfg.normal_path, map_location=device, strict=False
+    )
+    normal_net = normal_net.to(device)
+    normal_net.netG.eval()
+    print(
+        colored(
+            f"Resume Normal Estimator from {Format.start} {cfg.normal_path} {Format.end}", "green"
+        )
+    )
 
     # SMPLX object
     SMPLX_object = SMPLX()
@@ -89,16 +95,24 @@ if __name__ == "__main__":
     dataset_param = {
         "image_dir": args.in_dir,
         "seg_dir": args.seg_dir,
-        "use_seg": True,  # w/ or w/o segmentation
-        "hps_type": cfg.bni.hps_type,  # pymafx/pixie
+        "use_seg": True,    # w/ or w/o segmentation
+        "hps_type": cfg.bni.hps_type,    # pymafx/pixie
         "vol_res": cfg.vol_res,
         "single": args.multi,
     }
 
     if cfg.bni.use_ifnet:
-        print(colored("Use IF-Nets (Implicit)+ for completion", "green"))
+        # load IFGeo model
+        ifnet = IFGeo.load_from_checkpoint(
+            cfg=cfg, checkpoint_path=cfg.ifnet_path, map_location=device, strict=False
+        )
+        ifnet = ifnet.to(device)
+        ifnet.netG.eval()
+
+        print(colored(f"Resume IF-Net+ from {Format.start} {cfg.ifnet_path} {Format.end}", "green"))
+        print(colored(f"Complete with {Format.start} IF-Nets+ (Implicit) {Format.end}", "green"))
     else:
-        print(colored("Use SMPL-X (Explicit) for completion", "green"))
+        print(colored(f"Complete with {Format.start} SMPL-X (Explicit) {Format.end}", "green"))
 
     dataset = TestDataset(dataset_param, device)
 
@@ -125,13 +139,17 @@ if __name__ == "__main__":
         # 2. SMPL params (xxx_smpl.npy)
         # 3. d-BiNI surfaces (xxx_BNI.obj)
         # 4. seperate face/hand mesh (xxx_hand/face.obj)
-        # 5. full shape impainted by IF-Nets+, and remeshed shape (xxx_IF_(remesh).obj)
+        # 5. full shape impainted by IF-Nets+ after remeshing (xxx_IF.obj)
         # 6. sideded or occluded parts (xxx_side.obj)
         # 7. final reconstructed clothed human (xxx_full.obj)
 
         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)}
+        in_tensor = {
+            "smpl_faces": data["smpl_faces"],
+            "image": data["img_icon"].to(device),
+            "mask": data["img_mask"].to(device)
+        }
 
         # The optimizer and variables
         optimed_pose = data["body_pose"].requires_grad_(True)
@@ -139,7 +157,9 @@ 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",
@@ -156,10 +176,12 @@ if __name__ == "__main__":
 
         smpl_path = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl_00.obj"
 
+        # remove this line if you change the loop_smpl and obtain different SMPL-X fits
         if osp.exists(smpl_path):
 
             smpl_verts_lst = []
             smpl_faces_lst = []
+
             for idx in range(N_body):
 
                 smpl_obj = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl_{idx:02d}.obj"
@@ -173,10 +195,12 @@ if __name__ == "__main__":
             batch_smpl_faces = torch.stack(smpl_faces_lst)
 
             # render optimized mesh as normal [-1,1]
-            in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal(batch_smpl_verts, batch_smpl_faces)
+            in_tensor["T_normal_F"], in_tensor["T_normal_B"] = dataset.render_normal(
+                batch_smpl_verts, batch_smpl_faces
+            )
 
             with torch.no_grad():
-                in_tensor["normal_F"], in_tensor["normal_B"] = normal_model.netG(in_tensor)
+                in_tensor["normal_F"], in_tensor["normal_B"] = normal_net.netG(in_tensor)
 
             in_tensor["smpl_verts"] = batch_smpl_verts * torch.tensor([1., -1., 1.]).to(device)
             in_tensor["smpl_faces"] = batch_smpl_faces[:, :, [0, 2, 1]]
@@ -194,8 +218,10 @@ if __name__ == "__main__":
                 N_body, N_pose = optimed_pose.shape[:2]
 
                 # 6d_rot to rot_mat
-                optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1, 6)).view(N_body, 1, 3, 3)
-                optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1, 6)).view(N_body, N_pose, 3, 3)
+                optimed_orient_mat = rot6d_to_rotmat(optimed_orient.view(-1,
+                                                                         6)).view(N_body, 1, 3, 3)
+                optimed_pose_mat = rot6d_to_rotmat(optimed_pose.view(-1,
+                                                                     6)).view(N_body, N_pose, 3, 3)
 
                 smpl_verts, smpl_landmarks, smpl_joints = dataset.smpl_model(
                     shape_params=optimed_betas,
@@ -208,11 +234,16 @@ 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 = (smpl_joints[:, dataset.smpl_data.smpl_joint_ids_45_pixie, :] + 1.0) * 0.5
-                in_tensor["smpl_joint"] = smpl_joints[:, dataset.smpl_data.smpl_joint_ids_24_pixie, :]
+                smpl_joints_3d = (
+                    smpl_joints[:, dataset.smpl_data.smpl_joint_ids_45_pixie, :] + 1.0
+                ) * 0.5
+                in_tensor["smpl_joint"] = smpl_joints[:,
+                                                      dataset.smpl_data.smpl_joint_ids_24_pixie, :]
 
                 ghum_lmks = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], :2].to(device)
                 ghum_conf = data["landmark"][:, SMPLX_object.ghum_smpl_pairs[:, 0], -1].to(device)
@@ -227,7 +258,7 @@ if __name__ == "__main__":
                 T_mask_F, T_mask_B = dataset.render.get_image(type="mask")
 
                 with torch.no_grad():
-                    in_tensor["normal_F"], in_tensor["normal_B"] = normal_model.netG(in_tensor)
+                    in_tensor["normal_F"], in_tensor["normal_B"] = normal_net.netG(in_tensor)
 
                 diff_F_smpl = torch.abs(in_tensor["T_normal_F"] - in_tensor["normal_F"])
                 diff_B_smpl = torch.abs(in_tensor["T_normal_B"] - in_tensor["normal_B"])
@@ -249,25 +280,37 @@ 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])
+                body_overlap = (gt_arr * smpl_arr_fake.gt(0.0)
+                               ).sum(dim=[1, 2]) / smpl_arr_fake.gt(0.0).sum(dim=[1, 2])
                 body_overlap_mask = (gt_arr * smpl_arr_fake).unsqueeze(1)
                 body_overlap_flag = body_overlap < cfg.body_overlap_thres
 
-                losses["normal"]["value"] = (diff_F_smpl * body_overlap_mask[..., :512] + diff_B_smpl * body_overlap_mask[..., 512:]).mean() / 2.0
+                losses["normal"]["value"] = (
+                    diff_F_smpl * body_overlap_mask[..., :512] +
+                    diff_B_smpl * body_overlap_mask[..., 512:]
+                ).mean() / 2.0
 
                 losses["silhouette"]["weight"] = [0 if flag else 1.0 for flag in body_overlap_flag]
                 occluded_idx = torch.where(body_overlap_flag)[0]
                 ghum_conf[occluded_idx] *= ghum_conf[occluded_idx] > 0.95
-                losses["joint"]["value"] = (torch.norm(ghum_lmks - smpl_lmks, dim=2) * ghum_conf).mean(dim=1)
+                losses["joint"]["value"] = (torch.norm(ghum_lmks - smpl_lmks, dim=2) *
+                                            ghum_conf).mean(dim=1)
 
                 # Weighted sum of the losses
                 smpl_loss = 0.0
-                pbar_desc = "Body Fitting --- "
+                pbar_desc = "Body Fitting -- "
                 for k in ["normal", "silhouette", "joint"]:
-                    per_loop_loss = (losses[k]["value"] * torch.tensor(losses[k]["weight"]).to(device)).mean()
+                    per_loop_loss = (
+                        losses[k]["value"] * torch.tensor(losses[k]["weight"]).to(device)
+                    ).mean()
                     pbar_desc += f"{k}: {per_loop_loss:.3f} | "
                     smpl_loss += per_loop_loss
                 pbar_desc += f"Total: {smpl_loss:.3f}"
@@ -279,19 +322,25 @@ if __name__ == "__main__":
                 # save intermediate results / vis_freq and final_step
                 if (i % args.vis_freq == 0) or (i == args.loop_smpl - 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"))
+                    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")
+                    )
 
                 smpl_loss.backward()
                 optimizer_smpl.step()
@@ -304,14 +353,21 @@ if __name__ == "__main__":
         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)
+                [
+                    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)
         rgb_norm_B = blend_rgb_norm(in_tensor["normal_B"], data)
 
         img_overlap_path = osp.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png")
-        torchvision.utils.save_image(torch.Tensor([data["img_raw"], rgb_norm_F, rgb_norm_B]).permute(0, 3, 1, 2) / 255., img_overlap_path)
+        torchvision.utils.save_image(
+            torch.cat([data["img_raw"], rgb_norm_F, rgb_norm_B], dim=-1) / 255., img_overlap_path
+        )
 
         smpl_obj_lst = []
 
@@ -329,15 +385,28 @@ if __name__ == "__main__":
             if not osp.exists(smpl_obj_path):
                 smpl_obj.export(smpl_obj_path)
                 smpl_info = {
-                    "betas": optimed_betas[idx].detach().cpu().unsqueeze(0),
-                    "body_pose": 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),
-                    "transl": optimed_trans[idx].detach().cpu(),
-                    "expression": data["exp"][idx].cpu().unsqueeze(0),
-                    "jaw_pose": 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),
-                    "right_hand_pose": rotation_matrix_to_angle_axis(data["right_hand_pose"][idx]).cpu().unsqueeze(0),
-                    "scale": data["scale"][idx].cpu(),
+                    "betas":
+                        optimed_betas[idx].detach().cpu().unsqueeze(0),
+                    "body_pose":
+                        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),
+                    "transl":
+                        optimed_trans[idx].detach().cpu(),
+                    "expression":
+                        data["exp"][idx].cpu().unsqueeze(0),
+                    "jaw_pose":
+                        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),
+                    "right_hand_pose":
+                        rotation_matrix_to_angle_axis(data["right_hand_pose"][idx]
+                                                     ).cpu().unsqueeze(0),
+                    "scale":
+                        data["scale"][idx].cpu(),
                 }
                 np.save(
                     smpl_obj_path.replace(".obj", ".npy"),
@@ -359,10 +428,13 @@ 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(batch_smpl_verts, batch_smpl_faces)
+        in_tensor["depth_F"], in_tensor["depth_B"] = dataset.render_depth(
+            batch_smpl_verts, batch_smpl_faces
+        )
 
         per_loop_lst = []
 
@@ -389,7 +461,13 @@ if __name__ == "__main__":
             )
 
             # BNI process
-            BNI_object = BNI(dir_path=osp.join(args.out_dir, cfg.name, "BNI"), name=data["name"], BNI_dict=BNI_dict, cfg=cfg.bni, device=device)
+            BNI_object = BNI(
+                dir_path=osp.join(args.out_dir, cfg.name, "BNI"),
+                name=data["name"],
+                BNI_dict=BNI_dict,
+                cfg=cfg.bni,
+                device=device
+            )
 
             BNI_object.extract_surface(False)
 
@@ -406,29 +484,40 @@ if __name__ == "__main__":
                 side_mesh = apply_face_mask(side_mesh, ~SMPLX_object.smplx_eyeball_fid_mask)
 
                 # 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)}))
+                in_tensor.update(
+                    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=[
                     0,
                 ] * 3, scale=2.0).data.transpose(2, 1, 0)
                 occupancies = np.flip(occupancies, axis=1)
 
-                in_tensor["body_voxels"] = torch.tensor(occupancies.copy()).float().unsqueeze(0).to(device)
+                in_tensor["body_voxels"] = torch.tensor(occupancies.copy()
+                                                       ).float().unsqueeze(0).to(device)
 
                 with torch.no_grad():
-                    sdf = ifnet_model.reconEngine(netG=ifnet_model.netG, batch=in_tensor)
-                    verts_IF, faces_IF = ifnet_model.reconEngine.export_mesh(sdf)
+                    sdf = ifnet.reconEngine(netG=ifnet.netG, batch=in_tensor)
+                    verts_IF, faces_IF = ifnet.reconEngine.export_mesh(sdf)
 
-                if ifnet_model.clean_mesh_flag:
+                if ifnet.clean_mesh_flag:
                     verts_IF, faces_IF = clean_mesh(verts_IF, faces_IF)
 
                 side_mesh = trimesh.Trimesh(verts_IF, faces_IF)
-                side_mesh = remesh(side_mesh, side_mesh_path)
+                side_mesh = remesh_laplacian(side_mesh, side_mesh_path)
 
             else:
                 side_mesh = apply_vertex_mask(
                     side_mesh,
-                    (SMPLX_object.front_flame_vertex_mask + SMPLX_object.mano_vertex_mask + SMPLX_object.eyeball_vertex_mask).eq(0).float(),
+                    (
+                        SMPLX_object.front_flame_vertex_mask + SMPLX_object.mano_vertex_mask +
+                        SMPLX_object.eyeball_vertex_mask
+                    ).eq(0).float(),
                 )
 
                 #register side_mesh to BNI surfaces
@@ -448,7 +537,9 @@ if __name__ == "__main__":
             # 3. remove eyeball faces
 
             # export intermediate meshes
-            BNI_object.F_B_trimesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj")
+            BNI_object.F_B_trimesh.export(
+                f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj"
+            )
             full_lst = []
 
             if "face" in cfg.bni.use_smpl:
@@ -458,37 +549,63 @@ if __name__ == "__main__":
                 face_mesh.vertices = face_mesh.vertices - np.array([0, 0, cfg.bni.thickness])
 
                 # remove face neighbor triangles
-                BNI_object.F_B_trimesh = part_removal(BNI_object.F_B_trimesh, face_mesh, cfg.bni.face_thres, device, smplx_mesh, region="face")
-                side_mesh = part_removal(side_mesh, face_mesh, cfg.bni.face_thres, device, smplx_mesh, region="face")
+                BNI_object.F_B_trimesh = part_removal(
+                    BNI_object.F_B_trimesh,
+                    face_mesh,
+                    cfg.bni.face_thres,
+                    device,
+                    smplx_mesh,
+                    region="face"
+                )
+                side_mesh = part_removal(
+                    side_mesh, face_mesh, cfg.bni.face_thres, device, smplx_mesh, region="face"
+                )
                 face_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_face.obj")
                 full_lst += [face_mesh]
 
             if "hand" in cfg.bni.use_smpl and (True in data['hands_visibility'][idx]):
 
-                hand_mask = torch.zeros(SMPLX_object.smplx_verts.shape[0],)
+                hand_mask = torch.zeros(SMPLX_object.smplx_verts.shape[0], )
                 if data['hands_visibility'][idx][0]:
-                    hand_mask.index_fill_(0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["left_hand"]), 1.0)
+                    hand_mask.index_fill_(
+                        0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["left_hand"]), 1.0
+                    )
                 if data['hands_visibility'][idx][1]:
-                    hand_mask.index_fill_(0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["right_hand"]), 1.0)
+                    hand_mask.index_fill_(
+                        0, torch.tensor(SMPLX_object.smplx_mano_vid_dict["right_hand"]), 1.0
+                    )
 
                 # only hands
                 hand_mesh = apply_vertex_mask(hand_mesh, hand_mask)
 
                 # remove hand neighbor triangles
-                BNI_object.F_B_trimesh = part_removal(BNI_object.F_B_trimesh, hand_mesh, cfg.bni.hand_thres, device, smplx_mesh, region="hand")
-                side_mesh = part_removal(side_mesh, hand_mesh, cfg.bni.hand_thres, device, smplx_mesh, region="hand")
+                BNI_object.F_B_trimesh = part_removal(
+                    BNI_object.F_B_trimesh,
+                    hand_mesh,
+                    cfg.bni.hand_thres,
+                    device,
+                    smplx_mesh,
+                    region="hand"
+                )
+                side_mesh = part_removal(
+                    side_mesh, hand_mesh, cfg.bni.hand_thres, device, smplx_mesh, region="hand"
+                )
                 hand_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_hand.obj")
                 full_lst += [hand_mesh]
 
             full_lst += [BNI_object.F_B_trimesh]
 
             # initial side_mesh could be SMPLX or IF-net
-            side_mesh = part_removal(side_mesh, sum(full_lst), 2e-2, device, smplx_mesh, region="", clean=False)
+            side_mesh = part_removal(
+                side_mesh, sum(full_lst), 2e-2, device, smplx_mesh, region="", clean=False
+            )
 
             full_lst += [side_mesh]
 
             # # export intermediate meshes
-            BNI_object.F_B_trimesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj")
+            BNI_object.F_B_trimesh.export(
+                f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_BNI.obj"
+            )
             side_mesh.export(f"{args.out_dir}/{cfg.name}/obj/{data['name']}_{idx}_side.obj")
 
             if cfg.bni.use_poisson:
@@ -505,15 +622,22 @@ if __name__ == "__main__":
             rotate_recon_lst = dataset.render.get_image(cam_type="four")
             per_loop_lst.extend([in_tensor['image'][idx:idx + 1]] + rotate_recon_lst)
 
-            # coloring the final mesh
-            final_colors = query_color(
-                torch.tensor(final_mesh.vertices).float(),
-                torch.tensor(final_mesh.faces).long(),
-                in_tensor["image"][idx:idx + 1],
-                device=device,
-            )
-            final_mesh.visual.vertex_colors = final_colors
-            final_mesh.export(final_path)
+            if cfg.bni.texture_src == 'image':
+                
+                # coloring the final mesh (front: RGB pixels, back: normal colors)
+                final_colors = query_color(
+                    torch.tensor(final_mesh.vertices).float(),
+                    torch.tensor(final_mesh.faces).long(),
+                    in_tensor["image"][idx:idx + 1],
+                    device=device,
+                )
+                final_mesh.visual.vertex_colors = final_colors
+                final_mesh.export(final_path)
+                
+            elif cfg.bni.texture_src == 'SD':
+                
+                # !TODO: add texture from Stable Diffusion
+                pass
 
             # for video rendering
             in_tensor["BNI_verts"].append(torch.tensor(final_mesh.vertices).float())

apps/multi_render.py

@@ -20,6 +20,4 @@ faces_lst = in_tensor["body_faces"] + in_tensor["BNI_faces"]
 
 # self-rotated video
 render.load_meshes(verts_lst, faces_lst)
-render.get_rendered_video_multi(
-    in_tensor,
-    f"{root}/{args.name}_cloth.mp4")
+render.get_rendered_video_multi(in_tensor, f"{root}/{args.name}_cloth.mp4")

configs/econ.yaml

@@ -35,3 +35,4 @@ bni:
   face_thres: 6e-2
   thickness: 0.02
   hps_type: "pixie"
+  texture_src: "SD"

docs/installation.md

@@ -9,12 +9,11 @@ cd ECON
 
 ## Environment
 
-- Ubuntu 20 / 18
-- GCC=7 (required by [pypoisson](https://github.com/mmolero/pypoisson/issues/13))
+- Ubuntu 20 / 18, (Windows as well, see [issue#7](https://github.com/YuliangXiu/ECON/issues/7))
 - **CUDA=11.4, GPU Memory > 12GB**
 - Python = 3.8
 - PyTorch >= 1.13.0 (official [Get Started](https://pytorch.org/get-started/locally/))
-- CUPY >= 11.3.0 (offcial [Installation](https://docs.cupy.dev/en/stable/install.html#installing-cupy-from-pypi))
+- Cupy >= 11.3.0 (offcial [Installation](https://docs.cupy.dev/en/stable/install.html#installing-cupy-from-pypi))
 - PyTorch3D (official [INSTALL.md](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md), recommend [install-from-local-clone](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#2-install-from-a-local-clone))
 
 ```bash

lib/common/BNI.py

@@ -1,12 +1,12 @@
-from lib.common.BNI_utils import (verts_inverse_transform, depth_inverse_transform,
-                                  double_side_bilateral_normal_integration)
+from lib.common.BNI_utils import (
+    verts_inverse_transform, depth_inverse_transform, double_side_bilateral_normal_integration
+)
 
 import torch
 import trimesh
 
 
 class BNI:
-
     def __init__(self, dir_path, name, BNI_dict, cfg, device):
 
         self.scale = 256.0
@@ -64,22 +64,20 @@ class BNI:
 
         F_B_verts = torch.cat((F_verts, B_verts), dim=0)
         F_B_faces = torch.cat(
-            (bni_result["F_faces"], bni_result["B_faces"] + bni_result["F_faces"].max() + 1), dim=0)
+            (bni_result["F_faces"], bni_result["B_faces"] + bni_result["F_faces"].max() + 1), dim=0
+        )
 
-        self.F_B_trimesh = trimesh.Trimesh(F_B_verts.float(),
-                                           F_B_faces.long(),
-                                           process=False,
-                                           maintain_order=True)
+        self.F_B_trimesh = trimesh.Trimesh(
+            F_B_verts.float(), F_B_faces.long(), process=False, maintain_order=True
+        )
 
-        self.F_trimesh = trimesh.Trimesh(F_verts.float(),
-                                         bni_result["F_faces"].long(),
-                                         process=False,
-                                         maintain_order=True)
+        self.F_trimesh = trimesh.Trimesh(
+            F_verts.float(), bni_result["F_faces"].long(), process=False, maintain_order=True
+        )
 
-        self.B_trimesh = trimesh.Trimesh(B_verts.float(),
-                                         bni_result["B_faces"].long(),
-                                         process=False,
-                                         maintain_order=True)
+        self.B_trimesh = trimesh.Trimesh(
+            B_verts.float(), bni_result["B_faces"].long(), process=False, maintain_order=True
+        )
 
 
 if __name__ == "__main__":
@@ -93,16 +91,18 @@ if __name__ == "__main__":
     bni_dict = np.load(npy_file, allow_pickle=True).item()
 
     default_cfg = {'k': 2, 'lambda1': 1e-4, 'boundary_consist': 1e-6}
-    
+
     # for k in [1, 2, 4, 10, 100]:
     #     default_cfg['k'] = k
     # for k in [1e-8, 1e-4, 1e-2, 1e-1, 1]:
-        # default_cfg['lambda1'] = k
+    # default_cfg['lambda1'] = k
     # for k in [1e-4, 1e-2, 0]:
-        # default_cfg['boundary_consist'] = k
-        
-    bni_object = BNI(osp.dirname(npy_file), osp.basename(npy_file), bni_dict, default_cfg,
-                    torch.device('cuda:0'))
+    # default_cfg['boundary_consist'] = k
+
+    bni_object = BNI(
+        osp.dirname(npy_file), osp.basename(npy_file), bni_dict, default_cfg,
+        torch.device('cuda:0')
+    )
 
     bni_object.extract_surface()
     bni_object.F_trimesh.export(osp.join(osp.dirname(npy_file), "F.obj"))

lib/common/BNI_utils.py

@@ -53,8 +53,9 @@ def find_contour(mask, method='all'):
 
         contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
     else:
-        contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE,
-                                       cv2.CHAIN_APPROX_SIMPLE)
+        contours, _ = cv2.findContours(
+            mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+        )
 
     contour_cloth = np.array(find_max_list(contours))[:, 0, :]
 
@@ -67,16 +68,19 @@ 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])
     body_tan_half = np.sqrt(
-        (1. - np.clip(body_cos, a_max=1., a_min=-1.)) / np.clip(1. + body_cos, 1e-6, 2.))
+        (1. - np.clip(body_cos, a_max=1., a_min=-1.)) / np.clip(1. + body_cos, 1e-6, 2.)
+    )
 
     w = (body_tan_half + np.roll(body_tan_half, shift=1, axis=1)) / body_edges_a
     w /= w.sum(axis=1, keepdims=True)
@@ -97,16 +101,18 @@ def dispCorres(img_size, contour1, contour2, phi, dir_path):
     contour2 = contour2[None, :, None, :].astype(np.int32)
 
     disp = np.zeros((img_size, img_size, 3), dtype=np.uint8)
-    cv2.drawContours(disp, contour1, -1, (0, 255, 0), 1)  # green
-    cv2.drawContours(disp, contour2, -1, (255, 0, 0), 1)  # blue
+    cv2.drawContours(disp, contour1, -1, (0, 255, 0), 1)    # green
+    cv2.drawContours(disp, contour2, -1, (255, 0, 0), 1)    # blue
 
-    for i in range(contour1.shape[1]):  # do not show all the points when display
+    for i in range(contour1.shape[1]):    # do not show all the points when display
         # cv2.circle(disp, (contour1[0, i, 0, 0], contour1[0, i, 0, 1]), 1,
         #            (255, 0, 0), -1)
         corresPoint = contour2[0, phi[i], 0]
         # cv2.circle(disp, (corresPoint[0], corresPoint[1]), 1, (0, 255, 0), -1)
-        cv2.line(disp, (contour1[0, i, 0, 0], contour1[0, i, 0, 1]),
-                 (corresPoint[0], corresPoint[1]), (255, 255, 255), 1)
+        cv2.line(
+            disp, (contour1[0, i, 0, 0], contour1[0, i, 0, 1]), (corresPoint[0], corresPoint[1]),
+            (255, 255, 255), 1
+        )
 
     cv2.imwrite(osp.join(dir_path, "corres.png"), disp)
 
@@ -162,7 +168,8 @@ def verts_transform(t, depth_scale):
     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])
+        [0.0, 0.0, depth_scale]
+    )
 
     return t_copy
 
@@ -328,19 +335,22 @@ def construct_facets_from(mask):
     facet_move_top_mask = move_top(mask)
     facet_move_left_mask = move_left(mask)
     facet_move_top_left_mask = move_top_left(mask)
-    facet_top_left_mask = (facet_move_top_mask * facet_move_left_mask * facet_move_top_left_mask *
-                           mask)
+    facet_top_left_mask = (
+        facet_move_top_mask * facet_move_left_mask * facet_move_top_left_mask * mask
+    )
     facet_top_right_mask = move_right(facet_top_left_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):
@@ -364,8 +374,8 @@ def map_depth_map_to_point_clouds(depth_map, mask, K=None, step_size=1):
         u[..., 0] = xx
         u[..., 1] = yy
         u[..., 2] = 1
-        u = u[mask].T  # 3 x m
-        vertices = (cp.linalg.inv(K) @ u).T * depth_map[mask, cp.newaxis]  # m x 3
+        u = u[mask].T    # 3 x m
+        vertices = (cp.linalg.inv(K) @ u).T * depth_map[mask, cp.newaxis]    # m x 3
 
     return vertices
 
@@ -374,7 +384,6 @@ def sigmoid(x, k=1):
     return 1 / (1 + cp.exp(-k * x))
 
 
-
 def boundary_excluded_mask(mask):
     top_mask = cp.pad(mask, ((1, 0), (0, 0)), "constant", constant_values=0)[:-1, :]
     bottom_mask = cp.pad(mask, ((0, 1), (0, 0)), "constant", constant_values=0)[1:, :]
@@ -410,22 +419,24 @@ def create_boundary_matrix(mask):
     return B, B_full
 
 
-def double_side_bilateral_normal_integration(normal_front,
-                                             normal_back,
-                                             normal_mask,
-                                             depth_front=None,
-                                             depth_back=None,
-                                             depth_mask=None,
-                                             k=2,
-                                             lambda_normal_back=1,
-                                             lambda_depth_front=1e-4,
-                                             lambda_depth_back=1e-2,
-                                             lambda_boundary_consistency=1,
-                                             step_size=1,
-                                             max_iter=150,
-                                             tol=1e-4,
-                                             cg_max_iter=5000,
-                                             cg_tol=1e-3):
+def double_side_bilateral_normal_integration(
+    normal_front,
+    normal_back,
+    normal_mask,
+    depth_front=None,
+    depth_back=None,
+    depth_mask=None,
+    k=2,
+    lambda_normal_back=1,
+    lambda_depth_front=1e-4,
+    lambda_depth_back=1e-2,
+    lambda_boundary_consistency=1,
+    step_size=1,
+    max_iter=150,
+    tol=1e-4,
+    cg_max_iter=5000,
+    cg_tol=1e-3
+):
 
     # To avoid confusion, we list the coordinate systems in this code as follows
     #
@@ -467,14 +478,12 @@ def double_side_bilateral_normal_integration(normal_front,
     del normal_map_back
 
     # right, left, top, bottom
-    A3_f, A4_f, A1_f, A2_f = generate_dx_dy(normal_mask,
-                                            nz_horizontal=nz_front,
-                                            nz_vertical=nz_front,
-                                            step_size=step_size)
-    A3_b, A4_b, A1_b, A2_b = generate_dx_dy(normal_mask,
-                                            nz_horizontal=nz_back,
-                                            nz_vertical=nz_back,
-                                            step_size=step_size)
+    A3_f, A4_f, A1_f, A2_f = generate_dx_dy(
+        normal_mask, nz_horizontal=nz_front, nz_vertical=nz_front, step_size=step_size
+    )
+    A3_b, A4_b, A1_b, A2_b = generate_dx_dy(
+        normal_mask, nz_horizontal=nz_back, nz_vertical=nz_back, step_size=step_size
+    )
 
     has_left_mask = cp.logical_and(move_right(normal_mask), normal_mask)
     has_right_mask = cp.logical_and(move_left(normal_mask), normal_mask)
@@ -498,29 +507,25 @@ def double_side_bilateral_normal_integration(normal_front,
     b_back = cp.concatenate((-nx_back, -nx_back, -ny_back, -ny_back))
 
     # initialization
-    W_front = spdiags(0.5 * cp.ones(4 * num_normals),
-                      0,
-                      4 * num_normals,
-                      4 * num_normals,
-                      format="csr")
-    W_back = spdiags(0.5 * cp.ones(4 * num_normals),
-                     0,
-                     4 * num_normals,
-                     4 * num_normals,
-                     format="csr")
+    W_front = spdiags(
+        0.5 * cp.ones(4 * num_normals), 0, 4 * num_normals, 4 * num_normals, format="csr"
+    )
+    W_back = spdiags(
+        0.5 * cp.ones(4 * num_normals), 0, 4 * num_normals, 4 * num_normals, format="csr"
+    )
 
     z_front = cp.zeros(num_normals, float)
     z_back = cp.zeros(num_normals, float)
     z_combined = cp.concatenate((z_front, z_back))
 
     B, B_full = create_boundary_matrix(normal_mask)
-    B_mat = lambda_boundary_consistency * coo_matrix(B_full.get().T @ B_full.get())  #bug
+    B_mat = lambda_boundary_consistency * coo_matrix(B_full.get().T @ B_full.get())    #bug
 
     energy_list = []
 
     if depth_mask is not None:
-        depth_mask_flat = depth_mask[normal_mask].astype(bool)  # shape: (num_normals,)
-        z_prior_front = depth_map_front[normal_mask]  # shape: (num_normals,)
+        depth_mask_flat = depth_mask[normal_mask].astype(bool)    # shape: (num_normals,)
+        z_prior_front = depth_map_front[normal_mask]    # shape: (num_normals,)
         z_prior_front[~depth_mask_flat] = 0
         z_prior_back = depth_map_back[normal_mask]
         z_prior_back[~depth_mask_flat] = 0
@@ -554,40 +559,43 @@ def double_side_bilateral_normal_integration(normal_front,
                                  vstack((csr_matrix((num_normals, num_normals)), A_mat_back))]) + B_mat
         b_vec_combined = cp.concatenate((b_vec_front, b_vec_back))
 
-        D = spdiags(1 / cp.clip(A_mat_combined.diagonal(), 1e-5, None), 0, 2 * num_normals,
-                    2 * num_normals, "csr")  # Jacob preconditioner
+        D = spdiags(
+            1 / cp.clip(A_mat_combined.diagonal(), 1e-5, None), 0, 2 * num_normals, 2 * num_normals,
+            "csr"
+        )    # Jacob preconditioner
 
-        z_combined, _ = cg(A_mat_combined,
-                           b_vec_combined,
-                           M=D,
-                           x0=z_combined,
-                           maxiter=cg_max_iter,
-                           tol=cg_tol)
+        z_combined, _ = cg(
+            A_mat_combined, b_vec_combined, M=D, x0=z_combined, maxiter=cg_max_iter, tol=cg_tol
+        )
         z_front = z_combined[:num_normals]
         z_back = z_combined[num_normals:]
-        wu_f = sigmoid((A2_f.dot(z_front))**2 - (A1_f.dot(z_front))**2, k)  # top
-        wv_f = sigmoid((A4_f.dot(z_front))**2 - (A3_f.dot(z_front))**2, k)  # right
+        wu_f = sigmoid((A2_f.dot(z_front))**2 - (A1_f.dot(z_front))**2, k)    # top
+        wv_f = sigmoid((A4_f.dot(z_front))**2 - (A3_f.dot(z_front))**2, k)    # right
         wu_f[top_boundnary_mask] = 0.5
         wu_f[bottom_boundary_mask] = 0.5
         wv_f[left_boundary_mask] = 0.5
         wv_f[right_boudnary_mask] = 0.5
-        W_front = spdiags(cp.concatenate((wu_f, 1 - wu_f, wv_f, 1 - wv_f)),
-                          0,
-                          4 * num_normals,
-                          4 * num_normals,
-                          format="csr")
-
-        wu_b = sigmoid((A2_b.dot(z_back))**2 - (A1_b.dot(z_back))**2, k)  # top
-        wv_b = sigmoid((A4_b.dot(z_back))**2 - (A3_b.dot(z_back))**2, k)  # right
+        W_front = spdiags(
+            cp.concatenate((wu_f, 1 - wu_f, wv_f, 1 - wv_f)),
+            0,
+            4 * num_normals,
+            4 * num_normals,
+            format="csr"
+        )
+
+        wu_b = sigmoid((A2_b.dot(z_back))**2 - (A1_b.dot(z_back))**2, k)    # top
+        wv_b = sigmoid((A4_b.dot(z_back))**2 - (A3_b.dot(z_back))**2, k)    # right
         wu_b[top_boundnary_mask] = 0.5
         wu_b[bottom_boundary_mask] = 0.5
         wv_b[left_boundary_mask] = 0.5
         wv_b[right_boudnary_mask] = 0.5
-        W_back = spdiags(cp.concatenate((wu_b, 1 - wu_b, wv_b, 1 - wv_b)),
-                         0,
-                         4 * num_normals,
-                         4 * num_normals,
-                         format="csr")
+        W_back = spdiags(
+            cp.concatenate((wu_b, 1 - wu_b, wv_b, 1 - wv_b)),
+            0,
+            4 * num_normals,
+            4 * num_normals,
+            format="csr"
+        )
 
         energy_old = energy
         energy = (A_front_data @ z_front - b_front).T @ W_front @ (A_front_data @ z_front - b_front) + \
@@ -603,23 +611,26 @@ def double_side_bilateral_normal_integration(normal_front,
         if relative_energy < tol:
             break
     # del A1, A2, A3, A4, nx, ny
-    
+
     depth_map_front_est = cp.ones_like(normal_mask, float) * cp.nan
     depth_map_front_est[normal_mask] = z_front
 
     depth_map_back_est = cp.ones_like(normal_mask, float) * cp.nan
     depth_map_back_est[normal_mask] = z_back
-    
+
     # manually cut the intersection
-    normal_mask[depth_map_front_est>=depth_map_back_est] = False
+    normal_mask[depth_map_front_est >= depth_map_back_est] = False
     depth_map_front_est[~normal_mask] = cp.nan
     depth_map_back_est[~normal_mask] = cp.nan
 
     vertices_front = cp.asnumpy(
-        map_depth_map_to_point_clouds(depth_map_front_est, normal_mask, K=None,
-                                      step_size=step_size))
+        map_depth_map_to_point_clouds(
+            depth_map_front_est, normal_mask, K=None, step_size=step_size
+        )
+    )
     vertices_back = cp.asnumpy(
-        map_depth_map_to_point_clouds(depth_map_back_est, normal_mask, K=None, step_size=step_size))
+        map_depth_map_to_point_clouds(depth_map_back_est, normal_mask, K=None, step_size=step_size)
+    )
 
     facets_back = cp.asnumpy(construct_facets_from(normal_mask))
 
@@ -656,7 +667,7 @@ def save_normal_tensor(in_tensor, idx, png_path, thickness=0.0):
     depth_B_arr = depth2arr(in_tensor["depth_B"][idx])
 
     BNI_dict = {}
-    
+
     # clothed human
     BNI_dict["normal_F"] = normal_F_arr
     BNI_dict["normal_B"] = normal_B_arr

lib/common/blender_utils.py

@@ -3,6 +3,7 @@ import sys, os
 from math import radians
 import mathutils
 import bmesh
+
 print(sys.exec_prefix)
 from tqdm import tqdm
 import numpy as np
@@ -29,7 +30,6 @@ shadows = False
 # diffuse_color = (18/255., 139/255., 142/255.,1)     #correct
 # diffuse_color = (251/255., 60/255., 60/255.,1)    #wrong
 
-
 smooth = False
 
 wireframe = False
@@ -47,13 +47,16 @@ compositor_alpha = 0.7
 # Helper functions
 ##################################################
 
+
 def blender_print(*args, **kwargs):
-    print (*args, **kwargs, file=sys.stderr)
+    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'''
 
@@ -110,8 +113,10 @@ def setup_diffuse_transparent_material(target, color, object_transparent, backfa
     links.new(node_mix_backface.outputs[0], node_output.inputs[0])
     return
 
+
 ##################################################
 
+
 def setup_scene():
     global render
     global cycles_gpu
@@ -150,12 +155,13 @@ def setup_scene():
         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')
+                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()
-
+                (cuda_devices, opencl_devices
+                ) = bpy.context.preferences.addons['cycles'].preferences.get_devices()
 
             if (len(cuda_devices) < 1):
                 print('ERROR: CUDA GPU acceleration not available')
@@ -178,7 +184,7 @@ def setup_scene():
             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
@@ -226,7 +232,6 @@ def setup_scene():
             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
@@ -245,8 +250,10 @@ def setup_scene():
         # Output transparent image when no background is used
         scene.render.image_settings.color_mode = 'RGBA'
 
+
 ##################################################
 
+
 def setup_compositing():
 
     global compositor_image_scale
@@ -275,6 +282,7 @@ def setup_compositing():
 
     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
@@ -288,13 +296,13 @@ def render_file(input_file, input_dir, output_file, output_dir, yaw, correct):
     # 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_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')
@@ -309,11 +317,11 @@ def render_file(input_file, input_dir, output_file, output_dir, yaw, correct):
     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
+        diffuse_color = (18 / 255., 139 / 255., 142 / 255., 1)    #correct
     else:
-        diffuse_color = (251/255., 60/255., 60/255.,1)    #wrong
+        diffuse_color = (251 / 255., 60 / 255., 60 / 255., 1)    #wrong
 
     setup_diffuse_transparent_material(object, diffuse_color, object_transparent, mouth_transparent)
 
@@ -336,10 +344,10 @@ def render_file(input_file, input_dir, output_file, output_dir, yaw, correct):
     bpy.ops.render.render(write_still=True)
 
     # Remove temporary output redirection
-#    sys.stdout.flush()
-#    os.close(1)
-#    os.dup(old)
-#    os.close(old)
+    #    sys.stdout.flush()
+    #    os.close(1)
+    #    os.dup(old)
+    #    os.close(old)
 
     # Delete last selected object from scene
     object.select_set(True)
@@ -351,7 +359,7 @@ def process_file(input_file, input_dir, output_file, output_dir, correct=True):
     global quality_preview
 
     if not input_file.endswith('.obj'):
-        print('ERROR: Invalid input: ' + input_file )
+        print('ERROR: Invalid input: ' + input_file)
         return
 
     print('Processing: ' + input_file)
@@ -361,7 +369,7 @@ def process_file(input_file, input_dir, output_file, output_dir, correct=True):
     if quality_preview:
         output_file = output_file.replace('.png', '-preview.png')
 
-    angle = 360.0/views
+    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)}")
@@ -369,8 +377,7 @@ def process_file(input_file, input_dir, output_file, output_dir, correct=True):
         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

@@ -36,11 +36,13 @@ 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
 
@@ -56,9 +58,8 @@ def smpl_to_recon_labels(recon, smpl, k=1):
         Returns a dictionary containing the bodypart and the corresponding indices
     """
     smpl_vert_segmentation = json.load(
-        open(
-            os.path.join(os.path.dirname(__file__),
-                         "smpl_vert_segmentation.json")))
+        open(os.path.join(os.path.dirname(__file__), "smpl_vert_segmentation.json"))
+    )
     n = smpl.vertices.shape[0]
     y = np.array([None] * n)
     for key, val in smpl_vert_segmentation.items():
@@ -71,8 +72,7 @@ def smpl_to_recon_labels(recon, smpl, k=1):
 
     recon_labels = {}
     for key in smpl_vert_segmentation.keys():
-        recon_labels[key] = list(
-            np.argwhere(y_pred == key).flatten().astype(int))
+        recon_labels[key] = list(np.argwhere(y_pred == key).flatten().astype(int))
 
     return recon_labels
 
@@ -139,8 +139,7 @@ def extract_cloth(recon, segmentation, K, R, t, smpl=None):
         if type == 1 or type == 3 or type == 10:
             body_parts_to_remove += ["leftForeArm", "rightForeArm"]
         # No sleeves at all or lower body clothes
-        elif (type == 5 or type == 6 or type == 12 or type == 13 or type == 8
-              or type == 9):
+        elif (type == 5 or type == 6 or type == 12 or type == 13 or type == 8 or type == 9):
             body_parts_to_remove += [
                 "leftForeArm",
                 "rightForeArm",
@@ -159,8 +158,8 @@ def extract_cloth(recon, segmentation, K, R, t, smpl=None):
             ]
 
         verts_to_remove = list(
-            itertools.chain.from_iterable(
-                [recon_labels[part] for part in body_parts_to_remove]))
+            itertools.chain.from_iterable([recon_labels[part] for part in body_parts_to_remove])
+        )
 
         label_mask = np.zeros(num_verts, dtype=bool)
         label_mask[verts_to_remove] = True

lib/common/config.py

@@ -100,6 +100,7 @@ _C.bni.thickness = 0.00
 _C.bni.hand_thres = 4e-2
 _C.bni.face_thres = 6e-2
 _C.bni.hps_type = "pixie"
+_C.bni.texture_src = "image"
 
 # kernel_size, stride, dilation, padding
 
@@ -170,10 +171,10 @@ _C.dataset.rp_type = "pifu900"
 _C.dataset.th_type = "train"
 _C.dataset.input_size = 512
 _C.dataset.rotation_num = 3
-_C.dataset.num_precomp = 10  # Number of segmentation classifiers
-_C.dataset.num_multiseg = 500  # Number of categories per classifier
-_C.dataset.num_knn = 10  # for loss/error
-_C.dataset.num_knn_dis = 20  # for accuracy
+_C.dataset.num_precomp = 10    # Number of segmentation classifiers
+_C.dataset.num_multiseg = 500    # Number of categories per classifier
+_C.dataset.num_knn = 10    # for loss/error
+_C.dataset.num_knn_dis = 20    # for accuracy
 _C.dataset.num_verts_max = 20000
 _C.dataset.zray_type = False
 _C.dataset.online_smpl = False
@@ -210,8 +211,7 @@ def get_cfg_defaults():
 
 # Alternatively, provide a way to import the defaults as
 # a global singleton:
-cfg = _C  # users can `from config import cfg`
-
+cfg = _C    # users can `from config import cfg`
 
 # cfg = get_cfg_defaults()
 # cfg.merge_from_file('./configs/example.yaml')
@@ -244,9 +244,7 @@ def parse_args(args):
 def parse_args_extend(args):
     if args.resume:
         if not os.path.exists(args.log_dir):
-            raise ValueError(
-                "Experiment are set to resume mode, but log directory does not exist."
-            )
+            raise ValueError("Experiment are set to resume mode, but log directory does not exist.")
 
         # load log's cfg
         cfg_file = os.path.join(args.log_dir, "cfg.yaml")

lib/common/imutils.py

@@ -3,14 +3,13 @@ import mediapipe as mp
 import torch
 import numpy as np
 import torch.nn.functional as F
-from rembg import remove
-from rembg.session_factory import new_session
 from PIL import Image
-from torchvision.models import detection
-
 from lib.pymafx.core import constants
-from lib.common.cloth_extraction import load_segmentation
+
+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
 
 
 def transform_to_tensor(res, mean=None, std=None, is_tensor=False):
@@ -24,42 +23,40 @@ def transform_to_tensor(res, mean=None, std=None, is_tensor=False):
     return transforms.Compose(all_ops)
 
 
-def aug_matrix(w1, h1, w2, h2):
-    dx = (w2 - w1) / 2.0
-    dy = (h2 - h1) / 2.0
-
-    matrix_trans = np.array([[1.0, 0, dx], [0, 1.0, dy], [0, 0, 1.0]])
-
-    scale = np.min([float(w2) / w1, float(h2) / h1])
+def get_affine_matrix_wh(w1, h1, w2, h2):
 
-    M = get_affine_matrix(center=(w2 / 2.0, h2 / 2.0), translate=(0, 0), scale=scale)
-
-    M = np.array(M + [0.0, 0.0, 1.0]).reshape(3, 3)
-    M = M.dot(matrix_trans)
+    transl = torch.tensor([(w2 - w1) / 2.0, (h2 - h1) / 2.0]).unsqueeze(0)
+    center = torch.tensor([w1 / 2.0, h1 / 2.0]).unsqueeze(0)
+    scale = torch.min(torch.tensor([w2 / w1, h2 / h1])).repeat(2).unsqueeze(0)
+    M = get_affine_matrix2d(transl, center, scale, angle=torch.tensor([0.]))
 
     return M
 
 
-def get_affine_matrix(center, translate, scale):
-    cx, cy = center
-    tx, ty = translate
-
-    M = [1, 0, 0, 0, 1, 0]
-    M = [x * scale for x in M]
+def get_affine_matrix_box(boxes, w2, h2):
 
-    # Apply translation and of center translation: RSS * C^-1
-    M[2] += M[0] * (-cx) + M[1] * (-cy)
-    M[5] += M[3] * (-cx) + M[4] * (-cy)
+    # 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)
+    scale = torch.min(torch.tensor([w2 / width, h2 / height]),
+                      dim=0)[0].unsqueeze(1).repeat(1, 2) * 0.9    #(N,2)
+    transl = torch.tensor([w2 / 2.0 - center[:, 0], h2 / 2.0 - center[:, 1]]).unsqueeze(0)    #(N,2)
+    M = get_affine_matrix2d(transl, center, scale, angle=torch.tensor([0.]))
 
-    # Apply center translation: T * C * RSS * C^-1
-    M[2] += cx + tx
-    M[5] += cy + ty
     return M
 
 
 def load_img(img_file):
 
     img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
+
+    # considering 16-bit image
+    if img.dtype == np.uint16:
+        img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+
     if len(img.shape) == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
@@ -68,11 +65,10 @@ def load_img(img_file):
     else:
         img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
 
-    return img
+    return torch.tensor(img).permute(2, 0, 1).unsqueeze(0).float(), img.shape[:2]
 
 
 def get_keypoints(image):
-
     def collect_xyv(x, body=True):
         lmk = x.landmark
         all_lmks = []
@@ -84,8 +80,8 @@ def get_keypoints(image):
     mp_holistic = mp.solutions.holistic
 
     with mp_holistic.Holistic(
-            static_image_mode=True,
-            model_complexity=2,
+        static_image_mode=True,
+        model_complexity=2,
     ) as holistic:
         results = holistic.process(image)
 
@@ -93,9 +89,15 @@ def get_keypoints(image):
 
     result = {}
     result["body"] = collect_xyv(results.pose_landmarks) if results.pose_landmarks else fake_kps
-    result["lhand"] = collect_xyv(results.left_hand_landmarks, False) if results.left_hand_landmarks else fake_kps
-    result["rhand"] = collect_xyv(results.right_hand_landmarks, False) if results.right_hand_landmarks else fake_kps
-    result["face"] = collect_xyv(results.face_landmarks, False) if results.face_landmarks else fake_kps
+    result["lhand"] = collect_xyv(
+        results.left_hand_landmarks, False
+    ) if results.left_hand_landmarks else fake_kps
+    result["rhand"] = collect_xyv(
+        results.right_hand_landmarks, False
+    ) if results.right_hand_landmarks else fake_kps
+    result["face"] = collect_xyv(
+        results.face_landmarks, False
+    ) if results.face_landmarks else fake_kps
 
     return result
 
@@ -104,13 +106,21 @@ 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]}
+    item = {
+        'img_body':
+            F.interpolate(image.unsqueeze(0), size=224, mode='bicubic', align_corners=True)[0]
+    }
 
     for part in ['lhand', 'rhand', 'face']:
         kp2d = landmarks[part]
         kp2d_valid = kp2d[kp2d[:, 3] > 0.]
         if len(kp2d_valid) > 0:
-            bbox = [min(kp2d_valid[:, 0]), min(kp2d_valid[:, 1]), max(kp2d_valid[:, 0]), max(kp2d_valid[:, 1])]
+            bbox = [
+                min(kp2d_valid[:, 0]),
+                min(kp2d_valid[:, 1]),
+                max(kp2d_valid[:, 0]),
+                max(kp2d_valid[:, 1])
+            ]
             center_part = [(bbox[2] + bbox[0]) / 2., (bbox[3] + bbox[1]) / 2.]
             scale_part = 2. * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
 
@@ -141,20 +151,6 @@ def get_pymafx(image, landmarks):
     return item
 
 
-def expand_bbox(bbox, width, height, ratio=0.1):
-
-    bbox = np.around(bbox).astype(np.int16)
-    bbox_width = bbox[2] - bbox[0]
-    bbox_height = bbox[3] - bbox[1]
-
-    bbox[1] = max(bbox[1] - bbox_height * ratio, 0)
-    bbox[3] = min(bbox[3] + bbox_height * ratio, height)
-    bbox[0] = max(bbox[0] - bbox_width * ratio, 0)
-    bbox[2] = min(bbox[2] + bbox_width * ratio, width)
-
-    return bbox
-
-
 def remove_floats(mask):
 
     # 1. find all the contours
@@ -173,51 +169,48 @@ def remove_floats(mask):
     return new_mask
 
 
-def process_image(img_file, hps_type, single, input_res=512):
+def process_image(img_file, hps_type, single, input_res, detector):
 
-    img_raw = load_img(img_file)
-
-    in_height, in_width = img_raw.shape[:2]
-    M = aug_matrix(in_width, in_height, input_res * 2, input_res * 2)
-
-    # from rectangle to square by padding (input_res*2, input_res*2)
-    img_square = cv2.warpAffine(img_raw, M[0:2, :], (input_res * 2, input_res * 2), flags=cv2.INTER_CUBIC)
+    img_raw, (in_height, in_width) = load_img(img_file)
+    tgt_res = input_res * 2
+    M_square = get_affine_matrix_wh(in_width, in_height, tgt_res, tgt_res)
+    img_square = warp_affine(
+        img_raw,
+        M_square[:, :2], (tgt_res, ) * 2,
+        mode='bilinear',
+        padding_mode='zeros',
+        align_corners=True
+    )
 
     # detection for bbox
-    detector = detection.maskrcnn_resnet50_fpn(weights=detection.MaskRCNN_ResNet50_FPN_V2_Weights)
-    detector.eval()
-    predictions = detector([torch.from_numpy(img_square).permute(2, 0, 1) / 255.])[0]
+    predictions = detector(img_square / 255.)[0]
 
     if single:
         top_score = predictions["scores"][predictions["labels"] == 1].max()
         human_ids = torch.where(predictions["scores"] == top_score)[0]
     else:
-        human_ids = torch.logical_and(predictions["labels"] == 1, predictions["scores"] > 0.9).nonzero().squeeze(1)
+        human_ids = torch.logical_and(predictions["labels"] == 1,
+                                      predictions["scores"] > 0.9).nonzero().squeeze(1)
 
     boxes = predictions["boxes"][human_ids, :].detach().cpu().numpy()
     masks = predictions["masks"][human_ids, :, :].permute(0, 2, 3, 1).detach().cpu().numpy()
 
-    width = boxes[:, 2] - boxes[:, 0]  #(N,)
-    height = boxes[:, 3] - boxes[:, 1]  #(N,)
-    center = np.array([(boxes[:, 0] + boxes[:, 2]) / 2.0, (boxes[:, 1] + boxes[:, 3]) / 2.0]).T  #(N,2)
-    scale = np.array([width, height]).max(axis=0) / 90.
+    M_crop = get_affine_matrix_box(boxes, input_res, input_res)
 
     img_icon_lst = []
     img_crop_lst = []
     img_hps_lst = []
     img_mask_lst = []
-    uncrop_param_lst = []
     landmark_lst = []
     hands_visibility_lst = []
     img_pymafx_lst = []
 
     uncrop_param = {
-        "center": center,
-        "scale": scale,
         "ori_shape": [in_height, in_width],
         "box_shape": [input_res, input_res],
-        "crop_shape": [input_res * 2, input_res * 2, 3],
-        "M": M,
+        "square_shape": [tgt_res, tgt_res],
+        "M_square": M_square,
+        "M_crop": M_crop
     }
 
     for idx in range(len(boxes)):
@@ -228,59 +221,74 @@ def process_image(img_file, hps_type, single, input_res=512):
         else:
             mask_detection = masks[0] * 0.
 
-        img_crop, _ = crop(
-            np.concatenate([img_square, (mask_detection < 0.4) * 255], axis=2), center[idx], scale[idx], [input_res, input_res])
-
-        # get accurate segmentation mask of focus person
+        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),
+            M_crop[idx:idx + 1, :2], (input_res, ) * 2,
+            mode='bilinear',
+            padding_mode='zeros',
+            align_corners=True
+        ).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
+
+        # get accurate person segmentation mask
         img_rembg = remove(img_crop, post_process_mask=True, session=new_session("u2net"))
         img_mask = remove_floats(img_rembg[:, :, [3]])
 
-        # required image tensors / arrays
-
-        # img_icon  (tensor): (-1, 1),          [3,512,512]
-        # img_hps   (tensor): (-2.11, 2.44),    [3,224,224]
-
-        # img_np    (array): (0, 255),          [512,512,3]
-        # img_rembg (array): (0, 255),          [512,512,4]
-        # img_mask  (array): (0, 1),            [512,512,1]
-        # img_crop  (array): (0, 255),          [512,512,4]
-
         mean_icon = std_icon = (0.5, 0.5, 0.5)
         img_np = (img_rembg[..., :3] * img_mask).astype(np.uint8)
-        img_icon = transform_to_tensor(512, mean_icon, std_icon)(Image.fromarray(img_np)) * torch.tensor(img_mask).permute(
-            2, 0, 1)
-        img_hps = transform_to_tensor(224, constants.IMG_NORM_MEAN, constants.IMG_NORM_STD)(Image.fromarray(img_np))
+        img_icon = transform_to_tensor(512, mean_icon, std_icon)(
+            Image.fromarray(img_np)
+        ) * torch.tensor(img_mask).permute(2, 0, 1)
+        img_hps = transform_to_tensor(224, constants.IMG_NORM_MEAN,
+                                      constants.IMG_NORM_STD)(Image.fromarray(img_np))
 
         landmarks = get_keypoints(img_np)
 
+        # get hands visibility
+        hands_visibility = [True, True]
+        if landmarks['lhand'][:, -1].mean() == 0.:
+            hands_visibility[0] = False
+        if landmarks['rhand'][:, -1].mean() == 0.:
+            hands_visibility[1] = False
+        hands_visibility_lst.append(hands_visibility)
+
         if hps_type == 'pymafx':
             img_pymafx_lst.append(
                 get_pymafx(
-                    transform_to_tensor(512, constants.IMG_NORM_MEAN, constants.IMG_NORM_STD)(Image.fromarray(img_np)),
-                    landmarks))
+                    transform_to_tensor(512, constants.IMG_NORM_MEAN,
+                                        constants.IMG_NORM_STD)(Image.fromarray(img_np)), landmarks
+                )
+            )
 
         img_crop_lst.append(torch.tensor(img_crop).permute(2, 0, 1) / 255.0)
         img_icon_lst.append(img_icon)
         img_hps_lst.append(img_hps)
         img_mask_lst.append(torch.tensor(img_mask[..., 0]))
-        uncrop_param_lst.append(uncrop_param)
         landmark_lst.append(landmarks['body'])
 
-        hands_visibility = [True, True]
-        if landmarks['lhand'][:, -1].mean() == 0.:
-            hands_visibility[0] = False
-        if landmarks['rhand'][:, -1].mean() == 0.:
-            hands_visibility[1] = False
-        hands_visibility_lst.append(hands_visibility)
+    # required image tensors / arrays
+
+    # img_icon  (tensor): (-1, 1),          [3,512,512]
+    # img_hps   (tensor): (-2.11, 2.44),    [3,224,224]
+
+    # img_np    (array): (0, 255),          [512,512,3]
+    # img_rembg (array): (0, 255),          [512,512,4]
+    # img_mask  (array): (0, 1),            [512,512,1]
+    # img_crop  (array): (0, 255),          [512,512,4]
 
     return_dict = {
-        "img_icon": torch.stack(img_icon_lst).float(),  #[N, 3, res, res]
-        "img_crop": torch.stack(img_crop_lst).float(),  #[N, 4, res, res]               
-        "img_hps": torch.stack(img_hps_lst).float(),  #[N, 3, res, res]
-        "img_raw": img_raw,  #[H, W, 3]
-        "img_mask": torch.stack(img_mask_lst).float(),  #[N, res, res]
+        "img_icon": torch.stack(img_icon_lst).float(),    #[N, 3, res, res]
+        "img_crop": torch.stack(img_crop_lst).float(),    #[N, 4, res, res]               
+        "img_hps": torch.stack(img_hps_lst).float(),    #[N, 3, res, res]
+        "img_raw": img_raw,    #[1, 3, H, W]
+        "img_mask": torch.stack(img_mask_lst).float(),    #[N, res, res]
         "uncrop_param": uncrop_param,
-        "landmark": torch.stack(landmark_lst),  #[N, 33, 4]
+        "landmark": torch.stack(landmark_lst),    #[N, 33, 4]
         "hands_visibility": hands_visibility_lst,
     }
 
@@ -302,250 +310,51 @@ def process_image(img_file, hps_type, single, input_res=512):
     return return_dict
 
 
-def get_transform(center, scale, res):
-    """Generate transformation matrix."""
-    h = 100 * scale
-    t = np.zeros((3, 3))
-    t[0, 0] = float(res[1]) / h
-    t[1, 1] = float(res[0]) / h
-    t[0, 2] = res[1] * (-float(center[0]) / h + 0.5)
-    t[1, 2] = res[0] * (-float(center[1]) / h + 0.5)
-    t[2, 2] = 1
-
-    return t
-
-
-def transform(pt, center, scale, res, invert=0):
-    """Transform pixel location to different reference."""
-    t = get_transform(center, scale, res)
-    if invert:
-        t = np.linalg.inv(t)
-    new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.0]).T
-    new_pt = np.dot(t, new_pt)
-    return np.around(new_pt[:2]).astype(np.int16)
-
-
-def crop(img, center, scale, res):
-    """Crop image according to the supplied bounding box."""
-
-    img_height, img_width = img.shape[:2]
-
-    # Upper left point
-    ul = np.array(transform([0, 0], center, scale, res, invert=1))
-
-    # Bottom right point
-    br = np.array(transform(res, center, scale, res, invert=1))
-
-    new_shape = [br[1] - ul[1], br[0] - ul[0]]
-    if len(img.shape) > 2:
-        new_shape += [img.shape[2]]
-    new_img = np.zeros(new_shape)
-
-    # Range to fill new array
-    new_x = max(0, -ul[0]), min(br[0], img_width) - ul[0]
-    new_y = max(0, -ul[1]), min(br[1], img_height) - ul[1]
-
-    # Range to sample from original image
-    old_x = max(0, ul[0]), min(img_width, br[0])
-    old_y = max(0, ul[1]), min(img_height, br[1])
-
-    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1], old_x[0]:old_x[1]]
-    new_img = F.interpolate(
-        torch.tensor(new_img).permute(2, 0, 1).unsqueeze(0), res, mode='bilinear').permute(0, 2, 3,
-                                                                                           1)[0].numpy().astype(np.uint8)
-
-    return new_img, (old_x, new_x, old_y, new_y, new_shape)
-
-
-def crop_segmentation(org_coord, res, cropping_parameters):
-    old_x, new_x, old_y, new_y, new_shape = cropping_parameters
+def blend_rgb_norm(norms, data):
 
-    new_coord = np.zeros((org_coord.shape))
-    new_coord[:, 0] = new_x[0] + (org_coord[:, 0] - old_x[0])
-    new_coord[:, 1] = new_y[0] + (org_coord[:, 1] - old_y[0])
-
-    new_coord[:, 0] = res[0] * (new_coord[:, 0] / new_shape[1])
-    new_coord[:, 1] = res[1] * (new_coord[:, 1] / new_shape[0])
-
-    return new_coord
-
-
-def corner_align(ul, br):
-
-    if ul[1] - ul[0] != br[1] - br[0]:
-        ul[1] = ul[0] + br[1] - br[0]
-
-    return ul, br
-
-
-def uncrop(img, center, scale, orig_shape):
-    """'Undo' the image cropping/resizing.
-    This function is used when evaluating mask/part segmentation.
-    """
-
-    res = img.shape[:2]
-
-    # Upper left point
-    ul = np.array(transform([0, 0], center, scale, res, invert=1))
-    # Bottom right point
-    br = np.array(transform(res, center, scale, res, invert=1))
-
-    # quick fix
-    ul, br = corner_align(ul, br)
-
-    # size of cropped image
-    crop_shape = [br[1] - ul[1], br[0] - ul[0]]
-    new_img = np.zeros(orig_shape, dtype=np.uint8)
-
-    # Range to fill new array
-    new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0]
-    new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1]
+    # norms [N, 3, res, res]
+    masks = (norms.sum(dim=1) != norms[0, :, 0, 0].sum()).float().unsqueeze(1)
+    norm_mask = F.interpolate(
+        torch.cat([norms, masks], dim=1).detach(),
+        size=data["uncrop_param"]["box_shape"],
+        mode="bilinear",
+        align_corners=False
+    )
+    final = data["img_raw"].type_as(norm_mask)
 
-    # Range to sample from original image
-    old_x = max(0, ul[0]), min(orig_shape[1], br[0])
-    old_y = max(0, ul[1]), min(orig_shape[0], br[1])
+    for idx in range(len(norms)):
 
-    img = np.array(Image.fromarray(img.astype(np.uint8)).resize(crop_shape))
+        norm_pred = (norm_mask[idx:idx + 1, :3, :, :] + 1.0) * 255.0 / 2.0
+        mask_pred = norm_mask[idx:idx + 1, 3:4, :, :].repeat(1, 3, 1, 1)
 
-    new_img[old_y[0]:old_y[1], old_x[0]:old_x[1]] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]]
+        norm_ori = unwrap(norm_pred, data["uncrop_param"], idx)
+        mask_ori = unwrap(mask_pred, data["uncrop_param"], idx)
 
-    return new_img
+        final = final * (1.0 - mask_ori) + norm_ori * mask_ori
 
+    return final.detach().cpu()
 
-def rot_aa(aa, rot):
-    """Rotate axis angle parameters."""
-    # pose parameters
-    R = np.array([
-        [np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
-        [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
-        [0, 0, 1],
-    ])
-    # find the rotation of the body in camera frame
-    per_rdg, _ = cv2.Rodrigues(aa)
-    # apply the global rotation to the global orientation
-    resrot, _ = cv2.Rodrigues(np.dot(R, per_rdg))
-    aa = (resrot.T)[0]
-    return aa
 
+def unwrap(image, uncrop_param, idx):
 
-def flip_img(img):
-    """Flip rgb images or masks.
-    channels come last, e.g. (256,256,3).
-    """
-    img = np.fliplr(img)
-    return img
+    device = image.device
 
+    img_square = warp_affine(
+        image,
+        torch.inverse(uncrop_param["M_crop"])[idx:idx + 1, :2].to(device),
+        uncrop_param["square_shape"],
+        mode='bilinear',
+        padding_mode='zeros',
+        align_corners=True
+    )
 
-def flip_kp(kp, is_smpl=False):
-    """Flip keypoints."""
-    if len(kp) == 24:
-        if is_smpl:
-            flipped_parts = constants.SMPL_JOINTS_FLIP_PERM
-        else:
-            flipped_parts = constants.J24_FLIP_PERM
-    elif len(kp) == 49:
-        if is_smpl:
-            flipped_parts = constants.SMPL_J49_FLIP_PERM
-        else:
-            flipped_parts = constants.J49_FLIP_PERM
-    kp = kp[flipped_parts]
-    kp[:, 0] = -kp[:, 0]
-    return kp
-
-
-def flip_pose(pose):
-    """Flip pose.
-    The flipping is based on SMPL parameters.
-    """
-    flipped_parts = constants.SMPL_POSE_FLIP_PERM
-    pose = pose[flipped_parts]
-    # we also negate the second and the third dimension of the axis-angle
-    pose[1::3] = -pose[1::3]
-    pose[2::3] = -pose[2::3]
-    return pose
-
-
-def normalize_2d_kp(kp_2d, crop_size=224, inv=False):
-    # Normalize keypoints between -1, 1
-    if not inv:
-        ratio = 1.0 / crop_size
-        kp_2d = 2.0 * kp_2d * ratio - 1.0
-    else:
-        ratio = 1.0 / crop_size
-        kp_2d = (kp_2d + 1.0) / (2 * ratio)
-
-    return kp_2d
-
-
-def visualize_landmarks(image, joints, color):
-
-    img_w, img_h = image.shape[:2]
-
-    for joint in joints:
-        image = cv2.circle(image, (int(joint[0] * img_w), int(joint[1] * img_h)), 5, color)
-
-    return image
-
-
-def generate_heatmap(joints, heatmap_size, sigma=1, joints_vis=None):
-    """
-    param joints:  [num_joints, 3]
-    param joints_vis: [num_joints, 3]
-    return: target, target_weight(1: visible, 0: invisible)
-    """
-    num_joints = joints.shape[0]
-    device = joints.device
-    cur_device = torch.device(device.type, device.index)
-    if not hasattr(heatmap_size, "__len__"):
-        # width  height
-        heatmap_size = [heatmap_size, heatmap_size]
-    assert len(heatmap_size) == 2
-    target_weight = np.ones((num_joints, 1), dtype=np.float32)
-    if joints_vis is not None:
-        target_weight[:, 0] = joints_vis[:, 0]
-    target = torch.zeros(
-        (num_joints, heatmap_size[1], heatmap_size[0]),
-        dtype=torch.float32,
-        device=cur_device,
+    img_ori = warp_affine(
+        img_square,
+        torch.inverse(uncrop_param["M_square"])[:, :2].to(device),
+        uncrop_param["ori_shape"],
+        mode='bilinear',
+        padding_mode='zeros',
+        align_corners=True
     )
 
-    tmp_size = sigma * 3
-
-    for joint_id in range(num_joints):
-        mu_x = int(joints[joint_id][0] * heatmap_size[0] + 0.5)
-        mu_y = int(joints[joint_id][1] * heatmap_size[1] + 0.5)
-        # Check that any part of the gaussian is in-bounds
-        ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
-        br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
-        if (ul[0] >= heatmap_size[0] or ul[1] >= heatmap_size[1] or br[0] < 0 or br[1] < 0):
-            # If not, just return the image as is
-            target_weight[joint_id] = 0
-            continue
-
-        # # Generate gaussian
-        size = 2 * tmp_size + 1
-        # x = np.arange(0, size, 1, np.float32)
-        # y = x[:, np.newaxis]
-        # x0 = y0 = size // 2
-        # # The gaussian is not normalized, we want the center value to equal 1
-        # g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
-        # g = torch.from_numpy(g.astype(np.float32))
-
-        x = torch.arange(0, size, dtype=torch.float32, device=cur_device)
-        y = x.unsqueeze(-1)
-        x0 = y0 = size // 2
-        # The gaussian is not normalized, we want the center value to equal 1
-        g = torch.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))
-
-        # Usable gaussian range
-        g_x = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0]
-        g_y = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1]
-        # Image range
-        img_x = max(0, ul[0]), min(br[0], heatmap_size[0])
-        img_y = max(0, ul[1]), min(br[1], heatmap_size[1])
-
-        v = target_weight[joint_id]
-        if v > 0.5:
-            target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
-
-    return target, target_weight
+    return img_ori

lib/common/libmesh/inside_mesh.py

@@ -5,7 +5,7 @@ from .triangle_hash import TriangleHash as _TriangleHash
 def check_mesh_contains(mesh, points, hash_resolution=512):
     intersector = MeshIntersector(mesh, hash_resolution)
     contains, hole_points = intersector.query(points)
-    return contains,  hole_points
+    return contains, hole_points
 
 
 class MeshIntersector:
@@ -25,8 +25,7 @@ class MeshIntersector:
         # assert(np.allclose(triangles.reshape(-1, 3).max(0), resolution - 0.5))
 
         triangles2d = triangles[:, :, :2]
-        self._tri_intersector2d = TriangleIntersector2d(
-            triangles2d, resolution)
+        self._tri_intersector2d = TriangleIntersector2d(triangles2d, resolution)
 
     def query(self, points):
         # Rescale points
@@ -38,8 +37,7 @@ class MeshIntersector:
 
         # cull points outside of the axis aligned bounding box
         # this avoids running ray tests unless points are close
-        inside_aabb = np.all(
-            (0 <= points) & (points <= self.resolution), axis=1)
+        inside_aabb = np.all((0 <= points) & (points <= self.resolution), axis=1)
         if not inside_aabb.any():
             return contains, hole_points
 
@@ -48,14 +46,14 @@ class MeshIntersector:
         points = points[mask]
 
         # Compute intersection depth and check order
-        points_indices, tri_indices = self._tri_intersector2d.query(
-            points[:, :2])
+        points_indices, tri_indices = self._tri_intersector2d.query(points[:, :2])
 
         triangles_intersect = self._triangles[tri_indices]
         points_intersect = points[points_indices]
 
         depth_intersect, abs_n_2 = self.compute_intersection_depth(
-            points_intersect, triangles_intersect)
+            points_intersect, triangles_intersect
+        )
 
         # Count number of intersections in both directions
         smaller_depth = depth_intersect >= points_intersect[:, 2] * abs_n_2
@@ -73,7 +71,7 @@ class MeshIntersector:
         #     print('Warning: contains1 != contains2 for some points.')
         contains[mask] = (contains1 & contains2)
         hole_points[mask] = np.logical_xor(contains1, contains2)
-        return contains,  hole_points
+        return contains, hole_points
 
     def compute_intersection_depth(self, points, triangles):
         t1 = triangles[:, 0, :]
@@ -150,7 +148,7 @@ class TriangleIntersector2d:
 
         sum_uv = u + v
         contains[mask] = (
-            (0 < u) & (u < abs_detA) & (0 < v) & (v < abs_detA)
-            & (0 < sum_uv) & (sum_uv < abs_detA)
+            (0 < u) & (u < abs_detA) & (0 < v) & (v < abs_detA) & (0 < sum_uv) &
+            (sum_uv < abs_detA)
         )
         return contains

lib/common/libmesh/setup.py

@@ -2,7 +2,4 @@ from setuptools import setup
 from Cython.Build import cythonize
 import numpy
 
-
-setup(name = 'libmesh',
-      ext_modules = cythonize("*.pyx"),
-      include_dirs=[numpy.get_include()])
+setup(name='libmesh', ext_modules=cythonize("*.pyx"), include_dirs=[numpy.get_include()])

lib/common/libvoxelize/setup.py

@@ -1,5 +1,4 @@
 from setuptools import setup
 from Cython.Build import cythonize
 
-setup(name = 'libvoxelize',
-      ext_modules = cythonize("*.pyx"))
+setup(name='libvoxelize', ext_modules=cythonize("*.pyx"))

lib/common/local_affine.py

@@ -16,7 +16,6 @@ from lib.common.train_util import init_loss
 
 # reference: https://github.com/wuhaozhe/pytorch-nicp
 class LocalAffine(nn.Module):
-
     def __init__(self, num_points, batch_size=1, edges=None):
         '''
             specify the number of points, the number of points should be constant across the batch
@@ -26,8 +25,14 @@ class LocalAffine(nn.Module):
             add additional pooling on top of w matrix
         '''
         super(LocalAffine, self).__init__()
-        self.A = nn.Parameter(torch.eye(3).unsqueeze(0).unsqueeze(0).repeat(batch_size, num_points, 1, 1))
-        self.b = nn.Parameter(torch.zeros(3).unsqueeze(0).unsqueeze(0).unsqueeze(3).repeat(batch_size, num_points, 1, 1))
+        self.A = nn.Parameter(
+            torch.eye(3).unsqueeze(0).unsqueeze(0).repeat(batch_size, num_points, 1, 1)
+        )
+        self.b = nn.Parameter(
+            torch.zeros(3).unsqueeze(0).unsqueeze(0).unsqueeze(3).repeat(
+                batch_size, num_points, 1, 1
+            )
+        )
         self.edges = edges
         self.num_points = num_points
 
@@ -38,24 +43,23 @@ class LocalAffine(nn.Module):
         '''
         if self.edges is None:
             raise Exception("edges cannot be none when calculate stiff")
-        idx1 = self.edges[:, 0]
-        idx2 = self.edges[:, 1]
         affine_weight = torch.cat((self.A, self.b), dim=3)
-        w1 = torch.index_select(affine_weight, dim=1, index=idx1)
-        w2 = torch.index_select(affine_weight, dim=1, index=idx2)
+        w1 = torch.index_select(affine_weight, dim=1, index=self.edges[:, 0])
+        w2 = torch.index_select(affine_weight, dim=1, index=self.edges[:, 1])
         w_diff = (w1 - w2)**2
         w_rigid = (torch.linalg.det(self.A) - 1.0)**2
         return w_diff, w_rigid
 
     def forward(self, x):
         '''
-            x should have shape of B * N * 3
+            x should have shape of B * N * 3 * 1
         '''
         x = x.unsqueeze(3)
         out_x = torch.matmul(self.A, x)
         out_x = out_x + self.b
-        stiffness, rigid = self.stiffness()
         out_x.squeeze_(3)
+        stiffness, rigid = self.stiffness()
+
         return out_x, stiffness, rigid
 
 
@@ -75,10 +79,16 @@ def register(target_mesh, src_mesh, device):
     tgt_mesh = trimesh2meshes(target_mesh).to(device)
     src_verts = src_mesh.verts_padded().clone()
 
-    local_affine_model = LocalAffine(src_mesh.verts_padded().shape[1],
-                                     src_mesh.verts_padded().shape[0], src_mesh.edges_packed()).to(device)
+    local_affine_model = LocalAffine(
+        src_mesh.verts_padded().shape[1],
+        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",
@@ -90,28 +100,27 @@ def register(target_mesh, src_mesh, device):
 
     losses = init_loss()
 
-    loop_cloth = tqdm(range(200))
+    loop_cloth = tqdm(range(100))
 
     for i in loop_cloth:
 
         optimizer_cloth.zero_grad()
 
-        deformed_verts, stiffness, rigid = local_affine_model(src_verts)
+        deformed_verts, stiffness, rigid = local_affine_model(x=src_verts)
         src_mesh = src_mesh.update_padded(deformed_verts)
 
         # losses for laplacian, edge, normal consistency
         update_mesh_shape_prior_losses(src_mesh, losses)
 
         losses["cloth"]["value"] = chamfer_distance(
-            x=src_mesh.verts_padded(),
-            y=tgt_mesh.verts_padded())[0]
-        
-        losses["stiffness"]["value"] = torch.mean(stiffness)
+            x=src_mesh.verts_padded(), y=tgt_mesh.verts_padded()
+        )[0]
+        losses["stiff"]["value"] = torch.mean(stiffness)
         losses["rigid"]["value"] = torch.mean(rigid)
 
         # Weighted sum of the losses
         cloth_loss = torch.tensor(0.0, requires_grad=True).to(device)
-        pbar_desc = "Register SMPL-X towards ECON --- "
+        pbar_desc = "Register SMPL-X -> d-BiNI -- "
 
         for k in losses.keys():
             if losses[k]["weight"] > 0.0 and losses[k]["value"] != 0.0:
@@ -119,7 +128,7 @@ def register(target_mesh, src_mesh, device):
                     losses[k]["value"] * losses[k]["weight"]
                 pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.3f} | "
 
-        pbar_desc += f"Total: {cloth_loss:.5f}"
+        pbar_desc += f"TOTAL: {cloth_loss:.3f}"
         loop_cloth.set_description(pbar_desc)
 
         # update params
@@ -131,6 +140,7 @@ def register(target_mesh, src_mesh, device):
         src_mesh.verts_packed().detach().squeeze(0).cpu(),
         src_mesh.faces_packed().detach().squeeze(0).cpu(),
         process=False,
-        maintains_order=True)
+        maintains_order=True
+    )
 
     return final

lib/common/render.py

@@ -31,7 +31,8 @@ from pytorch3d.renderer import (
 )
 from pytorch3d.renderer.mesh import TexturesVertex
 from pytorch3d.structures import Meshes
-from lib.dataset.mesh_util import get_visibility, blend_rgb_norm
+from lib.dataset.mesh_util import get_visibility
+from lib.common.imutils import blend_rgb_norm
 
 import lib.common.render_utils as util
 import torch
@@ -74,20 +75,23 @@ def query_color(verts, faces, image, device):
 
     (xy, z) = verts.split([2, 1], dim=1)
     visibility = get_visibility(xy, z, faces[:, [0, 2, 1]]).flatten()
-    uv = xy.unsqueeze(0).unsqueeze(2)  # [B, N, 2]
+    uv = 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)
+        (
+            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]
+        0.5 * 255.0
+    )[visibility == 0.0]
 
     return colors.detach().cpu()
 
 
 class cleanShader(torch.nn.Module):
-
     def __init__(self, blend_params=None):
         super().__init__()
         self.blend_params = blend_params if blend_params is not None else BlendParams()
@@ -103,7 +107,6 @@ class cleanShader(torch.nn.Module):
 
 
 class Render:
-
     def __init__(self, size=512, device=torch.device("cuda:0")):
         self.device = device
         self.size = size
@@ -119,21 +122,30 @@ class Render:
 
         self.cam_pos = {
             "frontback":
-                torch.tensor([
-                    (0, self.mesh_y_center, self.dis),
-                    (0, self.mesh_y_center, -self.dis),
-                ]),
+                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),
-                ]),
+                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)])
+                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"
@@ -153,8 +165,8 @@ class Render:
 
         R, T = look_at_view_transform(
             eye=self.cam_pos[type][idx],
-            at=((0, self.mesh_y_center, 0),),
-            up=((0, 1, 0),),
+            at=((0, self.mesh_y_center, 0), ),
+            up=((0, 1, 0), ),
         )
 
         cameras = FoVOrthographicCameras(
@@ -167,7 +179,7 @@ class Render:
             min_y=-100.0,
             max_x=100.0,
             min_x=-100.0,
-            scale_xyz=(self.scale * np.ones(3),) * len(R),
+            scale_xyz=(self.scale * np.ones(3), ) * len(R),
         )
 
         return cameras
@@ -202,15 +214,17 @@ class Render:
                 cull_backfaces=True,
             )
 
-            self.silhouetteRas = MeshRasterizer(cameras=camera,
-                                                raster_settings=self.raster_settings_silhouette)
-            self.renderer = MeshRenderer(rasterizer=self.silhouetteRas,
-                                         shader=SoftSilhouetteShader())
+            self.silhouetteRas = MeshRasterizer(
+                cameras=camera, raster_settings=self.raster_settings_silhouette
+            )
+            self.renderer = MeshRenderer(
+                rasterizer=self.silhouetteRas, shader=SoftSilhouetteShader()
+            )
 
         elif type == "pointcloud":
-            self.raster_settings_pcd = PointsRasterizationSettings(image_size=self.size,
-                                                                   radius=0.006,
-                                                                   points_per_pixel=10)
+            self.raster_settings_pcd = PointsRasterizationSettings(
+                image_size=self.size, radius=0.006, points_per_pixel=10
+            )
 
             self.pcdRas = PointsRasterizer(cameras=camera, raster_settings=self.raster_settings_pcd)
             self.renderer = PointsRenderer(
@@ -230,8 +244,12 @@ class Render:
             V_lst = []
             F_lst = []
             for V, F in zip(verts, faces):
-                V_lst.append(torch.tensor(V).float().to(self.device))
-                F_lst.append(torch.tensor(F).long().to(self.device))
+                if not torch.is_tensor(V):
+                    V_lst.append(torch.tensor(V).float().to(self.device))
+                    F_lst.append(torch.tensor(F).long().to(self.device))
+                else:
+                    V_lst.append(V.float().to(self.device))
+                    F_lst.append(F.long().to(self.device))
             self.meshes = Meshes(V_lst, F_lst).to(self.device)
         else:
             # array or tensor
@@ -248,7 +266,8 @@ class Render:
         # texture only support single mesh
         if len(self.meshes) == 1:
             self.meshes.textures = TexturesVertex(
-                verts_features=(self.meshes.verts_normals_padded() + 1.0) * 0.5)
+                verts_features=(self.meshes.verts_normals_padded() + 1.0) * 0.5
+            )
 
     def get_image(self, cam_type="frontback", type="rgb", bg="gray"):
 
@@ -260,7 +279,8 @@ class Render:
 
             current_mesh = self.meshes[mesh_id]
             current_mesh.textures = TexturesVertex(
-                verts_features=(current_mesh.verts_normals_padded() + 1.0) * 0.5)
+                verts_features=(current_mesh.verts_normals_padded() + 1.0) * 0.5
+            )
 
             if type == "depth":
                 fragments = self.meshRas(current_mesh.extend(len(self.cam_pos[cam_type])))
@@ -276,7 +296,7 @@ class Render:
                 print(f"unknown {type}")
 
             if cam_type == 'frontback':
-                images[1] = torch.flip(images[1], dims=(-1,))
+                images[1] = torch.flip(images[1], dims=(-1, ))
 
             # images [N_render, 3, res, res]
             img_lst.append(images.unsqueeze(1))
@@ -287,9 +307,8 @@ class Render:
         return list(meshes)
 
     def get_rendered_video_multi(self, data, save_path):
-        
-        width = data["img_raw"].shape[1]
-        height = data["img_raw"].shape[0]
+
+        height, width = data["img_raw"].shape[2:]
 
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         video = cv2.VideoWriter(
@@ -302,14 +321,15 @@ class Render:
         pbar = tqdm(range(len(self.meshes)))
         pbar.set_description(colored(f"Normal Rendering {os.path.basename(save_path)}...", "blue"))
 
-        mesh_renders = []  #[(N_cam, 3, res, res)*N_mesh]
+        mesh_renders = []    #[(N_cam, 3, res, res)*N_mesh]
 
         # render all the normals
         for mesh_id in pbar:
 
             current_mesh = self.meshes[mesh_id]
             current_mesh.textures = TexturesVertex(
-                verts_features=(current_mesh.verts_normals_padded() + 1.0) * 0.5)
+                verts_features=(current_mesh.verts_normals_padded() + 1.0) * 0.5
+            )
 
             norm_lst = []
 
@@ -320,21 +340,33 @@ class Render:
                 self.init_renderer(batch_cams, "mesh", "gray")
 
                 norm_lst.append(
-                    self.renderer(current_mesh.extend(len(batch_cams_idx)))[..., :3].permute(
-                        0, 3, 1, 2))
+                    self.renderer(current_mesh.extend(len(batch_cams_idx))
+                                 )[..., :3].permute(0, 3, 1, 2)
+                )
             mesh_renders.append(torch.cat(norm_lst).detach().cpu())
 
         # generate video frame by frame
         pbar = tqdm(range(len(self.cam_pos["around"])))
         pbar.set_description(colored(f"Video Exporting {os.path.basename(save_path)}...", "blue"))
+
         for cam_id in pbar:
-            img_raw = data["img_raw"].astype(np.uint8)
+            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)
+            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
+            )
 
-            top_img = cv2.resize(np.concatenate([img_raw, img_smpl], axis=1), (width, height // 2))
-            final_img = np.concatenate([top_img, img_cloth], axis=0)
+            top_img = cv2.resize(
+                torch.cat([img_raw, img_smpl],
+                          dim=-1).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8),
+                (width, height // 2)
+            )
+            final_img = np.concatenate(
+                [top_img, img_cloth.squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)], axis=0
+            )
             video.write(final_img[:, :, ::-1])
 
         video.release()

lib/common/render_utils.py

@@ -25,9 +25,7 @@ from pytorch3d.renderer.mesh import rasterize_meshes
 Tensor = NewType("Tensor", torch.Tensor)
 
 
-def solid_angles(points: Tensor,
-                 triangles: Tensor,
-                 thresh: float = 1e-8) -> Tensor:
+def solid_angles(points: Tensor, triangles: Tensor, thresh: float = 1e-8) -> Tensor:
     """Compute solid angle between the input points and triangles
     Follows the method described in:
     The Solid Angle of a Plane Triangle
@@ -55,9 +53,7 @@ def solid_angles(points: Tensor,
     norms = torch.norm(centered_tris, dim=-1)
 
     # Should be BxQxFx3
-    cross_prod = torch.cross(centered_tris[:, :, :, 1],
-                             centered_tris[:, :, :, 2],
-                             dim=-1)
+    cross_prod = torch.cross(centered_tris[:, :, :, 1], centered_tris[:, :, :, 2], dim=-1)
     # Should be BxQxF
     numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1)
     del cross_prod
@@ -67,8 +63,10 @@ def solid_angles(points: Tensor,
     dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1)
     del centered_tris
 
-    denominator = (norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] +
-                   dot02 * norms[:, :, :, 1] + dot12 * norms[:, :, :, 0])
+    denominator = (
+        norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] + dot02 * norms[:, :, :, 1] +
+        dot12 * norms[:, :, :, 0]
+    )
     del dot01, dot12, dot02, norms
 
     # Should be BxQ
@@ -80,9 +78,7 @@ def solid_angles(points: Tensor,
     return 2 * solid_angle
 
 
-def winding_numbers(points: Tensor,
-                    triangles: Tensor,
-                    thresh: float = 1e-8) -> Tensor:
+def winding_numbers(points: Tensor, triangles: Tensor, thresh: float = 1e-8) -> Tensor:
     """Uses winding_numbers to compute inside/outside
     Robust inside-outside segmentation using generalized winding numbers
     Alec Jacobson,
@@ -109,8 +105,7 @@ def winding_numbers(points: Tensor,
     """
     # The generalized winding number is the sum of solid angles of the point
     # with respect to all triangles.
-    return (1 / (4 * math.pi) *
-            solid_angles(points, triangles, thresh=thresh).sum(dim=-1))
+    return (1 / (4 * math.pi) * solid_angles(points, triangles, thresh=thresh).sum(dim=-1))
 
 
 def batch_contains(verts, faces, points):
@@ -124,8 +119,7 @@ def batch_contains(verts, faces, points):
     contains = torch.zeros(B, N)
 
     for i in range(B):
-        contains[i] = torch.as_tensor(
-            trimesh.Trimesh(verts[i], faces[i]).contains(points[i]))
+        contains[i] = torch.as_tensor(trimesh.Trimesh(verts[i], faces[i]).contains(points[i]))
 
     return 2.0 * (contains - 0.5)
 
@@ -155,8 +149,7 @@ def face_vertices(vertices, faces):
     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]
+    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()]
@@ -168,7 +161,6 @@ 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):
         """
         use fixed raster_settings for rendering faces
@@ -189,8 +181,7 @@ class Pytorch3dRasterizer(nn.Module):
     def forward(self, vertices, faces, attributes=None):
         fixed_vertices = vertices.clone()
         fixed_vertices[..., :2] = -fixed_vertices[..., :2]
-        meshes_screen = Meshes(verts=fixed_vertices.float(),
-                               faces=faces.long())
+        meshes_screen = Meshes(verts=fixed_vertices.float(), faces=faces.long())
         raster_settings = self.raster_settings
         pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
             meshes_screen,
@@ -204,8 +195,9 @@ class Pytorch3dRasterizer(nn.Module):
         vismask = (pix_to_face > -1).float()
         D = attributes.shape[-1]
         attributes = attributes.clone()
-        attributes = attributes.view(attributes.shape[0] * attributes.shape[1],
-                                     3, attributes.shape[-1])
+        attributes = attributes.view(
+            attributes.shape[0] * attributes.shape[1], 3, attributes.shape[-1]
+        )
         N, H, W, K, _ = bary_coords.shape
         mask = pix_to_face == -1
         pix_to_face = pix_to_face.clone()
@@ -213,8 +205,7 @@ class Pytorch3dRasterizer(nn.Module):
         idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
         pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D)
         pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2)
-        pixel_vals[mask] = 0  # Replace masked values in output.
+        pixel_vals[mask] = 0    # Replace masked values in output.
         pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2)
-        pixel_vals = torch.cat(
-            [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1)
+        pixel_vals = torch.cat([pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1)
         return pixel_vals

lib/common/seg3d_lossless.py

@@ -31,7 +31,6 @@ logging.getLogger("lightning").setLevel(logging.ERROR)
 
 
 class Seg3dLossless(nn.Module):
-
     def __init__(
         self,
         query_func,
@@ -53,19 +52,14 @@ class Seg3dLossless(nn.Module):
         """
         super().__init__()
         self.query_func = query_func
-        self.register_buffer(
-            "b_min",
-            torch.tensor(b_min).float().unsqueeze(1))  # [bz, 1, 3]
-        self.register_buffer(
-            "b_max",
-            torch.tensor(b_max).float().unsqueeze(1))  # [bz, 1, 3]
+        self.register_buffer("b_min", torch.tensor(b_min).float().unsqueeze(1))    # [bz, 1, 3]
+        self.register_buffer("b_max", torch.tensor(b_max).float().unsqueeze(1))    # [bz, 1, 3]
 
         # ti.init(arch=ti.cuda)
         # self.mciso_taichi = MCISO(dim=3, N=resolutions[-1]-1)
 
         if type(resolutions[0]) is int:
-            resolutions = torch.tensor([(res, res, res)
-                                        for res in resolutions])
+            resolutions = torch.tensor([(res, res, res) for res in resolutions])
         else:
             resolutions = torch.tensor(resolutions)
         self.register_buffer("resolutions", resolutions)
@@ -87,45 +81,36 @@ class Seg3dLossless(nn.Module):
             ), f"resolution {resolution} need to be odd becuase of align_corner."
 
         # init first resolution
-        init_coords = create_grid3D(0,
-                                    resolutions[-1] - 1,
-                                    steps=resolutions[0])  # [N, 3]
-        init_coords = init_coords.unsqueeze(0).repeat(self.batchsize, 1,
-                                                      1)  # [bz, N, 3]
+        init_coords = create_grid3D(0, resolutions[-1] - 1, steps=resolutions[0])    # [N, 3]
+        init_coords = init_coords.unsqueeze(0).repeat(self.batchsize, 1, 1)    # [bz, N, 3]
         self.register_buffer("init_coords", init_coords)
 
         # some useful tensors
         calculated = torch.zeros(
-            (self.resolutions[-1][2], self.resolutions[-1][1],
-             self.resolutions[-1][0]),
+            (self.resolutions[-1][2], self.resolutions[-1][1], self.resolutions[-1][0]),
             dtype=torch.bool,
         )
         self.register_buffer("calculated", calculated)
 
-        gird8_offsets = (torch.stack(
-            torch.meshgrid(
-                [
-                    torch.tensor([-1, 0, 1]),
-                    torch.tensor([-1, 0, 1]),
-                    torch.tensor([-1, 0, 1]),
-                ],
-                indexing="ij",
-            )).int().view(3, -1).t())  # [27, 3]
+        gird8_offsets = (
+            torch.stack(
+                torch.meshgrid(
+                    [
+                        torch.tensor([-1, 0, 1]),
+                        torch.tensor([-1, 0, 1]),
+                        torch.tensor([-1, 0, 1]),
+                    ],
+                    indexing="ij",
+                )
+            ).int().view(3, -1).t()
+        )    # [27, 3]
         self.register_buffer("gird8_offsets", gird8_offsets)
 
         # smooth convs
-        self.smooth_conv3x3 = SmoothConv3D(in_channels=1,
-                                           out_channels=1,
-                                           kernel_size=3)
-        self.smooth_conv5x5 = SmoothConv3D(in_channels=1,
-                                           out_channels=1,
-                                           kernel_size=5)
-        self.smooth_conv7x7 = SmoothConv3D(in_channels=1,
-                                           out_channels=1,
-                                           kernel_size=7)
-        self.smooth_conv9x9 = SmoothConv3D(in_channels=1,
-                                           out_channels=1,
-                                           kernel_size=9)
+        self.smooth_conv3x3 = SmoothConv3D(in_channels=1, out_channels=1, kernel_size=3)
+        self.smooth_conv5x5 = SmoothConv3D(in_channels=1, out_channels=1, kernel_size=5)
+        self.smooth_conv7x7 = SmoothConv3D(in_channels=1, out_channels=1, kernel_size=7)
+        self.smooth_conv9x9 = SmoothConv3D(in_channels=1, out_channels=1, kernel_size=9)
 
     @torch.no_grad()
     def batch_eval(self, coords, **kwargs):
@@ -144,7 +129,7 @@ class Seg3dLossless(nn.Module):
         # query function
         occupancys = self.query_func(**kwargs, points=coords2D)
         if type(occupancys) is list:
-            occupancys = torch.stack(occupancys)  # [bz, C, N]
+            occupancys = torch.stack(occupancys)    # [bz, C, N]
         assert (
             len(occupancys.size()) == 3
         ), "query_func should return a occupancy with shape of [bz, C, N]"
@@ -175,10 +160,9 @@ class Seg3dLossless(nn.Module):
 
             # first step
             if torch.equal(resolution, self.resolutions[0]):
-                coords = self.init_coords.clone()  # torch.long
+                coords = self.init_coords.clone()    # torch.long
                 occupancys = self.batch_eval(coords, **kwargs)
-                occupancys = occupancys.view(self.batchsize, self.channels, D,
-                                             H, W)
+                occupancys = occupancys.view(self.batchsize, self.channels, D, H, W)
                 if (occupancys > 0.5).sum() == 0:
                     # return F.interpolate(
                     #     occupancys, size=(final_D, final_H, final_W),
@@ -239,23 +223,22 @@ class Seg3dLossless(nn.Module):
 
                 with torch.no_grad():
                     if torch.equal(resolution, self.resolutions[1]):
-                        is_boundary = (self.smooth_conv9x9(is_boundary.float())
-                                       > 0)[0, 0]
+                        is_boundary = (self.smooth_conv9x9(is_boundary.float()) > 0)[0, 0]
                     elif torch.equal(resolution, self.resolutions[2]):
-                        is_boundary = (self.smooth_conv7x7(is_boundary.float())
-                                       > 0)[0, 0]
+                        is_boundary = (self.smooth_conv7x7(is_boundary.float()) > 0)[0, 0]
                     else:
-                        is_boundary = (self.smooth_conv3x3(is_boundary.float())
-                                       > 0)[0, 0]
+                        is_boundary = (self.smooth_conv3x3(is_boundary.float()) > 0)[0, 0]
 
                     coords_accum = coords_accum.long()
                     is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1],
                                 coords_accum[0, :, 0], ] = False
-                    point_coords = (is_boundary.permute(
-                        2, 1, 0).nonzero(as_tuple=False).unsqueeze(0))
-                    point_indices = (point_coords[:, :, 2] * H * W +
-                                     point_coords[:, :, 1] * W +
-                                     point_coords[:, :, 0])
+                    point_coords = (
+                        is_boundary.permute(2, 1, 0).nonzero(as_tuple=False).unsqueeze(0)
+                    )
+                    point_indices = (
+                        point_coords[:, :, 2] * H * W + point_coords[:, :, 1] * W +
+                        point_coords[:, :, 0]
+                    )
 
                     R, C, D, H, W = occupancys.shape
 
@@ -269,13 +252,15 @@ class Seg3dLossless(nn.Module):
                 # put mask point predictions to the right places on the upsampled grid.
                 R, C, D, H, W = occupancys.shape
                 point_indices = point_indices.unsqueeze(1).expand(-1, C, -1)
-                occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
-                    2, point_indices, occupancys_topk).view(R, C, D, H, W))
+                occupancys = (
+                    occupancys.reshape(R, C,
+                                       D * H * W).scatter_(2, point_indices,
+                                                           occupancys_topk).view(R, C, D, H, W)
+                )
 
                 with torch.no_grad():
                     voxels = coords / stride
-                    coords_accum = torch.cat([voxels, coords_accum],
-                                             dim=1).unique(dim=1)
+                    coords_accum = torch.cat([voxels, coords_accum], dim=1).unique(dim=1)
 
         return occupancys[0, 0]
 
@@ -300,18 +285,16 @@ class Seg3dLossless(nn.Module):
 
             # first step
             if torch.equal(resolution, self.resolutions[0]):
-                coords = self.init_coords.clone()  # torch.long
+                coords = self.init_coords.clone()    # torch.long
                 occupancys = self.batch_eval(coords, **kwargs)
-                occupancys = occupancys.view(self.batchsize, self.channels, D,
-                                             H, W)
+                occupancys = occupancys.view(self.batchsize, self.channels, D, H, W)
 
                 if self.visualize:
                     self.plot(occupancys, coords, final_D, final_H, final_W)
 
                 with torch.no_grad():
                     coords_accum = coords / stride
-                    calculated[coords[0, :, 2], coords[0, :, 1],
-                               coords[0, :, 0]] = True
+                    calculated[coords[0, :, 2], coords[0, :, 1], coords[0, :, 0]] = True
 
             # next steps
             else:
@@ -338,35 +321,34 @@ class Seg3dLossless(nn.Module):
 
                 with torch.no_grad():
                     # TODO
-                    if self.use_shadow and torch.equal(resolution,
-                                                       self.resolutions[-1]):
+                    if self.use_shadow and torch.equal(resolution, self.resolutions[-1]):
                         # larger z means smaller depth here
                         depth_res = resolution[2].item()
-                        depth_index = torch.linspace(0,
-                                                     depth_res - 1,
-                                                     steps=depth_res).type_as(
-                                                         occupancys.device)
-                        depth_index_max = (torch.max(
-                            (occupancys > self.balance_value) *
-                            (depth_index + 1),
-                            dim=-1,
-                            keepdim=True,
-                        )[0] - 1)
+                        depth_index = torch.linspace(0, depth_res - 1,
+                                                     steps=depth_res).type_as(occupancys.device)
+                        depth_index_max = (
+                            torch.max(
+                                (occupancys > self.balance_value) * (depth_index + 1),
+                                dim=-1,
+                                keepdim=True,
+                            )[0] - 1
+                        )
                         shadow = depth_index < depth_index_max
                         is_boundary[shadow] = False
                         is_boundary = is_boundary[0, 0]
                     else:
-                        is_boundary = (self.smooth_conv3x3(is_boundary.float())
-                                       > 0)[0, 0]
+                        is_boundary = (self.smooth_conv3x3(is_boundary.float()) > 0)[0, 0]
                         # is_boundary = is_boundary[0, 0]
 
                     is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1],
                                 coords_accum[0, :, 0], ] = False
-                    point_coords = (is_boundary.permute(
-                        2, 1, 0).nonzero(as_tuple=False).unsqueeze(0))
-                    point_indices = (point_coords[:, :, 2] * H * W +
-                                     point_coords[:, :, 1] * W +
-                                     point_coords[:, :, 0])
+                    point_coords = (
+                        is_boundary.permute(2, 1, 0).nonzero(as_tuple=False).unsqueeze(0)
+                    )
+                    point_indices = (
+                        point_coords[:, :, 2] * H * W + point_coords[:, :, 1] * W +
+                        point_coords[:, :, 0]
+                    )
 
                     R, C, D, H, W = occupancys.shape
                     # interpolated value
@@ -388,28 +370,28 @@ class Seg3dLossless(nn.Module):
                 # put mask point predictions to the right places on the upsampled grid.
                 R, C, D, H, W = occupancys.shape
                 point_indices = point_indices.unsqueeze(1).expand(-1, C, -1)
-                occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
-                    2, point_indices, occupancys_topk).view(R, C, D, H, W))
+                occupancys = (
+                    occupancys.reshape(R, C,
+                                       D * H * W).scatter_(2, point_indices,
+                                                           occupancys_topk).view(R, C, D, H, W)
+                )
 
                 with torch.no_grad():
                     # conflicts
-                    conflicts = ((occupancys_interp - self.balance_value) *
-                                 (occupancys_topk - self.balance_value) < 0)[0,
-                                                                             0]
+                    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)
+                        self.plot(occupancys, coords, final_D, final_H, final_W)
 
                     voxels = coords / stride
-                    coords_accum = torch.cat([voxels, coords_accum],
-                                             dim=1).unique(dim=1)
-                    calculated[coords[0, :, 2], coords[0, :, 1],
-                               coords[0, :, 0]] = True
+                    coords_accum = torch.cat([voxels, coords_accum], dim=1).unique(dim=1)
+                    calculated[coords[0, :, 2], coords[0, :, 1], coords[0, :, 0]] = True
 
                 while conflicts.sum() > 0:
-                    if self.use_shadow and torch.equal(resolution,
-                                                       self.resolutions[-1]):
+                    if self.use_shadow and torch.equal(resolution, self.resolutions[-1]):
                         break
 
                     with torch.no_grad():
@@ -426,25 +408,27 @@ class Seg3dLossless(nn.Module):
                             )
 
                         conflicts_boundary = (
-                            (conflicts_coords.int() +
-                             self.gird8_offsets.unsqueeze(1) *
-                             stride.int()).reshape(-1, 3).long().unique(dim=0))
-                        conflicts_boundary[:,
-                                           0] = conflicts_boundary[:, 0].clamp(
-                                               0,
-                                               calculated.size(2) - 1)
-                        conflicts_boundary[:,
-                                           1] = conflicts_boundary[:, 1].clamp(
-                                               0,
-                                               calculated.size(1) - 1)
-                        conflicts_boundary[:,
-                                           2] = conflicts_boundary[:, 2].clamp(
-                                               0,
-                                               calculated.size(0) - 1)
-
-                        coords = conflicts_boundary[calculated[
-                            conflicts_boundary[:, 2], conflicts_boundary[:, 1],
-                            conflicts_boundary[:, 0], ] == False]
+                            (
+                                conflicts_coords.int() +
+                                self.gird8_offsets.unsqueeze(1) * stride.int()
+                            ).reshape(-1, 3).long().unique(dim=0)
+                        )
+                        conflicts_boundary[:, 0] = conflicts_boundary[:, 0].clamp(
+                            0,
+                            calculated.size(2) - 1
+                        )
+                        conflicts_boundary[:, 1] = conflicts_boundary[:, 1].clamp(
+                            0,
+                            calculated.size(1) - 1
+                        )
+                        conflicts_boundary[:, 2] = conflicts_boundary[:, 2].clamp(
+                            0,
+                            calculated.size(0) - 1
+                        )
+
+                        coords = conflicts_boundary[calculated[conflicts_boundary[:, 2],
+                                                               conflicts_boundary[:, 1],
+                                                               conflicts_boundary[:, 0], ] == False]
 
                         if self.debug:
                             self.plot(
@@ -458,9 +442,10 @@ class Seg3dLossless(nn.Module):
 
                         coords = coords.unsqueeze(0)
                         point_coords = coords / stride
-                        point_indices = (point_coords[:, :, 2] * H * W +
-                                         point_coords[:, :, 1] * W +
-                                         point_coords[:, :, 0])
+                        point_indices = (
+                            point_coords[:, :, 2] * H * W + point_coords[:, :, 1] * W +
+                            point_coords[:, :, 0]
+                        )
 
                         R, C, D, H, W = occupancys.shape
                         # interpolated value
@@ -481,44 +466,37 @@ 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)
-                    occupancys = (occupancys.reshape(R, C, D * H * W).scatter_(
-                        2, point_indices, occupancys_topk).view(R, C, D, H, W))
+                    point_indices = point_indices.unsqueeze(1).expand(-1, C, -1)
+                    occupancys = (
+                        occupancys.reshape(R, C,
+                                           D * H * W).scatter_(2, point_indices,
+                                                               occupancys_topk).view(R, C, D, H, W)
+                    )
 
                     with torch.no_grad():
                         voxels = coords / stride
-                        coords_accum = torch.cat([voxels, coords_accum],
-                                                 dim=1).unique(dim=1)
-                        calculated[coords[0, :, 2], coords[0, :, 1],
-                                   coords[0, :, 0]] = True
+                        coords_accum = torch.cat([voxels, coords_accum], dim=1).unique(dim=1)
+                        calculated[coords[0, :, 2], coords[0, :, 1], coords[0, :, 0]] = True
 
                 if self.visualize:
                     this_stage_coords = torch.cat(this_stage_coords, dim=1)
-                    self.plot(occupancys, this_stage_coords, final_D, final_H,
-                              final_W)
+                    self.plot(occupancys, this_stage_coords, final_D, final_H, final_W)
 
         return occupancys[0, 0]
 
-    def plot(self,
-             occupancys,
-             coords,
-             final_D,
-             final_H,
-             final_W,
-             title="",
-             **kwargs):
+    def plot(self, occupancys, coords, final_D, final_H, final_W, title="", **kwargs):
         final = F.interpolate(
             occupancys.float(),
             size=(final_D, final_H, final_W),
             mode="trilinear",
             align_corners=True,
-        )  # here true is correct!
+        )    # here true is correct!
         x = coords[0, :, 0].to("cpu")
         y = coords[0, :, 1].to("cpu")
         z = coords[0, :, 2].to("cpu")
@@ -548,20 +526,19 @@ class Seg3dLossless(nn.Module):
         sdf_all = sdf.permute(2, 1, 0)
 
         # shadow
-        grad_v = (sdf_all > 0.5) * torch.linspace(
-            resolution, 1, steps=resolution).to(sdf.device)
-        grad_c = torch.ones_like(sdf_all) * torch.linspace(
-            0, resolution - 1, steps=resolution).to(sdf.device)
+        grad_v = (sdf_all > 0.5) * torch.linspace(resolution, 1, steps=resolution).to(sdf.device)
+        grad_c = torch.ones_like(sdf_all) * torch.linspace(0, resolution - 1,
+                                                           steps=resolution).to(sdf.device)
         max_v, max_c = grad_v.max(dim=2)
         shadow = grad_c > max_c.view(resolution, resolution, 1)
         keep = (sdf_all > 0.5) & (~shadow)
 
-        p1 = keep.nonzero(as_tuple=False).t()  # [3, N]
-        p2 = p1.clone()  # z
+        p1 = keep.nonzero(as_tuple=False).t()    # [3, N]
+        p2 = p1.clone()    # z
         p2[2, :] = (p2[2, :] - 2).clamp(0, resolution)
-        p3 = p1.clone()  # y
+        p3 = p1.clone()    # y
         p3[1, :] = (p3[1, :] - 2).clamp(0, resolution)
-        p4 = p1.clone()  # x
+        p4 = p1.clone()    # x
         p4[0, :] = (p4[0, :] - 2).clamp(0, resolution)
 
         v1 = sdf_all[p1[0, :], p1[1, :], p1[2, :]]
@@ -569,10 +546,10 @@ class Seg3dLossless(nn.Module):
         v3 = sdf_all[p3[0, :], p3[1, :], p3[2, :]]
         v4 = sdf_all[p4[0, :], p4[1, :], p4[2, :]]
 
-        X = p1[0, :].long()  # [N,]
-        Y = p1[1, :].long()  # [N,]
-        Z = p2[2, :].float() * (0.5 - v1) / (v2 - v1) + p1[2, :].float() * (
-            v2 - 0.5) / (v2 - v1)  # [N,]
+        X = p1[0, :].long()    # [N,]
+        Y = p1[1, :].long()    # [N,]
+        Z = p2[2, :].float() * (0.5 - v1) / (v2 - v1) + p1[2, :].float() * (v2 - 0.5
+                                                                           ) / (v2 - v1)    # [N,]
         Z = Z.clamp(0, resolution)
 
         # normal
@@ -588,8 +565,7 @@ class Seg3dLossless(nn.Module):
 
     @torch.no_grad()
     def render_normal(self, resolution, X, Y, Z, norm):
-        image = torch.ones((1, 3, resolution, resolution),
-                           dtype=torch.float32).to(norm.device)
+        image = torch.ones((1, 3, resolution, resolution), dtype=torch.float32).to(norm.device)
         color = (norm + 1) / 2.0
         color = color.clamp(0, 1)
         image[0, :, Y, X] = color.t()
@@ -617,9 +593,9 @@ class Seg3dLossless(nn.Module):
     def export_mesh(self, occupancys):
 
         final = occupancys[1:, 1:, 1:].contiguous()
-        
+
         verts, faces = marching_cubes(final.unsqueeze(0), isolevel=0.5)
         verts = verts[0].cpu().float()
-        faces = faces[0].cpu().long()[:,[0,2,1]]
-        
+        faces = faces[0].cpu().long()[:, [0, 2, 1]]
+
         return verts, faces

lib/common/seg3d_utils.py

@@ -20,11 +20,7 @@ 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):
+def plot_mask2D(mask, title="", point_coords=None, figsize=10, point_marker_size=5):
     '''
     Simple plotting tool to show intermediate mask predictions and points 
     where PointRend is applied.
@@ -46,26 +42,19 @@ def plot_mask2D(mask,
     plt.xlabel(W, fontsize=30)
     plt.xticks([], [])
     plt.yticks([], [])
-    plt.imshow(mask.detach(),
-               interpolation="nearest",
-               cmap=plt.get_cmap('gray'))
+    plt.imshow(mask.detach(), interpolation="nearest", cmap=plt.get_cmap('gray'))
     if point_coords is not None:
-        plt.scatter(x=point_coords[0],
-                    y=point_coords[1],
-                    color="red",
-                    s=point_marker_size,
-                    clip_on=True)
+        plt.scatter(
+            x=point_coords[0], y=point_coords[1], color="red", s=point_marker_size, clip_on=True
+        )
     plt.xlim(-0.5, W - 0.5)
     plt.ylim(H - 0.5, -0.5)
     plt.show()
 
 
-def plot_mask3D(mask=None,
-                title="",
-                point_coords=None,
-                figsize=1500,
-                point_marker_size=8,
-                interactive=True):
+def plot_mask3D(
+    mask=None, title="", point_coords=None, figsize=1500, point_marker_size=8, interactive=True
+):
     '''
     Simple plotting tool to show intermediate mask predictions and points 
     where PointRend is applied.
@@ -90,7 +79,8 @@ def plot_mask3D(mask=None,
 
         # marching cube to find surface
         verts, faces, normals, values = measure.marching_cubes_lewiner(
-            mask, 0.5, gradient_direction='ascent')
+            mask, 0.5, gradient_direction='ascent'
+        )
 
         # create a mesh
         mesh = trimesh.Trimesh(verts, faces)
@@ -110,57 +100,49 @@ def plot_mask3D(mask=None,
         pc = vtkplotter.Points(point_coords, r=point_marker_size, c='red')
         vis_list.append(pc)
 
-    vp.show(*vis_list,
-            bg="white",
-            axes=1,
-            interactive=interactive,
-            azimuth=30,
-            elevation=30)
+    vp.show(*vis_list, bg="white", axes=1, interactive=interactive, azimuth=30, elevation=30)
 
 
 def create_grid3D(min, max, steps):
     if type(min) is int:
-        min = (min, min, min)  # (x, y, z)
+        min = (min, min, min)    # (x, y, z)
     if type(max) is int:
-        max = (max, max, max)  # (x, y)
+        max = (max, max, max)    # (x, y)
     if type(steps) is int:
-        steps = (steps, steps, steps)  # (x, y, z)
+        steps = (steps, steps, steps)    # (x, y, z)
     arrangeX = torch.linspace(min[0], max[0], steps[0]).long()
     arrangeY = torch.linspace(min[1], max[1], steps[1]).long()
     arrangeZ = torch.linspace(min[2], max[2], steps[2]).long()
-    gridD, girdH, gridW = torch.meshgrid([arrangeZ, arrangeY, arrangeX],
-                                         indexing='ij')
-    coords = torch.stack([gridW, girdH,
-                          gridD])  # [2, steps[0], steps[1], steps[2]]
-    coords = coords.view(3, -1).t()  # [N, 3]
+    gridD, girdH, gridW = torch.meshgrid([arrangeZ, arrangeY, arrangeX], indexing='ij')
+    coords = torch.stack([gridW, girdH, gridD])    # [2, steps[0], steps[1], steps[2]]
+    coords = coords.view(3, -1).t()    # [N, 3]
     return coords
 
 
 def create_grid2D(min, max, steps):
     if type(min) is int:
-        min = (min, min)  # (x, y)
+        min = (min, min)    # (x, y)
     if type(max) is int:
-        max = (max, max)  # (x, y)
+        max = (max, max)    # (x, y)
     if type(steps) is int:
-        steps = (steps, steps)  # (x, y)
+        steps = (steps, steps)    # (x, y)
     arrangeX = torch.linspace(min[0], max[0], steps[0]).long()
     arrangeY = torch.linspace(min[1], max[1], steps[1]).long()
     girdH, gridW = torch.meshgrid([arrangeY, arrangeX], indexing='ij')
-    coords = torch.stack([gridW, girdH])  # [2, steps[0], steps[1]]
-    coords = coords.view(2, -1).t()  # [N, 2]
+    coords = torch.stack([gridW, girdH])    # [2, steps[0], steps[1]]
+    coords = coords.view(2, -1).t()    # [N, 2]
     return coords
 
 
 class SmoothConv2D(nn.Module):
-
     def __init__(self, in_channels, out_channels, kernel_size=3):
         super().__init__()
         assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}"
         self.padding = (kernel_size - 1) // 2
 
         weight = torch.ones(
-            (in_channels, out_channels, kernel_size, kernel_size),
-            dtype=torch.float32) / (kernel_size**2)
+            (in_channels, out_channels, kernel_size, kernel_size), dtype=torch.float32
+        ) / (kernel_size**2)
         self.register_buffer('weight', weight)
 
     def forward(self, input):
@@ -168,53 +150,49 @@ class SmoothConv2D(nn.Module):
 
 
 class SmoothConv3D(nn.Module):
-
     def __init__(self, in_channels, out_channels, kernel_size=3):
         super().__init__()
         assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}"
         self.padding = (kernel_size - 1) // 2
 
         weight = torch.ones(
-            (in_channels, out_channels, kernel_size, kernel_size, kernel_size),
-            dtype=torch.float32) / (kernel_size**3)
+            (in_channels, out_channels, kernel_size, kernel_size, kernel_size), dtype=torch.float32
+        ) / (kernel_size**3)
         self.register_buffer('weight', weight)
 
     def forward(self, input):
         return F.conv3d(input, self.weight, padding=self.padding)
 
 
-def build_smooth_conv3D(in_channels=1,
-                        out_channels=1,
-                        kernel_size=3,
-                        padding=1):
-    smooth_conv = torch.nn.Conv3d(in_channels=in_channels,
-                                  out_channels=out_channels,
-                                  kernel_size=kernel_size,
-                                  padding=padding)
+def build_smooth_conv3D(in_channels=1, out_channels=1, kernel_size=3, padding=1):
+    smooth_conv = torch.nn.Conv3d(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        kernel_size=kernel_size,
+        padding=padding
+    )
     smooth_conv.weight.data = torch.ones(
-        (in_channels, out_channels, kernel_size, kernel_size, kernel_size),
-        dtype=torch.float32) / (kernel_size**3)
+        (in_channels, out_channels, kernel_size, kernel_size, kernel_size), dtype=torch.float32
+    ) / (kernel_size**3)
     smooth_conv.bias.data = torch.zeros(out_channels)
     return smooth_conv
 
 
-def build_smooth_conv2D(in_channels=1,
-                        out_channels=1,
-                        kernel_size=3,
-                        padding=1):
-    smooth_conv = torch.nn.Conv2d(in_channels=in_channels,
-                                  out_channels=out_channels,
-                                  kernel_size=kernel_size,
-                                  padding=padding)
+def build_smooth_conv2D(in_channels=1, out_channels=1, kernel_size=3, padding=1):
+    smooth_conv = torch.nn.Conv2d(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        kernel_size=kernel_size,
+        padding=padding
+    )
     smooth_conv.weight.data = torch.ones(
-        (in_channels, out_channels, kernel_size, kernel_size),
-        dtype=torch.float32) / (kernel_size**2)
+        (in_channels, out_channels, kernel_size, kernel_size), dtype=torch.float32
+    ) / (kernel_size**2)
     smooth_conv.bias.data = torch.zeros(out_channels)
     return smooth_conv
 
 
-def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points,
-                                         **kwargs):
+def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points, **kwargs):
     """
     Find `num_points` most uncertain points from `uncertainty_map` grid.
     Args:
@@ -233,28 +211,21 @@ def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points,
     # d_step = 1.0 / float(D)
 
     num_points = min(D * H * W, num_points)
-    point_scores, point_indices = torch.topk(uncertainty_map.view(
-        R, D * H * W),
-                                             k=num_points,
-                                             dim=1)
-    point_coords = torch.zeros(R,
-                               num_points,
-                               3,
-                               dtype=torch.float,
-                               device=uncertainty_map.device)
+    point_scores, point_indices = torch.topk(
+        uncertainty_map.view(R, D * H * W), k=num_points, dim=1
+    )
+    point_coords = torch.zeros(R, num_points, 3, dtype=torch.float, device=uncertainty_map.device)
     # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step
     # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step
     # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step
-    point_coords[:, :, 0] = (point_indices % W).to(torch.float)  # x
-    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)  # y
-    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)  # z
-    print(f"resolution {D} x {H} x {W}", point_scores.min(),
-          point_scores.max())
+    point_coords[:, :, 0] = (point_indices % W).to(torch.float)    # x
+    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)    # y
+    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)    # z
+    print(f"resolution {D} x {H} x {W}", point_scores.min(), point_scores.max())
     return point_indices, point_coords
 
 
-def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points,
-                                                clip_min):
+def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points, clip_min):
     """
     Find `num_points` most uncertain points from `uncertainty_map` grid.
     Args:
@@ -276,28 +247,21 @@ def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points,
     uncertainty_map = uncertainty_map.view(D * H * W)
     indices = (uncertainty_map >= clip_min).nonzero().squeeze(1)
     num_points = min(num_points, indices.size(0))
-    point_scores, point_indices = torch.topk(uncertainty_map[indices],
-                                             k=num_points,
-                                             dim=0)
+    point_scores, point_indices = torch.topk(uncertainty_map[indices], k=num_points, dim=0)
     point_indices = indices[point_indices].unsqueeze(0)
 
-    point_coords = torch.zeros(R,
-                               num_points,
-                               3,
-                               dtype=torch.float,
-                               device=uncertainty_map.device)
+    point_coords = torch.zeros(R, num_points, 3, dtype=torch.float, device=uncertainty_map.device)
     # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step
     # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step
     # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step
-    point_coords[:, :, 0] = (point_indices % W).to(torch.float)  # x
-    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)  # y
-    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)  # z
+    point_coords[:, :, 0] = (point_indices % W).to(torch.float)    # x
+    point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float)    # y
+    point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float)    # z
     # print (f"resolution {D} x {H} x {W}", point_scores.min(), point_scores.max())
     return point_indices, point_coords
 
 
-def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points,
-                                         **kwargs):
+def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points, **kwargs):
     """
     Find `num_points` most uncertain points from `uncertainty_map` grid.
     Args:
@@ -315,14 +279,8 @@ def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points,
     # w_step = 1.0 / float(W)
 
     num_points = min(H * W, num_points)
-    point_scores, point_indices = torch.topk(uncertainty_map.view(R, H * W),
-                                             k=num_points,
-                                             dim=1)
-    point_coords = torch.zeros(R,
-                               num_points,
-                               2,
-                               dtype=torch.long,
-                               device=uncertainty_map.device)
+    point_scores, point_indices = torch.topk(uncertainty_map.view(R, H * W), k=num_points, dim=1)
+    point_coords = torch.zeros(R, num_points, 2, dtype=torch.long, device=uncertainty_map.device)
     # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step
     # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step
     point_coords[:, :, 0] = (point_indices % W).to(torch.long)
@@ -331,8 +289,7 @@ def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points,
     return point_indices, point_coords
 
 
-def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points,
-                                                clip_min):
+def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points, clip_min):
     """
     Find `num_points` most uncertain points from `uncertainty_map` grid.
     Args:
@@ -353,16 +310,10 @@ def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points,
     uncertainty_map = uncertainty_map.view(H * W)
     indices = (uncertainty_map >= clip_min).nonzero().squeeze(1)
     num_points = min(num_points, indices.size(0))
-    point_scores, point_indices = torch.topk(uncertainty_map[indices],
-                                             k=num_points,
-                                             dim=0)
+    point_scores, point_indices = torch.topk(uncertainty_map[indices], k=num_points, dim=0)
     point_indices = indices[point_indices].unsqueeze(0)
 
-    point_coords = torch.zeros(R,
-                               num_points,
-                               2,
-                               dtype=torch.long,
-                               device=uncertainty_map.device)
+    point_coords = torch.zeros(R, num_points, 2, dtype=torch.long, device=uncertainty_map.device)
     # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step
     # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step
     point_coords[:, :, 0] = (point_indices % W).to(torch.long)
@@ -388,7 +339,6 @@ def calculate_uncertainty(logits, classes=None, balance_value=0.5):
     if logits.shape[1] == 1:
         gt_class_logits = logits
     else:
-        gt_class_logits = logits[
-            torch.arange(logits.shape[0], device=logits.device),
-            classes].unsqueeze(1)
+        gt_class_logits = logits[torch.arange(logits.shape[0], device=logits.device),
+                                 classes].unsqueeze(1)
     return -torch.abs(gt_class_logits - balance_value)

lib/common/train_util.py

@@ -14,63 +14,62 @@
 #
 # Contact: [email protected]
 
-import yaml
-import os.path as osp
 import torch
-import numpy as np
 from ..dataset.mesh_util import *
 from ..net.geometry import orthogonal
-import cv2, PIL
-from tqdm import tqdm
-import os
 from termcolor import colored
 import pytorch_lightning as pl
 
 
+class Format:
+    end = '\033[0m'
+    start = '\033[4m'
+
+
 def init_loss():
 
     losses = {
-        # Cloth: Normal_recon - Normal_pred
+    # Cloth: chamfer distance
         "cloth": {
             "weight": 1e3,
             "value": 0.0
         },
-        # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2)
-        "stiffness": {
+    # Stiffness: [RT]_v1 - [RT]_v2 (v1-edge-v2)
+        "stiff": {
             "weight": 1e5,
             "value": 0.0
         },
-        # Cloth: det(R) = 1
+    # Cloth: det(R) = 1
         "rigid": {
             "weight": 1e5,
             "value": 0.0
         },
-        # Cloth: edge length
+    # Cloth: edge length
         "edge": {
             "weight": 0,
             "value": 0.0
         },
-        # Cloth: normal consistency
+    # Cloth: normal consistency
         "nc": {
             "weight": 0,
             "value": 0.0
         },
-        # Cloth: laplacian smoonth
-        "laplacian": {
+    # Cloth: laplacian smoonth
+        "lapla": {
             "weight": 1e2,
             "value": 0.0
         },
-        # Body: Normal_pred - Normal_smpl
+    # Body: Normal_pred - Normal_smpl
         "normal": {
             "weight": 1e0,
             "value": 0.0
         },
-        # Body: Silhouette_pred - Silhouette_smpl
+    # Body: Silhouette_pred - Silhouette_smpl
         "silhouette": {
             "weight": 1e0,
             "value": 0.0
         },
-        # Joint: reprojected joints difference
+    # Joint: reprojected joints difference
         "joint": {
             "weight": 5e0,
             "value": 0.0
@@ -81,7 +80,6 @@ def init_loss():
 
 
 class SubTrainer(pl.Trainer):
-
     def save_checkpoint(self, filepath, weights_only=False):
         """Save model/training states as a checkpoint file through state-dump and file-write.
         Args:
@@ -101,214 +99,6 @@ class SubTrainer(pl.Trainer):
         pl.utilities.cloud_io.atomic_save(_checkpoint, filepath)
 
 
-def rename(old_dict, old_name, new_name):
-    new_dict = {}
-    for key, value in zip(old_dict.keys(), old_dict.values()):
-        new_key = key if key != old_name else new_name
-        new_dict[new_key] = old_dict[key]
-    return new_dict
-
-
-def load_normal_networks(model, normal_path):
-    
-    pretrained_dict = torch.load(
-            normal_path,
-            map_location=model.device)["state_dict"]
-    model_dict = model.state_dict()
-
-    # 1. filter out unnecessary keys
-    pretrained_dict = {
-        k: v
-        for k, v in pretrained_dict.items()
-        if k in model_dict and v.shape == model_dict[k].shape
-    }
-
-    # # 2. overwrite entries in the existing state dict
-    model_dict.update(pretrained_dict)
-    # 3. load the new state dict
-    model.load_state_dict(model_dict)
-
-    del pretrained_dict
-    del model_dict
-
-    print(colored(f"Resume Normal weights from {normal_path}", "green"))
-
-
-def load_networks(model, mlp_path, normal_path=None):
-
-    model_dict = model.state_dict()
-    main_dict = {}
-    normal_dict = {}
-    
-    # MLP part loading
-    if os.path.exists(mlp_path) and mlp_path.endswith("ckpt"):
-        main_dict = torch.load(
-            mlp_path,
-            map_location=model.device)["state_dict"]
-
-        main_dict = {
-            k: v
-            for k, v in main_dict.items()
-            if k in model_dict and v.shape == model_dict[k].shape and (
-                "reconEngine" not in k) and ("normal_filter" not in k) and (
-                    "voxelization" not in k)
-        }
-        print(colored(f"Resume MLP weights from {mlp_path}", "green"))
-
-    # normal network part loading
-    if normal_path is not None and os.path.exists(normal_path) and normal_path.endswith("ckpt"):
-        normal_dict = torch.load(
-            normal_path,
-            map_location=model.device)["state_dict"]
-
-        for key in normal_dict.keys():
-            normal_dict = rename(normal_dict, key,
-                                 key.replace("netG", "netG.normal_filter"))
-
-        normal_dict = {
-            k: v
-            for k, v in normal_dict.items()
-            if k in model_dict and v.shape == model_dict[k].shape
-        }
-        print(colored(f"Resume normal model from {normal_path}", "green"))
-
-    model_dict.update(main_dict)
-    model_dict.update(normal_dict)
-    model.load_state_dict(model_dict)
-
-    # clean unused GPU memory
-    del main_dict
-    del normal_dict
-    del model_dict
-    torch.cuda.empty_cache()
-
-
-def reshape_sample_tensor(sample_tensor, num_views):
-    if num_views == 1:
-        return sample_tensor
-    # Need to repeat sample_tensor along the batch dim num_views times
-    sample_tensor = sample_tensor.unsqueeze(dim=1)
-    sample_tensor = sample_tensor.repeat(1, num_views, 1, 1)
-    sample_tensor = sample_tensor.view(
-        sample_tensor.shape[0] * sample_tensor.shape[1],
-        sample_tensor.shape[2],
-        sample_tensor.shape[3],
-    )
-    return sample_tensor
-
-
-def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma):
-    """Sets the learning rate to the initial LR decayed by schedule"""
-    if epoch in schedule:
-        lr *= gamma
-        for param_group in optimizer.param_groups:
-            param_group["lr"] = lr
-    return lr
-
-
-def compute_acc(pred, gt, thresh=0.5):
-    """
-    return:
-        IOU, precision, and recall
-    """
-    with torch.no_grad():
-        vol_pred = pred > thresh
-        vol_gt = gt > thresh
-
-        union = vol_pred | vol_gt
-        inter = vol_pred & vol_gt
-
-        true_pos = inter.sum().float()
-
-        union = union.sum().float()
-        if union == 0:
-            union = 1
-        vol_pred = vol_pred.sum().float()
-        if vol_pred == 0:
-            vol_pred = 1
-        vol_gt = vol_gt.sum().float()
-        if vol_gt == 0:
-            vol_gt = 1
-        return true_pos / union, true_pos / vol_pred, true_pos / vol_gt
-
-def calc_error(opt, net, cuda, dataset, num_tests):
-    if num_tests > len(dataset):
-        num_tests = len(dataset)
-    with torch.no_grad():
-        erorr_arr, IOU_arr, prec_arr, recall_arr = [], [], [], []
-        for idx in tqdm(range(num_tests)):
-            data = dataset[idx * len(dataset) // num_tests]
-            # retrieve the data
-            image_tensor = data["img"].to(device=cuda)
-            calib_tensor = data["calib"].to(device=cuda)
-            sample_tensor = data["samples"].to(device=cuda).unsqueeze(0)
-            if opt.num_views > 1:
-                sample_tensor = reshape_sample_tensor(sample_tensor,
-                                                      opt.num_views)
-            label_tensor = data["labels"].to(device=cuda).unsqueeze(0)
-
-            res, error = net.forward(image_tensor,
-                                     sample_tensor,
-                                     calib_tensor,
-                                     labels=label_tensor)
-
-            IOU, prec, recall = compute_acc(res, label_tensor)
-
-            # print(
-            #     '{0}/{1} | Error: {2:06f} IOU: {3:06f} prec: {4:06f} recall: {5:06f}'
-            #         .format(idx, num_tests, error.item(), IOU.item(), prec.item(), recall.item()))
-            erorr_arr.append(error.item())
-            IOU_arr.append(IOU.item())
-            prec_arr.append(prec.item())
-            recall_arr.append(recall.item())
-
-    return (
-        np.average(erorr_arr),
-        np.average(IOU_arr),
-        np.average(prec_arr),
-        np.average(recall_arr),
-    )
-
-
-def calc_error_color(opt, netG, netC, cuda, dataset, num_tests):
-    if num_tests > len(dataset):
-        num_tests = len(dataset)
-    with torch.no_grad():
-        error_color_arr = []
-
-        for idx in tqdm(range(num_tests)):
-            data = dataset[idx * len(dataset) // num_tests]
-            # retrieve the data
-            image_tensor = data["img"].to(device=cuda)
-            calib_tensor = data["calib"].to(device=cuda)
-            color_sample_tensor = data["color_samples"].to(
-                device=cuda).unsqueeze(0)
-
-            if opt.num_views > 1:
-                color_sample_tensor = reshape_sample_tensor(
-                    color_sample_tensor, opt.num_views)
-
-            rgb_tensor = data["rgbs"].to(device=cuda).unsqueeze(0)
-
-            netG.filter(image_tensor)
-            _, errorC = netC.forward(
-                image_tensor,
-                netG.get_im_feat(),
-                color_sample_tensor,
-                calib_tensor,
-                labels=rgb_tensor,
-            )
-
-            # print('{0}/{1} | Error inout: {2:06f} | Error color: {3:06f}'
-            #       .format(idx, num_tests, errorG.item(), errorC.item()))
-            error_color_arr.append(errorC.item())
-
-    return np.average(error_color_arr)
-
-
-# pytorch lightning training related fucntions
-
-
 def query_func(opt, netG, features, points, proj_matrix=None):
     """
         - points: size of (bz, N, 3)
@@ -317,7 +107,7 @@ def query_func(opt, netG, features, points, proj_matrix=None):
     """
     assert len(points) == 1
     samples = points.repeat(opt.num_views, 1, 1)
-    samples = samples.permute(0, 2, 1)  # [bz, 3, N]
+    samples = samples.permute(0, 2, 1)    # [bz, 3, N]
 
     # view specific query
     if proj_matrix is not None:
@@ -337,85 +127,25 @@ def query_func(opt, netG, features, points, proj_matrix=None):
 
     return preds
 
+
 def query_func_IF(batch, netG, points):
     """
         - points: size of (bz, N, 3)
     return: size of (bz, 1, N)
     """
-    
+
     batch["samples_geo"] = points
     batch["calib"] = torch.stack([torch.eye(4).float()], dim=0).type_as(points)
-    
+
     preds = netG(batch)
 
     return preds.unsqueeze(1)
 
 
-def isin(ar1, ar2):
-    return (ar1[..., None] == ar2).any(-1)
-
-
-def in1d(ar1, ar2):
-    mask = ar2.new_zeros((max(ar1.max(), ar2.max()) + 1, ), dtype=torch.bool)
-    mask[ar2.unique()] = True
-    return mask[ar1]
-
 def 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()
-
-
-def tf_log_convert(log_dict):
-    new_log_dict = log_dict.copy()
-    for k, v in log_dict.items():
-        new_log_dict[k.replace("_", "/")] = v
-        del new_log_dict[k]
-
-    return new_log_dict
-
-
-def bar_log_convert(log_dict, name=None, rot=None):
-    from decimal import Decimal
-
-    new_log_dict = {}
-
-    if name is not None:
-        new_log_dict["name"] = name[0]
-    if rot is not None:
-        new_log_dict["rot"] = rot[0]
-
-    for k, v in log_dict.items():
-        color = "yellow"
-        if "loss" in k:
-            color = "red"
-            k = k.replace("loss", "L")
-        elif "acc" in k:
-            color = "green"
-            k = k.replace("acc", "A")
-        elif "iou" in k:
-            color = "green"
-            k = k.replace("iou", "I")
-        elif "prec" in k:
-            color = "green"
-            k = k.replace("prec", "P")
-        elif "recall" in k:
-            color = "green"
-            k = k.replace("recall", "R")
-
-        if "lr" not in k:
-            new_log_dict[colored(k.split("_")[1],
-                                 color)] = colored(f"{v:.3f}", color)
-        else:
-            new_log_dict[colored(k.split("_")[1],
-                                 color)] = colored(f"{Decimal(str(v)):.1E}",
-                                                   color)
-
-    if "loss" in new_log_dict.keys():
-        del new_log_dict["loss"]
-
-    return new_log_dict
+    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):
@@ -430,160 +160,10 @@ def accumulate(outputs, rot_num, split):
             keyword = f"{dataset}/{metric}"
             if keyword not in hparam_log_dict.keys():
                 hparam_log_dict[keyword] = 0
-            for idx in range(split[dataset][0] * rot_num,
-                             split[dataset][1] * rot_num):
+            for idx in range(split[dataset][0] * rot_num, split[dataset][1] * rot_num):
                 hparam_log_dict[keyword] += outputs[idx][metric].item()
-            hparam_log_dict[keyword] /= (split[dataset][1] -
-                                         split[dataset][0]) * rot_num
+            hparam_log_dict[keyword] /= (split[dataset][1] - split[dataset][0]) * rot_num
 
     print(colored(hparam_log_dict, "green"))
 
     return hparam_log_dict
-
-
-def calc_error_N(outputs, targets):
-    """calculate the error of normal (IGR)
-
-    Args:
-        outputs (torch.tensor): [B, 3, N]
-        target (torch.tensor): [B, N, 3]
-
-    # manifold loss and grad_loss in IGR paper
-    grad_loss = ((nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean()
-    normals_loss = ((mnfld_grad - normals).abs()).norm(2, dim=1).mean()
-
-    Returns:
-        torch.tensor: error of valid normals on the surface
-    """
-    # outputs = torch.tanh(-outputs.permute(0,2,1).reshape(-1,3))
-    outputs = -outputs.permute(0, 2, 1).reshape(-1, 1)
-    targets = targets.reshape(-1, 3)[:, 2:3]
-    with_normals = targets.sum(dim=1).abs() > 0.0
-
-    # eikonal loss
-    grad_loss = ((outputs[with_normals].norm(2, dim=-1) - 1)**2).mean()
-    # normals loss
-    normal_loss = (outputs - targets)[with_normals].abs().norm(2, dim=1).mean()
-
-    return grad_loss * 0.0 + normal_loss
-
-
-def calc_knn_acc(preds, carn_verts, labels, pick_num):
-    """calculate knn accuracy
-
-    Args:
-        preds (torch.tensor): [B, 3, N]
-        carn_verts (torch.tensor): [SMPLX_V_num, 3]
-        labels (torch.tensor): [B, N_knn, N]
-    """
-    N_knn_full = labels.shape[1]
-    preds = preds.permute(0, 2, 1).reshape(-1, 3)
-    labels = labels.permute(0, 2, 1).reshape(-1, N_knn_full)  # [BxN, num_knn]
-    labels = labels[:, :pick_num]
-
-    dist = torch.cdist(preds, carn_verts, p=2)  # [BxN, SMPL_V_num]
-    knn = dist.topk(k=pick_num, dim=1, largest=False)[1]  # [BxN, num_knn]
-    cat_mat = torch.sort(torch.cat((knn, labels), dim=1))[0]
-    bool_col = torch.zeros_like(cat_mat)[:, 0]
-    for i in range(pick_num * 2 - 1):
-        bool_col += cat_mat[:, i] == cat_mat[:, i + 1]
-    acc = (bool_col > 0).sum() / len(bool_col)
-
-    return acc
-
-
-def calc_acc_seg(output, target, num_multiseg):
-    from pytorch_lightning.metrics import Accuracy
-
-    return Accuracy()(output.reshape(-1, num_multiseg).cpu(),
-                      target.flatten().cpu())
-
-
-def add_watermark(imgs, titles):
-
-    # Write some Text
-
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    bottomLeftCornerOfText = (350, 50)
-    bottomRightCornerOfText = (800, 50)
-    fontScale = 1
-    fontColor = (1.0, 1.0, 1.0)
-    lineType = 2
-
-    for i in range(len(imgs)):
-
-        title = titles[i + 1]
-        cv2.putText(imgs[i], title, bottomLeftCornerOfText, font, fontScale,
-                    fontColor, lineType)
-
-        if i == 0:
-            cv2.putText(
-                imgs[i],
-                str(titles[i][0]),
-                bottomRightCornerOfText,
-                font,
-                fontScale,
-                fontColor,
-                lineType,
-            )
-
-    result = np.concatenate(imgs, axis=0).transpose(2, 0, 1)
-
-    return result
-
-
-def make_test_gif(img_dir):
-
-    if img_dir is not None and len(os.listdir(img_dir)) > 0:
-        for dataset in os.listdir(img_dir):
-            for subject in sorted(os.listdir(osp.join(img_dir, dataset))):
-                img_lst = []
-                im1 = None
-                for file in sorted(
-                        os.listdir(osp.join(img_dir, dataset, subject))):
-                    if file[-3:] not in ["obj", "gif"]:
-                        img_path = os.path.join(img_dir, dataset, subject,
-                                                file)
-                        if im1 == None:
-                            im1 = PIL.Image.open(img_path)
-                        else:
-                            img_lst.append(PIL.Image.open(img_path))
-
-                print(os.path.join(img_dir, dataset, subject, "out.gif"))
-                im1.save(
-                    os.path.join(img_dir, dataset, subject, "out.gif"),
-                    save_all=True,
-                    append_images=img_lst,
-                    duration=500,
-                    loop=0,
-                )
-
-
-def export_cfg(logger, dir, cfg):
-
-    cfg_export_file = osp.join(dir, f"cfg_{logger.version}.yaml")
-
-    if not osp.exists(cfg_export_file):
-        os.makedirs(osp.dirname(cfg_export_file), exist_ok=True)
-        with open(cfg_export_file, "w+") as file:
-            _ = yaml.dump(cfg, file)
-
-
-from yacs.config import CfgNode
-
-_VALID_TYPES = {tuple, list, str, int, float, bool}
-
-
-def convert_to_dict(cfg_node, key_list=[]):
-    """ Convert a config node to dictionary """
-    if not isinstance(cfg_node, CfgNode):
-        if type(cfg_node) not in _VALID_TYPES:
-            print(
-                "Key {} with value {} is not a valid type; valid types: {}".
-                format(".".join(key_list), type(cfg_node), _VALID_TYPES), )
-        return cfg_node
-    else:
-        cfg_dict = dict(cfg_node)
-        for k, v in cfg_dict.items():
-            cfg_dict[k] = convert_to_dict(v, key_list + [k])
-        return cfg_dict

lib/common/voxelize.py

@@ -13,6 +13,7 @@ from lib.common.libmesh.inside_mesh import check_mesh_contains
 
 # From Occupancy Networks, Mescheder et. al. CVPR'19
 
+
 def make_3d_grid(bb_min, bb_max, shape):
     ''' Makes a 3D grid.
 
@@ -37,7 +38,7 @@ def make_3d_grid(bb_min, bb_max, shape):
 
 class VoxelGrid:
     def __init__(self, data, loc=(0., 0., 0.), scale=1):
-        assert(data.shape[0] == data.shape[1] == data.shape[2])
+        assert (data.shape[0] == data.shape[1] == data.shape[2])
         data = np.asarray(data, dtype=np.bool)
         loc = np.asarray(loc)
         self.data = data
@@ -53,7 +54,7 @@ class VoxelGrid:
 
         # Default scale, scales the mesh to [-0.45, 0.45]^3
         if scale is None:
-            scale = (bounds[1] - bounds[0]).max()/0.9
+            scale = (bounds[1] - bounds[0]).max() / 0.9
 
         loc = np.asarray(loc)
         scale = float(scale)
@@ -61,7 +62,7 @@ class VoxelGrid:
         # Transform mesh
         mesh = mesh.copy()
         mesh.apply_translation(-loc)
-        mesh.apply_scale(1/scale)
+        mesh.apply_scale(1 / scale)
 
         # Apply method
         if method == 'ray':
@@ -75,7 +76,7 @@ class VoxelGrid:
     def down_sample(self, factor=2):
         if not (self.resolution % factor) == 0:
             raise ValueError('Resolution must be divisible by factor.')
-        new_data = block_reduce(self.data, (factor,) * 3, np.max)
+        new_data = block_reduce(self.data, (factor, ) * 3, np.max)
         return VoxelGrid(new_data, self.loc, self.scale)
 
     def to_mesh(self):
@@ -103,9 +104,9 @@ class VoxelGrid:
         f2 = f2_r | f2_l
         f3 = f3_r | f3_l
 
-        assert(f1.shape == (nx + 1, ny, nz))
-        assert(f2.shape == (nx, ny + 1, nz))
-        assert(f3.shape == (nx, ny, nz + 1))
+        assert (f1.shape == (nx + 1, ny, nz))
+        assert (f2.shape == (nx, ny + 1, nz))
+        assert (f3.shape == (nx, ny, nz + 1))
 
         # Determine if vertex present
         v = np.full(grid_shape, False)
@@ -146,53 +147,76 @@ 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)
@@ -200,7 +224,7 @@ class VoxelGrid:
 
     @property
     def resolution(self):
-        assert(self.data.shape[0] == self.data.shape[1] == self.data.shape[2])
+        assert (self.data.shape[0] == self.data.shape[1] == self.data.shape[2])
         return self.data.shape[0]
 
     def contains(self, points):
@@ -211,12 +235,9 @@ class VoxelGrid:
         # Discretize points to [0, nx-1]^3
         points_i = ((points + 0.5) * nx).astype(np.int32)
         # i1, i2, i3 have sizes (batch_size, T)
-        i1, i2, i3 = points_i[..., 0],  points_i[..., 1],  points_i[..., 2]
+        i1, i2, i3 = points_i[..., 0], points_i[..., 1], points_i[..., 2]
         # Only use indices inside bounding box
-        mask = (
-            (i1 >= 0) & (i2 >= 0) & (i3 >= 0)
-            & (nx > i1) & (nx > i2) & (nx > i3)
-        )
+        mask = ((i1 >= 0) & (i2 >= 0) & (i3 >= 0) & (nx > i1) & (nx > i2) & (nx > i3))
         # Prevent out of bounds error
         i1 = i1[mask]
         i2 = i2[mask]
@@ -254,7 +275,7 @@ def voxelize_surface(mesh, resolution):
     vertices = (vertices + 0.5) * resolution
 
     face_loc = vertices[faces]
-    occ = np.full((resolution,) * 3, 0, dtype=np.int32)
+    occ = np.full((resolution, ) * 3, 0, dtype=np.int32)
     face_loc = face_loc.astype(np.float32)
 
     voxelize_mesh_(occ, face_loc)
@@ -264,9 +285,9 @@ def voxelize_surface(mesh, resolution):
 
 
 def voxelize_interior(mesh, resolution):
-    shape = (resolution,) * 3
-    bb_min = (0.5,) * 3
-    bb_max = (resolution - 0.5,) * 3
+    shape = (resolution, ) * 3
+    bb_min = (0.5, ) * 3
+    bb_max = (resolution - 0.5, ) * 3
     # Create points. Add noise to break symmetry
     points = make_3d_grid(bb_min, bb_max, shape=shape).numpy()
     points = points + 0.1 * (np.random.rand(*points.shape) - 0.5)
@@ -280,14 +301,9 @@ def check_voxel_occupied(occupancy_grid):
     occ = occupancy_grid
 
     occupied = (
-        occ[..., :-1, :-1, :-1]
-        & occ[..., :-1, :-1, 1:]
-        & occ[..., :-1, 1:, :-1]
-        & occ[..., :-1, 1:, 1:]
-        & occ[..., 1:, :-1, :-1]
-        & occ[..., 1:, :-1, 1:]
-        & occ[..., 1:, 1:, :-1]
-        & occ[..., 1:, 1:, 1:]
+        occ[..., :-1, :-1, :-1] & occ[..., :-1, :-1, 1:] & occ[..., :-1, 1:, :-1] &
+        occ[..., :-1, 1:, 1:] & occ[..., 1:, :-1, :-1] & occ[..., 1:, :-1, 1:] &
+        occ[..., 1:, 1:, :-1] & occ[..., 1:, 1:, 1:]
     )
     return occupied
 
@@ -296,14 +312,9 @@ def check_voxel_unoccupied(occupancy_grid):
     occ = occupancy_grid
 
     unoccupied = ~(
-        occ[..., :-1, :-1, :-1]
-        | occ[..., :-1, :-1, 1:]
-        | occ[..., :-1, 1:, :-1]
-        | occ[..., :-1, 1:, 1:]
-        | occ[..., 1:, :-1, :-1]
-        | occ[..., 1:, :-1, 1:]
-        | occ[..., 1:, 1:, :-1]
-        | occ[..., 1:, 1:, 1:]
+        occ[..., :-1, :-1, :-1] | occ[..., :-1, :-1, 1:] | occ[..., :-1, 1:, :-1] |
+        occ[..., :-1, 1:, 1:] | occ[..., 1:, :-1, :-1] | occ[..., 1:, :-1, 1:] |
+        occ[..., 1:, 1:, :-1] | occ[..., 1:, 1:, 1:]
     )
     return unoccupied
 

lib/dataset/Evaluator.py

@@ -37,7 +37,6 @@ class _PointFaceDistance(Function):
     """
     Torch autograd Function wrapper PointFaceDistance Cuda implementation
     """
-
     @staticmethod
     def forward(
         ctx,
@@ -92,12 +91,15 @@ class _PointFaceDistance(Function):
         grad_dists = grad_dists.contiguous()
         points, tris, idxs = ctx.saved_tensors
         min_triangle_area = ctx.min_triangle_area
-        grad_points, grad_tris = _C.point_face_dist_backward(points, tris, idxs, grad_dists, min_triangle_area)
+        grad_points, grad_tris = _C.point_face_dist_backward(
+            points, tris, idxs, grad_dists, min_triangle_area
+        )
         return grad_points, None, grad_tris, None, None, None
 
 
-def _rand_barycentric_coords(size1, size2, dtype: torch.dtype,
-                             device: torch.device) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+def _rand_barycentric_coords(
+    size1, size2, dtype: torch.dtype, device: torch.device
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Helper function to generate random barycentric coordinates which are uniformly
     distributed over a triangle.
@@ -167,19 +169,21 @@ def sample_points_from_meshes(meshes, num_samples: int = 10000):
     faces = meshes.faces_packed()
     mesh_to_face = meshes.mesh_to_faces_packed_first_idx()
     num_meshes = len(meshes)
-    num_valid_meshes = torch.sum(meshes.valid)  # Non empty meshes.
+    num_valid_meshes = torch.sum(meshes.valid)    # Non empty meshes.
 
     # Initialize samples tensor with fill value 0 for empty meshes.
     samples = torch.zeros((num_meshes, num_samples, 3), device=meshes.device)
 
     # Only compute samples for non empty meshes
     with torch.no_grad():
-        areas, _ = mesh_face_areas_normals(verts, faces)  # Face areas can be zero.
+        areas, _ = mesh_face_areas_normals(verts, faces)    # Face areas can be zero.
         max_faces = meshes.num_faces_per_mesh().max().item()
-        areas_padded = packed_to_padded(areas, mesh_to_face[meshes.valid], max_faces)  # (N, F)
+        areas_padded = packed_to_padded(areas, mesh_to_face[meshes.valid], max_faces)    # (N, F)
 
         # TODO (gkioxari) Confirm multinomial bug is not present with real data.
-        samples_face_idxs = areas_padded.multinomial(num_samples, replacement=True)  # (N, num_samples)
+        samples_face_idxs = areas_padded.multinomial(
+            num_samples, replacement=True
+        )    # (N, num_samples)
         samples_face_idxs += mesh_to_face[meshes.valid].view(num_valid_meshes, 1)
 
     # Randomly generate barycentric coords.
@@ -200,23 +204,25 @@ def point_mesh_distance(meshes, pcls, weighted=True):
         raise ValueError("meshes and pointclouds must be equal sized batches")
 
     # packed representation for pointclouds
-    points = pcls.points_packed()  # (P, 3)
+    points = pcls.points_packed()    # (P, 3)
     points_first_idx = pcls.cloud_to_packed_first_idx()
     max_points = pcls.num_points_per_cloud().max().item()
 
     # packed representation for faces
     verts_packed = meshes.verts_packed()
     faces_packed = meshes.faces_packed()
-    tris = verts_packed[faces_packed]  # (T, 3, 3)
+    tris = verts_packed[faces_packed]    # (T, 3, 3)
     tris_first_idx = meshes.mesh_to_faces_packed_first_idx()
 
     # point to face distance: shape (P,)
-    point_to_face, idxs = _PointFaceDistance.apply(points, points_first_idx, tris, tris_first_idx, max_points, 5e-3)
+    point_to_face, idxs = _PointFaceDistance.apply(
+        points, points_first_idx, tris, tris_first_idx, max_points, 5e-3
+    )
 
     if weighted:
         # weight each example by the inverse of number of points in the example
-        point_to_cloud_idx = pcls.packed_to_cloud_idx()  # (sum(P_i),)
-        num_points_per_cloud = pcls.num_points_per_cloud()  # (N,)
+        point_to_cloud_idx = pcls.packed_to_cloud_idx()    # (sum(P_i),)
+        num_points_per_cloud = pcls.num_points_per_cloud()    # (N,)
         weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
         weights_p = 1.0 / weights_p.float()
         point_to_face = torch.sqrt(point_to_face) * weights_p
@@ -225,7 +231,6 @@ def point_mesh_distance(meshes, pcls, weighted=True):
 
 
 class Evaluator:
-
     def __init__(self, device):
 
         self.render = Render(size=512, device=device)
@@ -253,8 +258,8 @@ class Evaluator:
         self.render.meshes = self.tgt_mesh
         tgt_normal_imgs = self.render.get_image(cam_type="four", bg="black")
 
-        src_normal_arr = make_grid(torch.cat(src_normal_imgs, dim=0), nrow=4, padding=0)  # [-1,1]
-        tgt_normal_arr = make_grid(torch.cat(tgt_normal_imgs, dim=0), nrow=4, padding=0)  # [-1,1]
+        src_normal_arr = make_grid(torch.cat(src_normal_imgs, dim=0), nrow=4, padding=0)    # [-1,1]
+        tgt_normal_arr = make_grid(torch.cat(tgt_normal_imgs, dim=0), nrow=4, padding=0)    # [-1,1]
         src_norm = torch.norm(src_normal_arr, dim=0, keepdim=True)
         tgt_norm = torch.norm(tgt_normal_arr, dim=0, keepdim=True)
 
@@ -274,8 +279,11 @@ 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))
+            (
+                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
@@ -291,7 +299,9 @@ class Evaluator:
         p2s_dist_all, _ = point_mesh_distance(self.src_mesh, tgt_points) * 100.0
         p2s_dist = p2s_dist_all.sum()
 
-        chamfer_dist = (point_mesh_distance(self.tgt_mesh, src_points)[0].sum() * 100.0 + p2s_dist) * 0.5
+        chamfer_dist = (
+            point_mesh_distance(self.tgt_mesh, src_points)[0].sum() * 100.0 + p2s_dist
+        ) * 0.5
 
         return chamfer_dist, p2s_dist
 

lib/dataset/NormalDataset.py

@@ -23,7 +23,6 @@ import torchvision.transforms as transforms
 
 
 class NormalDataset:
-
     def __init__(self, cfg, split="train"):
 
         self.split = split
@@ -44,8 +43,7 @@ class NormalDataset:
         if self.split != "train":
             self.rotations = range(0, 360, 120)
         else:
-            self.rotations = np.arange(0, 360, 360 //
-                                       self.opt.rotation_num).astype(np.int)
+            self.rotations = np.arange(0, 360, 360 // self.opt.rotation_num).astype(np.int)
 
         self.datasets_dict = {}
 
@@ -54,26 +52,29 @@ class NormalDataset:
             dataset_dir = osp.join(self.root, dataset)
 
             self.datasets_dict[dataset] = {
-                "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"),
-                                       dtype=str),
+                "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str),
                 "scale": self.scales[dataset_id],
             }
 
         self.subject_list = self.get_subject_list(split)
 
         # PIL to tensor
-        self.image_to_tensor = transforms.Compose([
-            transforms.Resize(self.input_size),
-            transforms.ToTensor(),
-            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-        ])
+        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):
 
@@ -88,16 +89,12 @@ class NormalDataset:
                 subject_list += np.loadtxt(split_txt, dtype=str).tolist()
 
         if self.split != "test":
-            subject_list += subject_list[:self.bsize -
-                                         len(subject_list) % self.bsize]
+            subject_list += subject_list[:self.bsize - len(subject_list) % self.bsize]
             print(colored(f"total: {len(subject_list)}", "yellow"))
 
-        bug_list = sorted(
-            np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist())
+        bug_list = sorted(np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist())
 
-        subject_list = [
-            subject for subject in subject_list if (subject not in bug_list)
-        ]
+        subject_list = [subject for subject in subject_list if (subject not in bug_list)]
 
         # subject_list = ["thuman2/0008"]
         return subject_list
@@ -113,48 +110,41 @@ class NormalDataset:
         rotation = self.rotations[rid]
         subject = self.subject_list[mid].split("/")[1]
         dataset = self.subject_list[mid].split("/")[0]
-        render_folder = "/".join(
-            [dataset + f"_{self.opt.rotation_num}views", subject])
+        render_folder = "/".join([dataset + f"_{self.opt.rotation_num}views", subject])
 
         if not osp.exists(osp.join(self.root, render_folder)):
             render_folder = "/".join([dataset + f"_36views", subject])
 
         # setup paths
         data_dict = {
-            "dataset":
-            dataset,
-            "subject":
-            subject,
-            "rotation":
-            rotation,
-            "scale":
-            self.datasets_dict[dataset]["scale"],
-            "image_path":
-            osp.join(self.root, render_folder, "render",
-                     f"{rotation:03d}.png"),
+            "dataset": dataset,
+            "subject": subject,
+            "rotation": rotation,
+            "scale": self.datasets_dict[dataset]["scale"],
+            "image_path": osp.join(self.root, render_folder, "render", f"{rotation:03d}.png"),
         }
 
         # image/normal/depth loader
         for name, channel in zip(self.in_total, self.in_total_dim):
 
             if f"{name}_path" not in data_dict.keys():
-                data_dict.update({
-                    f"{name}_path":
-                    osp.join(self.root, render_folder, name,
-                             f"{rotation:03d}.png")
-                })
-
-            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
-        ]
+                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]
 
         for key in path_keys:
             del data_dict[key]
@@ -172,10 +162,9 @@ class NormalDataset:
 
         # simulate occlusion
         if erasing:
-            mask = kornia.augmentation.RandomErasing(p=0.2,
-                                                     scale=(0.01, 0.2),
-                                                     ratio=(0.3, 3.3),
-                                                     keepdim=True)(mask)
+            mask = kornia.augmentation.RandomErasing(
+                p=0.2, scale=(0.01, 0.2), ratio=(0.3, 3.3), keepdim=True
+            )(mask)
         image = (image * mask)[:channel]
 
         return (image * (0.5 - inv) * 2.0).float()

lib/dataset/NormalModule.py

@@ -22,7 +22,6 @@ import pytorch_lightning as pl
 
 
 class NormalModule(pl.LightningDataModule):
-
     def __init__(self, cfg):
         super(NormalModule, self).__init__()
         self.cfg = cfg
@@ -40,7 +39,7 @@ class NormalModule(pl.LightningDataModule):
         self.train_dataset = NormalDataset(cfg=self.cfg, split="train")
         self.val_dataset = NormalDataset(cfg=self.cfg, split="val")
         self.test_dataset = NormalDataset(cfg=self.cfg, split="test")
-        
+
         self.data_size = {
             "train": len(self.train_dataset),
             "val": len(self.val_dataset),
@@ -69,7 +68,7 @@ class NormalModule(pl.LightningDataModule):
         )
 
         return val_data_loader
-    
+
     def val_dataloader(self):
 
         test_data_loader = DataLoader(

lib/dataset/PointFeat.py

@@ -6,7 +6,6 @@ from lib.dataset.mesh_util import SMPLX, barycentric_coordinates_of_projection
 
 
 class PointFeat:
-
     def __init__(self, verts, faces):
 
         # verts [B, N_vert, 3]
@@ -23,7 +22,10 @@ class PointFeat:
 
         if verts.shape[1] == 10475:
             faces = faces[:, ~SMPLX().smplx_eyeball_fid_mask]
-            mouth_faces = (torch.as_tensor(SMPLX().smplx_mouth_fid).unsqueeze(0).repeat(self.Bsize, 1, 1).to(self.device))
+            mouth_faces = (
+                torch.as_tensor(SMPLX().smplx_mouth_fid).unsqueeze(0).repeat(self.Bsize, 1,
+                                                                             1).to(self.device)
+            )
             self.faces = torch.cat([faces, mouth_faces], dim=1).long()
 
         self.verts = verts.float()
@@ -35,11 +37,15 @@ class PointFeat:
         points = points.float()
         residues, pts_ind = point_mesh_distance(self.mesh, Pointclouds(points), weighted=False)
 
-        closest_triangles = torch.gather(self.triangles, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
+        closest_triangles = torch.gather(
+            self.triangles, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3)
+        ).view(-1, 3, 3)
         bary_weights = barycentric_coordinates_of_projection(points.view(-1, 3), closest_triangles)
 
         feat_normals = face_vertices(self.mesh.verts_normals_padded(), self.faces)
-        closest_normals = torch.gather(feat_normals, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3)
+        closest_normals = torch.gather(
+            feat_normals, 1, pts_ind[None, :, None, None].expand(-1, -1, 3, 3)
+        ).view(-1, 3, 3)
         shoot_verts = ((closest_triangles * bary_weights[:, :, None]).sum(1).unsqueeze(0))
 
         pts2shoot_normals = points - shoot_verts

lib/dataset/TestDataset.py

@@ -25,6 +25,7 @@ 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
@@ -36,8 +37,9 @@ from lib.dataset.body_model import TetraSMPLModel
 from lib.dataset.mesh_util import get_visibility, SMPLX
 import torch.nn.functional as F
 from torchvision import transforms
+from torchvision.models import detection
+
 import os.path as osp
-import os
 import torch
 import glob
 import numpy as np
@@ -48,7 +50,6 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class TestDataset:
-
     def __init__(self, cfg, device):
 
         self.image_dir = cfg["image_dir"]
@@ -65,7 +66,9 @@ class TestDataset:
         keep_lst = sorted(glob.glob(f"{self.image_dir}/*"))
         img_fmts = ["jpg", "png", "jpeg", "JPG", "bmp"]
 
-        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()
@@ -80,7 +83,16 @@ class TestDataset:
 
         self.smpl_model = PIXIE_SMPLX(pixie_cfg.model).to(self.device)
 
-        print(colored(f"Use {self.hps_type.upper()} to estimate human pose and shape", "green"))
+        self.detector = detection.maskrcnn_resnet50_fpn(
+            weights=detection.MaskRCNN_ResNet50_FPN_V2_Weights
+        )
+        self.detector.eval()
+
+        print(
+            colored(
+                f"SMPL-X estimate with {Format.start} {self.hps_type.upper()} {Format.end}", "green"
+            )
+        )
 
         self.render = Render(size=512, device=self.device)
 
@@ -90,7 +102,9 @@ class TestDataset:
     def compute_vis_cmap(self, smpl_verts, smpl_faces):
 
         (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=-1)
-        smpl_vis = get_visibility(xy, z, torch.as_tensor(smpl_faces).long()[:, :, [0, 2, 1]]).unsqueeze(-1)
+        smpl_vis = get_visibility(xy, z,
+                                  torch.as_tensor(smpl_faces).long()[:, :,
+                                                                     [0, 2, 1]]).unsqueeze(-1)
         smpl_cmap = self.smpl_data.cmap_smpl_vids(self.smpl_type).unsqueeze(0)
 
         return {
@@ -109,7 +123,8 @@ class TestDataset:
         depth_FB[:, ~depth_mask[0]] = 0.
 
         # Important: index_long = depth_value - 1
-        index_z = (((depth_FB + 1.) * 0.5 * self.vol_res) - 1).clip(0, self.vol_res - 1).permute(1, 2, 0)
+        index_z = (((depth_FB + 1.) * 0.5 * self.vol_res) - 1).clip(0, self.vol_res -
+                                                                    1).permute(1, 2, 0)
         index_z_ceil = torch.ceil(index_z).long()
         index_z_floor = torch.floor(index_z).long()
         index_z_frac = torch.frac(index_z)
@@ -121,7 +136,7 @@ class TestDataset:
             F.one_hot(index_z_floor[..., 1], self.vol_res) * (1.0 - index_z_frac[..., 1])
 
         voxels[index_mask] *= 0
-        voxels = torch.flip(voxels, [2]).permute(2, 0, 1).float()  #[x-2, y-0, z-1]
+        voxels = torch.flip(voxels, [2]).permute(2, 0, 1).float()    #[x-2, y-0, z-1]
 
         return {
             "depth_voxels": voxels.flip([
@@ -139,18 +154,25 @@ class TestDataset:
         smpl_model.set_params(rotation_matrix_to_angle_axis(rot6d_to_rotmat(pose)), beta=betas[0])
 
         verts = (
-            np.concatenate([smpl_model.verts, smpl_model.verts_added], axis=0) * scale.item() + trans.detach().cpu().numpy())
+            np.concatenate([smpl_model.verts, smpl_model.verts_added], axis=0) * scale.item() +
+            trans.detach().cpu().numpy()
+        )
         faces = (
             np.loadtxt(
                 osp.join(self.smpl_data.tedra_dir, "tetrahedrons_neutral_adult.txt"),
                 dtype=np.int32,
-            ) - 1)
+            ) - 1
+        )
 
         pad_v_num = int(8000 - verts.shape[0])
         pad_f_num = int(25100 - faces.shape[0])
 
-        verts = (np.pad(verts, ((0, pad_v_num), (0, 0)), mode="constant", constant_values=0.0).astype(np.float32) * 0.5)
-        faces = np.pad(faces, ((0, pad_f_num), (0, 0)), mode="constant", constant_values=0.0).astype(np.int32)
+        verts = (
+            np.pad(verts, ((0, pad_v_num),
+                           (0, 0)), mode="constant", constant_values=0.0).astype(np.float32) * 0.5
+        )
+        faces = np.pad(faces, ((0, pad_f_num), (0, 0)), mode="constant",
+                       constant_values=0.0).astype(np.int32)
 
         verts[:, 2] *= -1.0
 
@@ -168,7 +190,7 @@ class TestDataset:
         img_path = self.subject_list[index]
         img_name = img_path.split("/")[-1].rsplit(".", 1)[0]
 
-        arr_dict = process_image(img_path, self.hps_type, self.single, 512)
+        arr_dict = process_image(img_path, self.hps_type, self.single, 512, self.detector)
         arr_dict.update({"name": img_name})
 
         with torch.no_grad():
@@ -179,7 +201,10 @@ class TestDataset:
                 preds_dict, _ = self.hps.forward(batch)
 
         arr_dict["smpl_faces"] = (
-            torch.as_tensor(self.smpl_data.smplx_faces.astype(np.int64)).unsqueeze(0).long().to(self.device))
+            torch.as_tensor(self.smpl_data.smplx_faces.astype(np.int64)).unsqueeze(0).long().to(
+                self.device
+            )
+        )
         arr_dict["type"] = self.smpl_type
 
         if self.hps_type == "pymafx":
@@ -198,13 +223,16 @@ class TestDataset:
         elif self.hps_type == "pixie":
             arr_dict.update(preds_dict)
             arr_dict["global_orient"] = preds_dict["global_pose"]
-            arr_dict["betas"] = preds_dict["shape"]  #200
+            arr_dict["betas"] = preds_dict["shape"]    #200
             arr_dict["smpl_verts"] = preds_dict["vertices"]
             scale, tranX, tranY = preds_dict["cam"].split(1, dim=1)
             # 1.1435, 0.0128, 0.3520
 
         arr_dict["scale"] = scale.unsqueeze(1)
-        arr_dict["trans"] = (torch.cat([tranX, tranY, torch.zeros_like(tranX)], dim=1).unsqueeze(1).to(self.device).float())
+        arr_dict["trans"] = (
+            torch.cat([tranX, tranY, torch.zeros_like(tranX)],
+                      dim=1).unsqueeze(1).to(self.device).float()
+        )
 
         # data_dict info (key-shape):
         # scale, tranX, tranY - tensor.float
@@ -230,4 +258,4 @@ class TestDataset:
 
         # render optimized mesh (normal, T_normal, image [-1,1])
         self.render.load_meshes(verts, faces)
-        return self.render.get_image(type="depth")
+        return self.render.get_image(type="depth")

lib/dataset/body_model.py

@@ -21,7 +21,6 @@ import os
 
 
 class SMPLModel:
-
     def __init__(self, model_path, age):
         """
         SMPL model.
@@ -49,20 +48,16 @@ class SMPLModel:
 
         if age == "kid":
             v_template_smil = np.load(
-                os.path.join(os.path.dirname(model_path),
-                             "smpl/smpl_kid_template.npy"))
+                os.path.join(os.path.dirname(model_path), "smpl/smpl_kid_template.npy")
+            )
             v_template_smil -= np.mean(v_template_smil, axis=0)
-            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
-                                             axis=2)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template, axis=2)
             self.shapedirs = np.concatenate(
-                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
-                axis=2)
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), axis=2
+            )
             self.beta_shape[0] += 1
 
-        id_to_col = {
-            self.kintree_table[1, i]: i
-            for i in range(self.kintree_table.shape[1])
-        }
+        id_to_col = {self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])}
         self.parent = {
             i: id_to_col[self.kintree_table[0, i]]
             for i in range(1, self.kintree_table.shape[1])
@@ -121,33 +116,30 @@ class SMPLModel:
         pose_cube = self.pose.reshape((-1, 1, 3))
         # rotation matrix for each joint
         self.R = self.rodrigues(pose_cube)
-        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
-                                 (self.R.shape[0] - 1, 3, 3))
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), (self.R.shape[0] - 1, 3, 3))
         lrotmin = (self.R[1:] - I_cube).ravel()
         # how pose affect body shape in zero pose
         v_posed = v_shaped + self.posedirs.dot(lrotmin)
         # world transformation of each joint
         G = np.empty((self.kintree_table.shape[1], 4, 4))
-        G[0] = self.with_zeros(
-            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
         for i in range(1, self.kintree_table.shape[1]):
             G[i] = G[self.parent[i]].dot(
                 self.with_zeros(
-                    np.hstack([
-                        self.R[i],
-                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
-                            [3, 1])),
-                    ])))
+                    np.hstack(
+                        [
+                            self.R[i],
+                            ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1])),
+                        ]
+                    )
+                )
+            )
         # remove the transformation due to the rest pose
-        G = G - self.pack(
-            np.matmul(
-                G,
-                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        G = G - self.pack(np.matmul(G, np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
         # transformation of each vertex
         T = np.tensordot(self.weights, G, axes=[[1], [0]])
         rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
-        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
-                                                                    4])[:, :3]
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3]
         self.verts = v + self.trans.reshape([1, 3])
         self.G = G
 
@@ -171,19 +163,20 @@ 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])
-        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
-                                 [theta.shape[0], 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
         dot = np.matmul(A, B)
@@ -238,12 +231,7 @@ class SMPLModel:
 
 
 class TetraSMPLModel:
-
-    def __init__(self,
-                 model_path,
-                 model_addition_path,
-                 age="adult",
-                 v_template=None):
+    def __init__(self, model_path, model_addition_path, age="adult", v_template=None):
         """
         SMPL model.
 
@@ -276,10 +264,7 @@ class TetraSMPLModel:
         self.posedirs_added = params_added["posedirs_added"]
         self.tetrahedrons = params_added["tetrahedrons"]
 
-        id_to_col = {
-            self.kintree_table[1, i]: i
-            for i in range(self.kintree_table.shape[1])
-        }
+        id_to_col = {self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])}
         self.parent = {
             i: id_to_col[self.kintree_table[0, i]]
             for i in range(1, self.kintree_table.shape[1])
@@ -291,14 +276,13 @@ class TetraSMPLModel:
 
         if age == "kid":
             v_template_smil = np.load(
-                os.path.join(os.path.dirname(model_path),
-                             "smpl_kid_template.npy"))
+                os.path.join(os.path.dirname(model_path), "smpl_kid_template.npy")
+            )
             v_template_smil -= np.mean(v_template_smil, axis=0)
-            v_template_diff = np.expand_dims(v_template_smil - self.v_template,
-                                             axis=2)
+            v_template_diff = np.expand_dims(v_template_smil - self.v_template, axis=2)
             self.shapedirs = np.concatenate(
-                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff),
-                axis=2)
+                (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), axis=2
+            )
             self.beta_shape[0] += 1
 
         self.pose = np.zeros(self.pose_shape)
@@ -356,50 +340,42 @@ class TetraSMPLModel:
         """
         # how beta affect body shape
         v_shaped = self.shapedirs.dot(self.beta) + self.v_template
-        v_shaped_added = self.shapedirs_added.dot(
-            self.beta) + self.v_template_added
+        v_shaped_added = self.shapedirs_added.dot(self.beta) + self.v_template_added
         # joints location
         self.J = self.J_regressor.dot(v_shaped)
         pose_cube = self.pose.reshape((-1, 1, 3))
         # rotation matrix for each joint
         self.R = self.rodrigues(pose_cube)
-        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
-                                 (self.R.shape[0] - 1, 3, 3))
+        I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), (self.R.shape[0] - 1, 3, 3))
         lrotmin = (self.R[1:] - I_cube).ravel()
         # how pose affect body shape in zero pose
         v_posed = v_shaped + self.posedirs.dot(lrotmin)
         v_posed_added = v_shaped_added + self.posedirs_added.dot(lrotmin)
         # world transformation of each joint
         G = np.empty((self.kintree_table.shape[1], 4, 4))
-        G[0] = self.with_zeros(
-            np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+        G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
         for i in range(1, self.kintree_table.shape[1]):
             G[i] = G[self.parent[i]].dot(
                 self.with_zeros(
-                    np.hstack([
-                        self.R[i],
-                        ((self.J[i, :] - self.J[self.parent[i], :]).reshape(
-                            [3, 1])),
-                    ])))
+                    np.hstack(
+                        [
+                            self.R[i],
+                            ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1])),
+                        ]
+                    )
+                )
+            )
         # remove the transformation due to the rest pose
-        G = G - self.pack(
-            np.matmul(
-                G,
-                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
+        G = G - self.pack(np.matmul(G, np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])))
         self.G = G
         # transformation of each vertex
         T = np.tensordot(self.weights, G, axes=[[1], [0]])
         rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
-        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1,
-                                                                    4])[:, :3]
+        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3]
         self.verts = v + self.trans.reshape([1, 3])
         T_added = np.tensordot(self.weights_added, G, axes=[[1], [0]])
-        rest_shape_added_h = np.hstack(
-            (v_posed_added, np.ones([v_posed_added.shape[0], 1])))
-        v_added = np.matmul(T_added,
-                            rest_shape_added_h.reshape([-1, 4,
-                                                        1])).reshape([-1, 4
-                                                                      ])[:, :3]
+        rest_shape_added_h = np.hstack((v_posed_added, np.ones([v_posed_added.shape[0], 1])))
+        v_added = np.matmul(T_added, rest_shape_added_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3]
         self.verts_added = v_added + self.trans.reshape([1, 3])
 
     def rodrigues(self, r):
@@ -422,19 +398,20 @@ 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])
-        i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0),
-                                 [theta.shape[0], 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
         dot = np.matmul(A, B)

lib/dataset/mesh_util.py

@@ -14,32 +14,33 @@
 #
 # Contact: [email protected]
 
+import os
 import numpy as np
-import cv2
-import pymeshlab
 import torch
 import torchvision
 import trimesh
-import os
-from termcolor import colored
+import open3d as o3d
+import tinyobjloader
 import os.path as osp
 import _pickle as cPickle
+from termcolor import colored
 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
-import tinyobjloader
 
-from lib.common.imutils import uncrop
-from lib.common.render_utils import Pytorch3dRasterizer
 
+class Format:
+    end = '\033[0m'
+    start = '\033[4m'
 
-class SMPLX:
 
+class SMPLX:
     def __init__(self):
 
         self.current_dir = osp.join(osp.dirname(__file__), "../../data/smpl_related")
@@ -54,10 +55,14 @@ class SMPLX:
 
         self.smplx_eyeball_fid_path = osp.join(self.current_dir, "smpl_data/eyeball_fid.npy")
         self.smplx_fill_mouth_fid_path = osp.join(self.current_dir, "smpl_data/fill_mouth_fid.npy")
-        self.smplx_flame_vid_path = osp.join(self.current_dir, "smpl_data/FLAME_SMPLX_vertex_ids.npy")
+        self.smplx_flame_vid_path = osp.join(
+            self.current_dir, "smpl_data/FLAME_SMPLX_vertex_ids.npy"
+        )
         self.smplx_mano_vid_path = osp.join(self.current_dir, "smpl_data/MANO_SMPLX_vertex_ids.pkl")
         self.front_flame_path = osp.join(self.current_dir, "smpl_data/FLAME_face_mask_ids.npy")
-        self.smplx_vertex_lmkid_path = osp.join(self.current_dir, "smpl_data/smplx_vertex_lmkid.npy")
+        self.smplx_vertex_lmkid_path = osp.join(
+            self.current_dir, "smpl_data/smplx_vertex_lmkid.npy"
+        )
 
         self.smplx_faces = np.load(self.smplx_faces_path)
         self.smplx_verts = np.load(self.smplx_verts_path)
@@ -68,84 +73,51 @@ class SMPLX:
         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)]
 
         # hands
-        self.mano_vertex_mask = torch.zeros(self.smplx_verts.shape[0],).index_fill_(0, torch.tensor(self.smplx_mano_vid), 1.0)
+        self.mano_vertex_mask = torch.zeros(self.smplx_verts.shape[0], ).index_fill_(
+            0, torch.tensor(self.smplx_mano_vid), 1.0
+        )
         # face
-        self.front_flame_vertex_mask = torch.zeros(self.smplx_verts.shape[0],).index_fill_(
-            0, torch.tensor(self.smplx_front_flame_vid), 1.0)
-        self.eyeball_vertex_mask = torch.zeros(self.smplx_verts.shape[0],).index_fill_(
-            0, torch.tensor(self.smplx_faces[self.smplx_eyeball_fid_mask].flatten()), 1.0)
+        self.front_flame_vertex_mask = torch.zeros(self.smplx_verts.shape[0], ).index_fill_(
+            0, torch.tensor(self.smplx_front_flame_vid), 1.0
+        )
+        self.eyeball_vertex_mask = torch.zeros(self.smplx_verts.shape[0], ).index_fill_(
+            0, torch.tensor(self.smplx_faces[self.smplx_eyeball_fid_mask].flatten()), 1.0
+        )
 
         self.smplx_to_smpl = cPickle.load(open(self.smplx_to_smplx_path, "rb"))
 
         self.model_dir = osp.join(self.current_dir, "models")
         self.tedra_dir = osp.join(self.current_dir, "../tedra_data")
 
-        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,
-            ]) + 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 += 68
 
         self.smpl_joint_ids_45_pixie = (np.arange(22).tolist() + self.extra_joint_ids.tolist())
 
@@ -222,27 +194,6 @@ def load_fit_body(fitted_path, scale, smpl_type="smplx", smpl_gender="neutral",
     return smpl_mesh, smpl_joints
 
 
-def create_grid_points_from_xyz_bounds(bound, res):
-
-    min_x, max_x, min_y, max_y, min_z, max_z = bound
-    x = torch.linspace(min_x, max_x, res)
-    y = torch.linspace(min_y, max_y, res)
-    z = torch.linspace(min_z, max_z, res)
-    X, Y, Z = torch.meshgrid(x, y, z, indexing='ij')
-
-    return torch.stack([X, Y, Z], dim=-1)
-
-
-def create_grid_points_from_xy_bounds(bound, res):
-
-    min_x, max_x, min_y, max_y = bound
-    x = torch.linspace(min_x, max_x, res)
-    y = torch.linspace(min_y, max_y, res)
-    X, Y = torch.meshgrid(x, y, indexing='ij')
-
-    return torch.stack([X, Y], dim=-1)
-
-
 def apply_face_mask(mesh, face_mask):
 
     mesh.update_faces(face_mask)
@@ -277,7 +228,8 @@ def part_removal(full_mesh, part_mesh, thres, device, smpl_obj, region, clean=Tr
 
     part_extractor = PointFeat(
         torch.tensor(part_mesh.vertices).unsqueeze(0).to(device),
-        torch.tensor(part_mesh.faces).unsqueeze(0).to(device))
+        torch.tensor(part_mesh.faces).unsqueeze(0).to(device)
+    )
 
     (part_dist, _) = part_extractor.query(torch.tensor(full_mesh.vertices).unsqueeze(0).to(device))
 
@@ -286,12 +238,20 @@ def part_removal(full_mesh, part_mesh, thres, device, smpl_obj, region, clean=Tr
     if region == "hand":
         _, idx = smpl_tree.query(full_mesh.vertices, k=1)
         full_lmkid = SMPL_container.smplx_vertex_lmkid[idx]
-        remove_mask = torch.logical_and(remove_mask, torch.tensor(full_lmkid >= 20).type_as(remove_mask).unsqueeze(0))
+        remove_mask = torch.logical_and(
+            remove_mask,
+            torch.tensor(full_lmkid >= 20).type_as(remove_mask).unsqueeze(0)
+        )
 
     elif region == "face":
         _, idx = smpl_tree.query(full_mesh.vertices, k=5)
-        face_space_mask = torch.isin(torch.tensor(idx), torch.tensor(SMPL_container.smplx_front_flame_vid))
-        remove_mask = torch.logical_and(remove_mask, face_space_mask.any(dim=1).type_as(remove_mask).unsqueeze(0))
+        face_space_mask = torch.isin(
+            torch.tensor(idx), torch.tensor(SMPL_container.smplx_front_flame_vid)
+        )
+        remove_mask = torch.logical_and(
+            remove_mask,
+            face_space_mask.any(dim=1).type_as(remove_mask).unsqueeze(0)
+        )
 
     BNI_part_mask = ~(remove_mask).flatten()[full_mesh.faces].any(dim=1)
     full_mesh.update_faces(BNI_part_mask.detach().cpu())
@@ -303,109 +263,6 @@ def part_removal(full_mesh, part_mesh, thres, device, smpl_obj, region, clean=Tr
     return full_mesh
 
 
-def cross(triangles):
-    """
-    Returns the cross product of two edges from input triangles
-    Parameters
-    --------------
-    triangles: (n, 3, 3) float
-      Vertices of triangles
-    Returns
-    --------------
-    crosses : (n, 3) float
-      Cross product of two edge vectors
-    """
-    vectors = np.diff(triangles, axis=1)
-    crosses = np.cross(vectors[:, 0], vectors[:, 1])
-    return crosses
-
-
-def tri_area(triangles=None, crosses=None, sum=False):
-    """
-    Calculates the sum area of input triangles
-    Parameters
-    ----------
-    triangles : (n, 3, 3) float
-      Vertices of triangles
-    crosses : (n, 3) float or None
-      As a speedup don't re- compute cross products
-    sum : bool
-      Return summed area or individual triangle area
-    Returns
-    ----------
-    area : (n,) float or float
-      Individual or summed area depending on `sum` argument
-    """
-    if crosses is None:
-        crosses = cross(triangles)
-    area = (np.sum(crosses**2, axis=1)**.5) * .5
-    if sum:
-        return np.sum(area)
-    return area
-
-
-def sample_surface(triangles, count, area=None):
-    """
-    Sample the surface of a mesh, returning the specified
-    number of points
-    For individual triangle sampling uses this method:
-    http://mathworld.wolfram.com/TrianglePointPicking.html
-    Parameters
-    ---------
-    triangles : (n, 3, 3) float
-      Vertices of triangles
-    count : int
-      Number of points to return
-    Returns
-    ---------
-    samples : (count, 3) float
-      Points in space on the surface of mesh
-    face_index : (count,) int
-      Indices of faces for each sampled point
-    """
-
-    # len(mesh.faces) float, array of the areas
-    # of each face of the mesh
-    if area is None:
-        area = tri_area(triangles)
-
-    # total area (float)
-    area_sum = np.sum(area)
-    # cumulative area (len(mesh.faces))
-    area_cum = np.cumsum(area)
-    face_pick = np.random.random(count) * area_sum
-    face_index = np.searchsorted(area_cum, face_pick)
-
-    # pull triangles into the form of an origin + 2 vectors
-    tri_origins = triangles[:, 0]
-    tri_vectors = triangles[:, 1:].copy()
-    tri_vectors -= np.tile(tri_origins, (1, 2)).reshape((-1, 2, 3))
-
-    # pull the vectors for the faces we are going to sample from
-    tri_origins = tri_origins[face_index]
-    tri_vectors = tri_vectors[face_index]
-
-    # randomly generate two 0-1 scalar components to multiply edge vectors by
-    random_lengths = np.random.random((len(tri_vectors), 2, 1))
-
-    # points will be distributed on a quadrilateral if we use 2 0-1 samples
-    # if the two scalar components sum less than 1.0 the point will be
-    # inside the triangle, so we find vectors longer than 1.0 and
-    # transform them to be inside the triangle
-    random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0
-    random_lengths[random_test] -= 1.0
-    random_lengths = np.abs(random_lengths)
-
-    # multiply triangle edge vectors by the random lengths and sum
-    sample_vector = (tri_vectors * random_lengths).sum(axis=1)
-
-    # finally, offset by the origin to generate
-    # (n,3) points in space on the triangle
-    samples = torch.tensor(sample_vector + tri_origins).float()
-
-    return samples, face_index
-
-
 def obj_loader(path, with_uv=True):
     # Create reader.
     reader = tinyobjloader.ObjReader()
@@ -424,8 +281,8 @@ def obj_loader(path, with_uv=True):
     f_vt = tri[:, [2, 5, 8]]
 
     if with_uv:
-        face_uvs = vt[f_vt].mean(axis=1)  #[m, 2]
-        vert_uvs = np.zeros((v.shape[0], 2), dtype=np.float32)  #[n, 2]
+        face_uvs = vt[f_vt].mean(axis=1)    #[m, 2]
+        vert_uvs = np.zeros((v.shape[0], 2), dtype=np.float32)    #[n, 2]
         vert_uvs[f_v.reshape(-1)] = vt[f_vt.reshape(-1)]
 
         return v, f_v, vert_uvs, face_uvs
@@ -434,7 +291,6 @@ def obj_loader(path, with_uv=True):
 
 
 class HoppeMesh:
-
     def __init__(self, verts, faces, uvs=None, texture=None):
         """
         The HoppeSDF calculates signed distance towards a predefined oriented point cloud
@@ -459,34 +315,20 @@ class HoppeMesh:
         - points: [n, 3]
         - return: [n, 4] rgba
         """
-        triangles = self.verts[faces]  #[n, 3, 3]
-        barycentric = trimesh.triangles.points_to_barycentric(triangles, points)  #[n, 3]
-        vert_colors = self.vertex_colors[faces]  #[n, 3, 4]
+        triangles = self.verts[faces]    #[n, 3, 3]
+        barycentric = trimesh.triangles.points_to_barycentric(triangles, points)    #[n, 3]
+        vert_colors = self.vertex_colors[faces]    #[n, 3, 4]
         point_colors = torch.tensor((barycentric[:, :, None] * vert_colors).sum(axis=1)).float()
         return point_colors
 
     def triangles(self):
-        return self.verts[self.faces].numpy()  #[n, 3, 3]
+        return self.verts[self.faces].numpy()    #[n, 3, 3]
 
 
 def tensor2variable(tensor, device):
     return tensor.requires_grad_(True).to(device)
 
 
-class GMoF(torch.nn.Module):
-
-    def __init__(self, rho=1):
-        super(GMoF, self).__init__()
-        self.rho = rho
-
-    def extra_repr(self):
-        return "rho = {}".format(self.rho)
-
-    def forward(self, residual):
-        dist = torch.div(residual, residual + self.rho**2)
-        return self.rho**2 * dist
-
-
 def mesh_edge_loss(meshes, target_length: float = 0.0):
     """
     Computes mesh edge length regularization loss averaged across all meshes
@@ -508,10 +350,10 @@ def mesh_edge_loss(meshes, target_length: float = 0.0):
         return torch.tensor([0.0], dtype=torch.float32, device=meshes.device, requires_grad=True)
 
     N = len(meshes)
-    edges_packed = meshes.edges_packed()  # (sum(E_n), 3)
-    verts_packed = meshes.verts_packed()  # (sum(V_n), 3)
-    edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx()  # (sum(E_n), )
-    num_edges_per_mesh = meshes.num_edges_per_mesh()  # N
+    edges_packed = meshes.edges_packed()    # (sum(E_n), 3)
+    verts_packed = meshes.verts_packed()    # (sum(V_n), 3)
+    edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx()    # (sum(E_n), )
+    num_edges_per_mesh = meshes.num_edges_per_mesh()    # N
 
     # Determine the weight for each edge based on the number of edges in the
     # mesh it corresponds to.
@@ -531,99 +373,37 @@ def mesh_edge_loss(meshes, target_length: float = 0.0):
     return loss_all
 
 
-def remesh(obj, obj_path):
-
-    obj.export(obj_path)
-    ms = pymeshlab.MeshSet()
-    ms.load_new_mesh(obj_path)
-    # ms.meshing_decimation_quadric_edge_collapse(targetfacenum=100000)
-    ms.meshing_isotropic_explicit_remeshing(targetlen=pymeshlab.Percentage(0.5), adaptive=True)
-    ms.apply_coord_laplacian_smoothing()
-    ms.save_current_mesh(obj_path[:-4] + "_remesh.obj")
-    polished_mesh = trimesh.load_mesh(obj_path[:-4] + "_remesh.obj")
+def remesh_laplacian(mesh, obj_path):
 
-    return polished_mesh
-
-
-def poisson_remesh(obj_path):
-
-    ms = pymeshlab.MeshSet()
-    ms.load_new_mesh(obj_path)
-    ms.meshing_decimation_quadric_edge_collapse(targetfacenum=50000)
-    # ms.apply_coord_laplacian_smoothing()
-    ms.save_current_mesh(obj_path)
-    # ms.save_current_mesh(obj_path.replace(".obj", ".ply"))
-    polished_mesh = trimesh.load_mesh(obj_path)
+    mesh = mesh.simplify_quadratic_decimation(50000)
+    mesh = trimesh.smoothing.filter_humphrey(
+        mesh, alpha=0.1, beta=0.5, iterations=10, laplacian_operator=None
+    )
+    mesh.export(obj_path)
 
-    return polished_mesh
+    return mesh
 
 
 def poisson(mesh, obj_path, depth=10):
 
-    from pypoisson import poisson_reconstruction
-    faces, vertices = poisson_reconstruction(mesh.vertices, mesh.vertex_normals, depth=depth)
-
-    new_meshes = trimesh.Trimesh(vertices, faces)
-    new_mesh_lst = new_meshes.split(only_watertight=False)
-    comp_num = [new_mesh.vertices.shape[0] for new_mesh in new_mesh_lst]
-    final_mesh = new_mesh_lst[comp_num.index(max(comp_num))]
-    final_mesh.export(obj_path)
+    pcd_path = obj_path[:-4] + ".ply"
+    assert (mesh.vertex_normals.shape[1] == 3)
+    mesh.export(pcd_path)
+    pcl = o3d.io.read_point_cloud(pcd_path)
+    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Error) as cm:
+        mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+            pcl, depth=depth, n_threads=-1
+        )
+    print(colored(f"\n Poisson completion to {Format.start} {obj_path} {Format.end}", "yellow"))
 
-    final_mesh = poisson_remesh(obj_path)
+    # only keep the largest component
+    largest_mesh = keep_largest(trimesh.Trimesh(np.array(mesh.vertices), np.array(mesh.triangles)))
+    largest_mesh.export(obj_path)
 
-    return final_mesh
+    # mesh decimation for faster rendering
+    low_res_mesh = largest_mesh.simplify_quadratic_decimation(50000)
 
-
-def get_mask(tensor, dim):
-
-    mask = torch.abs(tensor).sum(dim=dim, keepdims=True) > 0.0
-    mask = mask.type_as(tensor)
-
-    return mask
-
-
-def blend_rgb_norm(norms, data):
-
-    # norms [N, 3, res, res]
-
-    masks = (norms.sum(dim=1) != norms[0, :, 0, 0].sum()).float().unsqueeze(1)
-    norm_mask = F.interpolate(
-        torch.cat([norms, masks], dim=1).detach().cpu(),
-        size=data["uncrop_param"]["box_shape"],
-        mode="bilinear",
-        align_corners=False).permute(0, 2, 3, 1).numpy()
-    final = data["img_raw"]
-
-    for idx in range(len(norms)):
-
-        norm_pred = (norm_mask[idx, :, :, :3] + 1.0) * 255.0 / 2.0
-        mask_pred = np.repeat(norm_mask[idx, :, :, 3:4], 3, axis=-1)
-
-        norm_ori = unwrap(norm_pred, data["uncrop_param"], idx)
-        mask_ori = unwrap(mask_pred, data["uncrop_param"], idx)
-
-        final = final * (1.0 - mask_ori) + norm_ori * mask_ori
-
-    return final.astype(np.uint8)
-
-
-def unwrap(image, uncrop_param, idx):
-
-    img_uncrop = uncrop(
-        image,
-        uncrop_param["center"][idx],
-        uncrop_param["scale"][idx],
-        uncrop_param["crop_shape"],
-    )
-
-    img_orig = cv2.warpAffine(
-        img_uncrop,
-        np.linalg.inv(uncrop_param["M"])[:2, :],
-        uncrop_param["ori_shape"][::-1],
-        flags=cv2.INTER_CUBIC,
-    )
-
-    return img_orig
+    return low_res_mesh
 
 
 # Losses to smooth / regularize the mesh shape
@@ -634,60 +414,7 @@ def update_mesh_shape_prior_losses(mesh, losses):
     # mesh normal consistency
     losses["nc"]["value"] = mesh_normal_consistency(mesh)
     # mesh laplacian smoothing
-    losses["laplacian"]["value"] = mesh_laplacian_smoothing(mesh, method="uniform")
-
-
-def rename(old_dict, old_name, new_name):
-    new_dict = {}
-    for key, value in zip(old_dict.keys(), old_dict.values()):
-        new_key = key if key != old_name else new_name
-        new_dict[new_key] = old_dict[key]
-    return new_dict
-
-
-def load_checkpoint(model, cfg):
-
-    model_dict = model.state_dict()
-    main_dict = {}
-    normal_dict = {}
-
-    device = torch.device(f"cuda:{cfg['test_gpus'][0]}")
-
-    if os.path.exists(cfg.resume_path) and cfg.resume_path.endswith("ckpt"):
-        main_dict = torch.load(cfg.resume_path, map_location=device)["state_dict"]
-
-        main_dict = {
-            k: v for k, v in main_dict.items() if k in model_dict and v.shape == model_dict[k].shape and
-            ("reconEngine" not in k) and ("normal_filter" not in k) and ("voxelization" not in k)
-        }
-        print(colored(f"Resume MLP weights from {cfg.resume_path}", "green"))
-
-    if os.path.exists(cfg.normal_path) and cfg.normal_path.endswith("ckpt"):
-        normal_dict = torch.load(cfg.normal_path, map_location=device)["state_dict"]
-
-        for key in normal_dict.keys():
-            normal_dict = rename(normal_dict, key, key.replace("netG", "netG.normal_filter"))
-
-        normal_dict = {k: v for k, v in normal_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
-        print(colored(f"Resume normal model from {cfg.normal_path}", "green"))
-
-    model_dict.update(main_dict)
-    model_dict.update(normal_dict)
-    model.load_state_dict(model_dict)
-
-    model.netG = model.netG.to(device)
-    model.reconEngine = model.reconEngine.to(device)
-
-    model.netG.training = False
-    model.netG.eval()
-
-    del main_dict
-    del normal_dict
-    del model_dict
-
-    torch.cuda.empty_cache()
-
-    return model
+    losses["lapla"]["value"] = mesh_laplacian_smoothing(mesh, method="uniform")
 
 
 def read_smpl_constants(folder):
@@ -706,8 +433,10 @@ def read_smpl_constants(folder):
     smpl_vertex_code = np.float32(np.copy(smpl_vtx_std))
     """Load smpl faces & tetrahedrons"""
     smpl_faces = np.loadtxt(os.path.join(folder, "faces.txt"), dtype=np.int32) - 1
-    smpl_face_code = (smpl_vertex_code[smpl_faces[:, 0]] + smpl_vertex_code[smpl_faces[:, 1]] +
-                      smpl_vertex_code[smpl_faces[:, 2]]) / 3.0
+    smpl_face_code = (
+        smpl_vertex_code[smpl_faces[:, 0]] + smpl_vertex_code[smpl_faces[:, 1]] +
+        smpl_vertex_code[smpl_faces[:, 2]]
+    ) / 3.0
     smpl_tetras = (np.loadtxt(os.path.join(folder, "tetrahedrons.txt"), dtype=np.int32) - 1)
 
     return_dict = {
@@ -720,19 +449,6 @@ def read_smpl_constants(folder):
     return return_dict
 
 
-def feat_select(feat, select):
-
-    # feat [B, featx2, N]
-    # select [B, 1, N]
-    # return [B, feat, N]
-
-    dim = feat.shape[1] // 2
-    idx = torch.tile((1 - select), (1, dim, 1)) * dim + torch.arange(0, dim).unsqueeze(0).unsqueeze(2).type_as(select)
-    feat_select = torch.gather(feat, 1, idx.long())
-
-    return feat_select
-
-
 def get_visibility(xy, z, faces, img_res=2**12, blur_radius=0.0, faces_per_pixel=1):
     """get the visibility of vertices
 
@@ -771,7 +487,9 @@ def get_visibility(xy, z, faces, img_res=2**12, blur_radius=0.0, faces_per_pixel
 
     for idx in range(N_body):
         Num_faces = len(faces[idx])
-        vis_vertices_id = torch.unique(faces[idx][torch.unique(pix_to_face[idx][pix_to_face[idx] != -1]) - Num_faces * idx, :])
+        vis_vertices_id = torch.unique(
+            faces[idx][torch.unique(pix_to_face[idx][pix_to_face[idx] != -1]) - Num_faces * idx, :]
+        )
         vis_mask[idx, vis_vertices_id] = 1.0
 
     # print("------------------------\n")
@@ -825,7 +543,7 @@ def orthogonal(points, calibrations, transforms=None):
     """
     rot = calibrations[:, :3, :3]
     trans = calibrations[:, :3, 3:4]
-    pts = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    pts = torch.baddbmm(trans, rot, points)    # [B, 3, N]
     if transforms is not None:
         scale = transforms[:2, :2]
         shift = transforms[:2, 2:3]
@@ -925,37 +643,14 @@ def compute_normal_batch(vertices, faces):
     return vert_norm
 
 
-def calculate_mIoU(outputs, labels):
-
-    SMOOTH = 1e-6
-
-    outputs = outputs.int()
-    labels = labels.int()
-
-    intersection = ((outputs & labels).float().sum())  # Will be zero if Truth=0 or Prediction=0
-    union = (outputs | labels).float().sum()  # Will be zzero if both are 0
-
-    iou = (intersection + SMOOTH) / (union + SMOOTH)  # We smooth our devision to avoid 0/0
-
-    thresholded = (torch.clamp(20 * (iou - 0.5), 0, 10).ceil() / 10)  # This is equal to comparing with thresolds
-
-    return (thresholded.mean().detach().cpu().numpy()
-           )  # Or thresholded.mean() if you are interested in average across the batch
-
-
-def add_alpha(colors, alpha=0.7):
-
-    colors_pad = np.pad(colors, ((0, 0), (0, 1)), mode="constant", constant_values=alpha)
-
-    return colors_pad
-
-
 def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type="smpl"):
 
     font_path = os.path.join(os.path.dirname(__file__), "tbfo.ttf")
     font = ImageFont.truetype(font_path, 30)
     grid_img = torchvision.utils.make_grid(torch.cat(per_loop_lst, dim=0), nrow=nrow, padding=0)
-    grid_img = Image.fromarray(((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 * 255.0).astype(np.uint8))
+    grid_img = Image.fromarray(
+        ((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 * 255.0).astype(np.uint8)
+    )
 
     if False:
         # add text
@@ -965,16 +660,20 @@ 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([
+            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(
@@ -996,12 +695,9 @@ def clean_mesh(verts, faces):
     device = verts.device
 
     mesh_lst = trimesh.Trimesh(verts.detach().cpu().numpy(), faces.detach().cpu().numpy())
-    mesh_lst = mesh_lst.split(only_watertight=False)
-    comp_num = [mesh.vertices.shape[0] for mesh in mesh_lst]
-
-    mesh_clean = mesh_lst[comp_num.index(max(comp_num))]
-    final_verts = torch.as_tensor(mesh_clean.vertices).float().to(device)
-    final_faces = torch.as_tensor(mesh_clean.faces).long().to(device)
+    largest_mesh = keep_largest(mesh_lst)
+    final_verts = torch.as_tensor(largest_mesh.vertices).float().to(device)
+    final_faces = torch.as_tensor(largest_mesh.faces).long().to(device)
 
     return final_verts, final_faces
 

lib/net/BasePIFuNet.py

@@ -21,11 +21,10 @@ from .geometry import index, orthogonal, perspective
 
 
 class BasePIFuNet(pl.LightningModule):
-
     def __init__(
-            self,
-            projection_mode="orthogonal",
-            error_term=nn.MSELoss(),
+        self,
+        projection_mode="orthogonal",
+        error_term=nn.MSELoss(),
     ):
         """
         :param projection_mode:

lib/net/Discriminator.py

@@ -9,17 +9,18 @@ from lib.torch_utils.ops.native_ops import FusedLeakyReLU, fused_leaky_relu, upf
 
 
 class DiscriminatorHead(nn.Module):
-
     def __init__(self, in_channel, disc_stddev=False):
         super().__init__()
 
         self.disc_stddev = disc_stddev
         stddev_dim = 1 if disc_stddev else 0
 
-        self.conv_stddev = ConvLayer2d(in_channel=in_channel + stddev_dim,
-                                       out_channel=in_channel,
-                                       kernel_size=3,
-                                       activate=True)
+        self.conv_stddev = ConvLayer2d(
+            in_channel=in_channel + stddev_dim,
+            out_channel=in_channel,
+            kernel_size=3,
+            activate=True
+        )
 
         self.final_linear = nn.Sequential(
             nn.Flatten(),
@@ -32,8 +33,8 @@ class DiscriminatorHead(nn.Module):
         inv_perm = torch.argsort(perm)
 
         batch, channel, height, width = x.shape
-        x = x[
-            perm]  # shuffle inputs so that all views in a single trajectory don't get put together
+        x = x[perm
+             ]    # shuffle inputs so that all views in a single trajectory don't get put together
 
         group = min(batch, stddev_group)
         stddev = x.view(group, -1, stddev_feat, channel // stddev_feat, height, width)
@@ -41,7 +42,7 @@ class DiscriminatorHead(nn.Module):
         stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
         stddev = stddev.repeat(group, 1, height, width)
 
-        stddev = stddev[inv_perm]  # reorder inputs
+        stddev = stddev[inv_perm]    # reorder inputs
         x = x[inv_perm]
 
         out = torch.cat([x, stddev], 1)
@@ -56,7 +57,6 @@ class DiscriminatorHead(nn.Module):
 
 
 class ConvDecoder(nn.Module):
-
     def __init__(self, in_channel, out_channel, in_res, out_res):
         super().__init__()
 
@@ -68,20 +68,22 @@ class ConvDecoder(nn.Module):
         for i in range(log_size_in, log_size_out):
             out_ch = in_ch // 2
             self.layers.append(
-                ConvLayer2d(in_channel=in_ch,
-                            out_channel=out_ch,
-                            kernel_size=3,
-                            upsample=True,
-                            bias=True,
-                            activate=True))
+                ConvLayer2d(
+                    in_channel=in_ch,
+                    out_channel=out_ch,
+                    kernel_size=3,
+                    upsample=True,
+                    bias=True,
+                    activate=True
+                )
+            )
             in_ch = out_ch
 
         self.layers.append(
-            ConvLayer2d(in_channel=in_ch,
-                        out_channel=out_channel,
-                        kernel_size=3,
-                        bias=True,
-                        activate=False))
+            ConvLayer2d(
+                in_channel=in_ch, out_channel=out_channel, kernel_size=3, bias=True, activate=False
+            )
+        )
         self.layers = nn.Sequential(*self.layers)
 
     def forward(self, x):
@@ -89,7 +91,6 @@ class ConvDecoder(nn.Module):
 
 
 class StyleDiscriminator(nn.Module):
-
     def __init__(self, in_channel, in_res, ch_mul=64, ch_max=512, **kwargs):
         super().__init__()
 
@@ -104,7 +105,8 @@ class StyleDiscriminator(nn.Module):
         for i in range(log_size_in, log_size_out, -1):
             out_channels = int(min(in_channels * 2, ch_max))
             self.layers.append(
-                ConvResBlock2d(in_channel=in_channels, out_channel=out_channels, downsample=True))
+                ConvResBlock2d(in_channel=in_channels, out_channel=out_channels, downsample=True)
+            )
             in_channels = out_channels
         self.layers = nn.Sequential(*self.layers)
 
@@ -147,7 +149,6 @@ class Blur(nn.Module):
         Upsample factor.
 
     """
-
     def __init__(self, kernel, pad, upsample_factor=1):
         super().__init__()
 
@@ -177,7 +178,6 @@ class Upsample(nn.Module):
         Upsampling factor.
 
     """
-
     def __init__(self, kernel=[1, 3, 3, 1], factor=2):
         super().__init__()
 
@@ -208,7 +208,6 @@ class Downsample(nn.Module):
         Downsampling factor.
 
     """
-
     def __init__(self, kernel=[1, 3, 3, 1], factor=2):
         super().__init__()
 
@@ -250,7 +249,6 @@ class EqualLinear(nn.Module):
         Apply leakyReLU activation.
 
     """
-
     def __init__(self, in_channel, out_channel, bias=True, bias_init=0, lr_mul=1, activate=False):
         super().__init__()
 
@@ -300,7 +298,6 @@ class EqualConv2d(nn.Module):
         Use bias term.
 
     """
-
     def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
         super().__init__()
 
@@ -316,16 +313,20 @@ class EqualConv2d(nn.Module):
             self.bias = None
 
     def forward(self, input):
-        out = F.conv2d(input,
-                       self.weight * self.scale,
-                       bias=self.bias,
-                       stride=self.stride,
-                       padding=self.padding)
+        out = F.conv2d(
+            input,
+            self.weight * self.scale,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding
+        )
         return out
 
     def __repr__(self):
-        return (f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
-                f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})")
+        return (
+            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
+            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
+        )
 
 
 class EqualConvTranspose2d(nn.Module):
@@ -353,15 +354,16 @@ class EqualConvTranspose2d(nn.Module):
         Use bias term.
 
     """
-
-    def __init__(self,
-                 in_channel,
-                 out_channel,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 output_padding=0,
-                 bias=True):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        stride=1,
+        padding=0,
+        output_padding=0,
+        bias=True
+    ):
         super().__init__()
 
         self.weight = nn.Parameter(torch.randn(in_channel, out_channel, kernel_size, kernel_size))
@@ -388,12 +390,13 @@ class EqualConvTranspose2d(nn.Module):
         return out
 
     def __repr__(self):
-        return (f'{self.__class__.__name__}({self.weight.shape[0]}, {self.weight.shape[1]},'
-                f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})')
+        return (
+            f'{self.__class__.__name__}({self.weight.shape[0]}, {self.weight.shape[1]},'
+            f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
+        )
 
 
 class ConvLayer2d(nn.Sequential):
-
     def __init__(
         self,
         in_channel,
@@ -415,12 +418,15 @@ class ConvLayer2d(nn.Sequential):
             pad1 = p // 2 + 1
 
             layers.append(
-                EqualConvTranspose2d(in_channel,
-                                     out_channel,
-                                     kernel_size,
-                                     padding=0,
-                                     stride=2,
-                                     bias=bias and not activate))
+                EqualConvTranspose2d(
+                    in_channel,
+                    out_channel,
+                    kernel_size,
+                    padding=0,
+                    stride=2,
+                    bias=bias and not activate
+                )
+            )
             layers.append(Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor))
 
         if downsample:
@@ -431,23 +437,29 @@ class ConvLayer2d(nn.Sequential):
 
             layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
             layers.append(
-                EqualConv2d(in_channel,
-                            out_channel,
-                            kernel_size,
-                            padding=0,
-                            stride=2,
-                            bias=bias and not activate))
+                EqualConv2d(
+                    in_channel,
+                    out_channel,
+                    kernel_size,
+                    padding=0,
+                    stride=2,
+                    bias=bias and not activate
+                )
+            )
 
         if (not downsample) and (not upsample):
             padding = kernel_size // 2
 
             layers.append(
-                EqualConv2d(in_channel,
-                            out_channel,
-                            kernel_size,
-                            padding=padding,
-                            stride=1,
-                            bias=bias and not activate))
+                EqualConv2d(
+                    in_channel,
+                    out_channel,
+                    kernel_size,
+                    padding=padding,
+                    stride=1,
+                    bias=bias and not activate
+                )
+            )
 
         if activate:
             layers.append(FusedLeakyReLU(out_channel, bias=bias))
@@ -472,7 +484,6 @@ class ConvResBlock2d(nn.Module):
         Apply downsampling via strided convolution in the second conv.
 
     """
-
     def __init__(self, in_channel, out_channel, upsample=False, downsample=False):
         super().__init__()
 

lib/net/FBNet.py

@@ -51,17 +51,17 @@ def get_norm_layer(norm_type="instance"):
 
 
 def define_G(
-        input_nc,
-        output_nc,
-        ngf,
-        netG,
-        n_downsample_global=3,
-        n_blocks_global=9,
-        n_local_enhancers=1,
-        n_blocks_local=3,
-        norm="instance",
-        gpu_ids=[],
-        last_op=nn.Tanh(),
+    input_nc,
+    output_nc,
+    ngf,
+    netG,
+    n_downsample_global=3,
+    n_blocks_global=9,
+    n_local_enhancers=1,
+    n_blocks_local=3,
+    norm="instance",
+    gpu_ids=[],
+    last_op=nn.Tanh(),
 ):
     norm_layer = get_norm_layer(norm_type=norm)
     if netG == "global":
@@ -97,17 +97,20 @@ def define_G(
     return netG
 
 
-def define_D(input_nc,
-             ndf,
-             n_layers_D,
-             norm='instance',
-             use_sigmoid=False,
-             num_D=1,
-             getIntermFeat=False,
-             gpu_ids=[]):
+def define_D(
+    input_nc,
+    ndf,
+    n_layers_D,
+    norm='instance',
+    use_sigmoid=False,
+    num_D=1,
+    getIntermFeat=False,
+    gpu_ids=[]
+):
     norm_layer = get_norm_layer(norm_type=norm)
-    netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D,
-                                   getIntermFeat)
+    netD = MultiscaleDiscriminator(
+        input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat
+    )
     if len(gpu_ids) > 0:
         assert (torch.cuda.is_available())
         netD.cuda(gpu_ids[0])
@@ -129,7 +132,6 @@ def print_network(net):
 # Generator
 ##############################################################################
 class LocalEnhancer(pl.LightningModule):
-
     def __init__(
         self,
         input_nc,
@@ -155,8 +157,9 @@ class LocalEnhancer(pl.LightningModule):
             n_blocks_global,
             norm_layer,
         ).model
-        model_global = [model_global[i] for i in range(len(model_global) - 3)
-                       ]  # get rid of final convolution layers
+        model_global = [
+            model_global[i] for i in range(len(model_global) - 3)
+        ]    # get rid of final convolution layers
         self.model = nn.Sequential(*model_global)
 
         ###### local enhancer layers #####
@@ -224,17 +227,16 @@ class LocalEnhancer(pl.LightningModule):
 
 
 class GlobalGenerator(pl.LightningModule):
-
     def __init__(
-            self,
-            input_nc,
-            output_nc,
-            ngf=64,
-            n_downsampling=3,
-            n_blocks=9,
-            norm_layer=nn.BatchNorm2d,
-            padding_type="reflect",
-            last_op=nn.Tanh(),
+        self,
+        input_nc,
+        output_nc,
+        ngf=64,
+        n_downsampling=3,
+        n_blocks=9,
+        norm_layer=nn.BatchNorm2d,
+        padding_type="reflect",
+        last_op=nn.Tanh(),
     ):
         assert n_blocks >= 0
         super(GlobalGenerator, self).__init__()
@@ -296,42 +298,49 @@ class GlobalGenerator(pl.LightningModule):
 
 # Defines the PatchGAN discriminator with the specified arguments.
 class NLayerDiscriminator(nn.Module):
-
-    def __init__(self,
-                 input_nc,
-                 ndf=64,
-                 n_layers=3,
-                 norm_layer=nn.BatchNorm2d,
-                 use_sigmoid=False,
-                 getIntermFeat=False):
+    def __init__(
+        self,
+        input_nc,
+        ndf=64,
+        n_layers=3,
+        norm_layer=nn.BatchNorm2d,
+        use_sigmoid=False,
+        getIntermFeat=False
+    ):
         super(NLayerDiscriminator, self).__init__()
         self.getIntermFeat = getIntermFeat
         self.n_layers = n_layers
 
         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)]]
 
@@ -359,27 +368,30 @@ class NLayerDiscriminator(nn.Module):
 
 
 class MultiscaleDiscriminator(pl.LightningModule):
-
-    def __init__(self,
-                 input_nc,
-                 ndf=64,
-                 n_layers=3,
-                 norm_layer=nn.BatchNorm2d,
-                 use_sigmoid=False,
-                 num_D=3,
-                 getIntermFeat=False):
+    def __init__(
+        self,
+        input_nc,
+        ndf=64,
+        n_layers=3,
+        norm_layer=nn.BatchNorm2d,
+        use_sigmoid=False,
+        num_D=3,
+        getIntermFeat=False
+    ):
         super(MultiscaleDiscriminator, self).__init__()
         self.num_D = num_D
         self.n_layers = n_layers
         self.getIntermFeat = getIntermFeat
 
         for i in range(num_D):
-            netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid,
-                                       getIntermFeat)
+            netD = NLayerDiscriminator(
+                input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat
+            )
             if getIntermFeat:
                 for j in range(n_layers + 2):
-                    setattr(self, 'scale' + str(i) + '_layer' + str(j),
-                            getattr(netD, 'model' + str(j)))
+                    setattr(
+                        self, 'scale' + str(i) + '_layer' + str(j), getattr(netD, 'model' + str(j))
+                    )
             else:
                 setattr(self, 'layer' + str(i), netD.model)
 
@@ -414,11 +426,11 @@ class MultiscaleDiscriminator(pl.LightningModule):
 
 # Define a resnet block
 class ResnetBlock(pl.LightningModule):
-
     def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
         super(ResnetBlock, self).__init__()
-        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation,
-                                                use_dropout)
+        self.conv_block = self.build_conv_block(
+            dim, padding_type, norm_layer, activation, use_dropout
+        )
 
     def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
         conv_block = []
@@ -459,7 +471,6 @@ class ResnetBlock(pl.LightningModule):
 
 
 class Encoder(pl.LightningModule):
-
     def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d):
         super(Encoder, self).__init__()
         self.output_nc = output_nc
@@ -510,18 +521,17 @@ class Encoder(pl.LightningModule):
         inst_list = np.unique(inst.cpu().numpy().astype(int))
         for i in inst_list:
             for b in range(input.size()[0]):
-                indices = (inst[b:b + 1] == int(i)).nonzero()  # n x 4
+                indices = (inst[b:b + 1] == int(i)).nonzero()    # n x 4
                 for j in range(self.output_nc):
                     output_ins = outputs[indices[:, 0] + b, indices[:, 1] + j, indices[:, 2],
-                                         indices[:, 3],]
+                                         indices[:, 3], ]
                     mean_feat = torch.mean(output_ins).expand_as(output_ins)
                     outputs_mean[indices[:, 0] + b, indices[:, 1] + j, indices[:, 2],
-                                 indices[:, 3],] = mean_feat
+                                 indices[:, 3], ] = mean_feat
         return outputs_mean
 
 
 class Vgg19(nn.Module):
-
     def __init__(self, requires_grad=False):
         super(Vgg19, self).__init__()
         vgg_pretrained_features = models.vgg19(weights=models.VGG19_Weights.DEFAULT).features
@@ -555,7 +565,6 @@ class Vgg19(nn.Module):
 
 
 class VGG19FeatLayer(nn.Module):
-
     def __init__(self):
         super(VGG19FeatLayer, self).__init__()
         self.vgg19 = models.vgg19(weights=models.VGG19_Weights.DEFAULT).features.eval()
@@ -593,7 +602,6 @@ class VGG19FeatLayer(nn.Module):
 
 
 class VGGLoss(pl.LightningModule):
-
     def __init__(self):
         super(VGGLoss, self).__init__()
         self.vgg = Vgg19().eval()
@@ -609,11 +617,7 @@ class VGGLoss(pl.LightningModule):
 
 
 class GANLoss(pl.LightningModule):
-
-    def __init__(self,
-                 use_lsgan=True,
-                 target_real_label=1.0,
-                 target_fake_label=0.0):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
         super(GANLoss, self).__init__()
         self.real_label = target_real_label
         self.fake_label = target_fake_label
@@ -628,16 +632,18 @@ 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
@@ -659,7 +665,6 @@ class GANLoss(pl.LightningModule):
 
 
 class IDMRFLoss(pl.LightningModule):
-
     def __init__(self, featlayer=VGG19FeatLayer):
         super(IDMRFLoss, self).__init__()
         self.featlayer = featlayer()
@@ -678,7 +683,8 @@ class IDMRFLoss(pl.LightningModule):
         patch_size = 1
         patch_stride = 1
         patches_as_depth_vectors = featmaps.unfold(2, patch_size, patch_stride).unfold(
-            3, patch_size, patch_stride)
+            3, patch_size, patch_stride
+        )
         self.patches_OIHW = patches_as_depth_vectors.permute(0, 2, 3, 1, 4, 5)
         dims = self.patches_OIHW.size()
         self.patches_OIHW = self.patches_OIHW.view(-1, dims[3], dims[4], dims[5])
@@ -743,7 +749,8 @@ class IDMRFLoss(pl.LightningModule):
             self.mrf_loss(gen_vgg_feats[layer], tar_vgg_feats[layer])
             for layer in self.feat_content_layers
         ]
-        self.content_loss = functools.reduce(lambda x, y: x + y,
-                                             content_loss_list) * self.lambda_content
+        self.content_loss = functools.reduce(
+            lambda x, y: x + y, content_loss_list
+        ) * self.lambda_content
 
         return self.style_loss + self.content_loss

lib/net/GANLoss.py

@@ -32,13 +32,12 @@ def logistic_loss(fake_pred, real_pred, mode):
 
 
 def r1_loss(real_pred, real_img):
-    (grad_real,) = autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True)
+    (grad_real, ) = autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True)
     grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()
     return grad_penalty
 
 
 class GANLoss(nn.Module):
-
     def __init__(
         self,
         opt,
@@ -64,7 +63,7 @@ class GANLoss(nn.Module):
         logits_fake = self.discriminator(disc_in_fake)
 
         disc_loss = self.disc_loss(fake_pred=logits_fake, real_pred=logits_real, mode='d')
-        
+
         log = {
             "disc_loss": disc_loss.detach(),
             "logits_real": logits_real.mean().detach(),

lib/net/IFGeoNet.py

@@ -8,20 +8,17 @@ from lib.dataset.mesh_util import read_smpl_constants, SMPLX
 
 
 class SelfAttention(torch.nn.Module):
-
     def __init__(self, in_channels, out_channels):
         super().__init__()
-        self.conv = nn.Conv3d(in_channels,
-                                    out_channels,
-                                    3,
-                                    padding=1,
-                                    padding_mode='replicate')
-        self.attention = nn.Conv3d(in_channels,
-                                         out_channels,
-                                         kernel_size=3,
-                                         padding=1,
-                                         padding_mode='replicate',
-                                         bias=False)
+        self.conv = nn.Conv3d(in_channels, out_channels, 3, padding=1, padding_mode='replicate')
+        self.attention = nn.Conv3d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            padding_mode='replicate',
+            bias=False
+        )
         with torch.no_grad():
             self.attention.weight.copy_(torch.zeros_like(self.attention.weight))
 
@@ -32,38 +29,45 @@ class SelfAttention(torch.nn.Module):
 
 
 class IFGeoNet(nn.Module):
-
     def __init__(self, cfg, hidden_dim=256):
         super(IFGeoNet, self).__init__()
 
-        self.conv_in_partial = nn.Conv3d(1, 16, 3, padding=1,
-                                         padding_mode='replicate')  # out: 256 ->m.p. 128
+        self.conv_in_partial = nn.Conv3d(
+            1, 16, 3, padding=1, padding_mode='replicate'
+        )    # out: 256 ->m.p. 128
 
-        self.conv_in_smpl = nn.Conv3d(1, 4, 3, padding=1,
-                                      padding_mode='replicate')  # out: 256 ->m.p. 128
+        self.conv_in_smpl = nn.Conv3d(
+            1, 4, 3, padding=1, padding_mode='replicate'
+        )    # out: 256 ->m.p. 128
 
         self.SA = SelfAttention(4, 4)
-        self.conv_0_fusion = nn.Conv3d(16 + 4, 32, 3, padding=1,
-                                       padding_mode='replicate')  # out: 128
-        self.conv_0_1_fusion = nn.Conv3d(32, 32, 3, padding=1,
-                                         padding_mode='replicate')  # out: 128 ->m.p. 64
-
-        self.conv_0 = nn.Conv3d(32, 32, 3, padding=1, padding_mode='replicate')  # out: 128
-        self.conv_0_1 = nn.Conv3d(32, 32, 3, padding=1,
-                                  padding_mode='replicate')  # out: 128 ->m.p. 64
-
-        self.conv_1 = nn.Conv3d(32, 64, 3, padding=1, padding_mode='replicate')  # out: 64
-        self.conv_1_1 = nn.Conv3d(64, 64, 3, padding=1,
-                                  padding_mode='replicate')  # out: 64 -> mp 32
-
-        self.conv_2 = nn.Conv3d(64, 128, 3, padding=1, padding_mode='replicate')  # out: 32
-        self.conv_2_1 = nn.Conv3d(128, 128, 3, padding=1,
-                                  padding_mode='replicate')  # out: 32 -> mp 16
-        self.conv_3 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 16
-        self.conv_3_1 = nn.Conv3d(128, 128, 3, padding=1,
-                                  padding_mode='replicate')  # out: 16 -> mp 8
-        self.conv_4 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 8
-        self.conv_4_1 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 8
+        self.conv_0_fusion = nn.Conv3d(
+            16 + 4, 32, 3, padding=1, padding_mode='replicate'
+        )    # out: 128
+        self.conv_0_1_fusion = nn.Conv3d(
+            32, 32, 3, padding=1, padding_mode='replicate'
+        )    # out: 128 ->m.p. 64
+
+        self.conv_0 = nn.Conv3d(32, 32, 3, padding=1, padding_mode='replicate')    # out: 128
+        self.conv_0_1 = nn.Conv3d(
+            32, 32, 3, padding=1, padding_mode='replicate'
+        )    # out: 128 ->m.p. 64
+
+        self.conv_1 = nn.Conv3d(32, 64, 3, padding=1, padding_mode='replicate')    # out: 64
+        self.conv_1_1 = nn.Conv3d(
+            64, 64, 3, padding=1, padding_mode='replicate'
+        )    # out: 64 -> mp 32
+
+        self.conv_2 = nn.Conv3d(64, 128, 3, padding=1, padding_mode='replicate')    # out: 32
+        self.conv_2_1 = nn.Conv3d(
+            128, 128, 3, padding=1, padding_mode='replicate'
+        )    # out: 32 -> mp 16
+        self.conv_3 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 16
+        self.conv_3_1 = nn.Conv3d(
+            128, 128, 3, padding=1, padding_mode='replicate'
+        )    # out: 16 -> mp 8
+        self.conv_4 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 8
+        self.conv_4_1 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 8
 
         feature_size = (1 + 32 + 32 + 64 + 128 + 128 + 128) + 3
         self.fc_0 = nn.Conv1d(feature_size, hidden_dim * 2, 1)
@@ -97,21 +101,21 @@ class IFGeoNet(nn.Module):
             smooth_kernel_size=7,
             batch_size=cfg.batch_size,
         )
-        
+
         self.l1_loss = nn.SmoothL1Loss()
 
     def forward(self, batch):
-        
+
         if "body_voxels" in batch.keys():
             x_smpl = batch["body_voxels"]
         else:
             with torch.no_grad():
                 self.voxelization.update_param(batch["voxel_faces"])
-                x_smpl = self.voxelization(batch["voxel_verts"])[:, 0]  #[B, 128, 128, 128]
-        
+                x_smpl = self.voxelization(batch["voxel_verts"])[:, 0]    #[B, 128, 128, 128]
+
         p = orthogonal(batch["samples_geo"].permute(0, 2, 1),
-                       batch["calib"]).permute(0, 2, 1)  #[2, 60000, 3]
-        x = batch["depth_voxels"]  #[B, 128, 128, 128]
+                       batch["calib"]).permute(0, 2, 1)    #[2, 60000, 3]
+        x = batch["depth_voxels"]    #[B, 128, 128, 128]
 
         x = x.unsqueeze(1)
         x_smpl = x_smpl.unsqueeze(1)
@@ -119,63 +123,67 @@ class IFGeoNet(nn.Module):
         p = p.unsqueeze(1).unsqueeze(1)
 
         # partial inputs feature extraction
-        feature_0_partial = F.grid_sample(x, p, padding_mode='border', align_corners = True)
+        feature_0_partial = F.grid_sample(x, p, padding_mode='border', align_corners=True)
         net_partial = self.actvn(self.conv_in_partial(x))
         net_partial = self.partial_conv_in_bn(net_partial)
-        net_partial = self.maxpool(net_partial)  # out 64
+        net_partial = self.maxpool(net_partial)    # out 64
 
         # smpl inputs feature extraction
         # feature_0_smpl = F.grid_sample(x_smpl, p, padding_mode='border', align_corners = True)
         net_smpl = self.actvn(self.conv_in_smpl(x_smpl))
         net_smpl = self.smpl_conv_in_bn(net_smpl)
-        net_smpl = self.maxpool(net_smpl)  # out 64
+        net_smpl = self.maxpool(net_smpl)    # out 64
         net_smpl = self.SA(net_smpl)
-        
+
         # Feature fusion
         net = self.actvn(self.conv_0_fusion(torch.concat([net_partial, net_smpl], dim=1)))
         net = self.actvn(self.conv_0_1_fusion(net))
         net = self.conv0_1_bn_fusion(net)
-        feature_1_fused = F.grid_sample(net, p, padding_mode='border', align_corners = True)
+        feature_1_fused = F.grid_sample(net, p, padding_mode='border', align_corners=True)
         # net = self.maxpool(net)  # out 64
 
         net = self.actvn(self.conv_0(net))
         net = self.actvn(self.conv_0_1(net))
         net = self.conv0_1_bn(net)
-        feature_2 = F.grid_sample(net, p, padding_mode='border', align_corners = True)
-        net = self.maxpool(net)  # out 32
+        feature_2 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
+        net = self.maxpool(net)    # out 32
 
         net = self.actvn(self.conv_1(net))
         net = self.actvn(self.conv_1_1(net))
         net = self.conv1_1_bn(net)
-        feature_3 = F.grid_sample(net, p, padding_mode='border', align_corners = True)
-        net = self.maxpool(net)  # out 16
+        feature_3 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
+        net = self.maxpool(net)    # out 16
 
         net = self.actvn(self.conv_2(net))
         net = self.actvn(self.conv_2_1(net))
         net = self.conv2_1_bn(net)
-        feature_4 = F.grid_sample(net, p, padding_mode='border', align_corners = True)
-        net = self.maxpool(net)  # out 8
+        feature_4 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
+        net = self.maxpool(net)    # out 8
 
         net = self.actvn(self.conv_3(net))
         net = self.actvn(self.conv_3_1(net))
         net = self.conv3_1_bn(net)
-        feature_5 = F.grid_sample(net, p, padding_mode='border', align_corners = True)
-        net = self.maxpool(net)  # out 4
+        feature_5 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
+        net = self.maxpool(net)    # out 4
 
         net = self.actvn(self.conv_4(net))
         net = self.actvn(self.conv_4_1(net))
         net = self.conv4_1_bn(net)
-        feature_6 = F.grid_sample(net, p, padding_mode='border', align_corners = True)  # out 2
+        feature_6 = F.grid_sample(net, p, padding_mode='border', align_corners=True)    # out 2
 
         # 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]))  # (B, featues_per_sample, samples_num)
+            features, (shape[0], shape[1] * shape[3], shape[4])
+        )    # (B, featues_per_sample, samples_num)
         # (B, featue_size, samples_num)
         features = torch.cat((features, p_features), dim=1)
 
@@ -183,7 +191,7 @@ class IFGeoNet(nn.Module):
         net = self.actvn(self.fc_1(net))
         net = self.actvn(self.fc_2(net))
         net = self.fc_out(net).squeeze(1)
-        
+
         return net
 
     def compute_loss(self, prds, tgts):

lib/net/IFGeoNet_nobody.py

@@ -8,16 +8,17 @@ from lib.dataset.mesh_util import read_smpl_constants, SMPLX
 
 
 class SelfAttention(torch.nn.Module):
-
     def __init__(self, in_channels, out_channels):
         super().__init__()
         self.conv = nn.Conv3d(in_channels, out_channels, 3, padding=1, padding_mode='replicate')
-        self.attention = nn.Conv3d(in_channels,
-                                   out_channels,
-                                   kernel_size=3,
-                                   padding=1,
-                                   padding_mode='replicate',
-                                   bias=False)
+        self.attention = nn.Conv3d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            padding_mode='replicate',
+            bias=False
+        )
         with torch.no_grad():
             self.attention.weight.copy_(torch.zeros_like(self.attention.weight))
 
@@ -28,34 +29,39 @@ class SelfAttention(torch.nn.Module):
 
 
 class IFGeoNet(nn.Module):
-
     def __init__(self, cfg, hidden_dim=256):
         super(IFGeoNet, self).__init__()
 
-        self.conv_in_partial = nn.Conv3d(1, 16, 3, padding=1,
-                                         padding_mode='replicate')  # out: 256 ->m.p. 128
+        self.conv_in_partial = nn.Conv3d(
+            1, 16, 3, padding=1, padding_mode='replicate'
+        )    # out: 256 ->m.p. 128
 
         self.SA = SelfAttention(4, 4)
-        self.conv_0_fusion = nn.Conv3d(16, 32, 3, padding=1, padding_mode='replicate')  # out: 128
-        self.conv_0_1_fusion = nn.Conv3d(32, 32, 3, padding=1,
-                                         padding_mode='replicate')  # out: 128 ->m.p. 64
-
-        self.conv_0 = nn.Conv3d(32, 32, 3, padding=1, padding_mode='replicate')  # out: 128
-        self.conv_0_1 = nn.Conv3d(32, 32, 3, padding=1,
-                                  padding_mode='replicate')  # out: 128 ->m.p. 64
-
-        self.conv_1 = nn.Conv3d(32, 64, 3, padding=1, padding_mode='replicate')  # out: 64
-        self.conv_1_1 = nn.Conv3d(64, 64, 3, padding=1,
-                                  padding_mode='replicate')  # out: 64 -> mp 32
-
-        self.conv_2 = nn.Conv3d(64, 128, 3, padding=1, padding_mode='replicate')  # out: 32
-        self.conv_2_1 = nn.Conv3d(128, 128, 3, padding=1,
-                                  padding_mode='replicate')  # out: 32 -> mp 16
-        self.conv_3 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 16
-        self.conv_3_1 = nn.Conv3d(128, 128, 3, padding=1,
-                                  padding_mode='replicate')  # out: 16 -> mp 8
-        self.conv_4 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 8
-        self.conv_4_1 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')  # out: 8
+        self.conv_0_fusion = nn.Conv3d(16, 32, 3, padding=1, padding_mode='replicate')    # out: 128
+        self.conv_0_1_fusion = nn.Conv3d(
+            32, 32, 3, padding=1, padding_mode='replicate'
+        )    # out: 128 ->m.p. 64
+
+        self.conv_0 = nn.Conv3d(32, 32, 3, padding=1, padding_mode='replicate')    # out: 128
+        self.conv_0_1 = nn.Conv3d(
+            32, 32, 3, padding=1, padding_mode='replicate'
+        )    # out: 128 ->m.p. 64
+
+        self.conv_1 = nn.Conv3d(32, 64, 3, padding=1, padding_mode='replicate')    # out: 64
+        self.conv_1_1 = nn.Conv3d(
+            64, 64, 3, padding=1, padding_mode='replicate'
+        )    # out: 64 -> mp 32
+
+        self.conv_2 = nn.Conv3d(64, 128, 3, padding=1, padding_mode='replicate')    # out: 32
+        self.conv_2_1 = nn.Conv3d(
+            128, 128, 3, padding=1, padding_mode='replicate'
+        )    # out: 32 -> mp 16
+        self.conv_3 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 16
+        self.conv_3_1 = nn.Conv3d(
+            128, 128, 3, padding=1, padding_mode='replicate'
+        )    # out: 16 -> mp 8
+        self.conv_4 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 8
+        self.conv_4_1 = nn.Conv3d(128, 128, 3, padding=1, padding_mode='replicate')    # out: 8
 
         feature_size = (1 + 32 + 32 + 64 + 128 + 128 + 128) + 3
         self.fc_0 = nn.Conv1d(feature_size, hidden_dim * 2, 1)
@@ -95,8 +101,8 @@ class IFGeoNet(nn.Module):
     def forward(self, batch):
 
         p = orthogonal(batch["samples_geo"].permute(0, 2, 1),
-                       batch["calib"]).permute(0, 2, 1)  #[2, 60000, 3]
-        x = batch["depth_voxels"]  #[B, 128, 128, 128]
+                       batch["calib"]).permute(0, 2, 1)    #[2, 60000, 3]
+        x = batch["depth_voxels"]    #[B, 128, 128, 128]
 
         x = x.unsqueeze(1)
         p_features = p.transpose(1, -1)
@@ -106,7 +112,7 @@ class IFGeoNet(nn.Module):
         feature_0_partial = F.grid_sample(x, p, padding_mode='border', align_corners=True)
         net_partial = self.actvn(self.conv_in_partial(x))
         net_partial = self.partial_conv_in_bn(net_partial)
-        net_partial = self.maxpool(net_partial)  # out 64
+        net_partial = self.maxpool(net_partial)    # out 64
 
         # Feature fusion
         net = self.actvn(self.conv_0_fusion(net_partial))
@@ -119,40 +125,44 @@ class IFGeoNet(nn.Module):
         net = self.actvn(self.conv_0_1(net))
         net = self.conv0_1_bn(net)
         feature_2 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
-        net = self.maxpool(net)  # out 32
+        net = self.maxpool(net)    # out 32
 
         net = self.actvn(self.conv_1(net))
         net = self.actvn(self.conv_1_1(net))
         net = self.conv1_1_bn(net)
         feature_3 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
-        net = self.maxpool(net)  # out 16
+        net = self.maxpool(net)    # out 16
 
         net = self.actvn(self.conv_2(net))
         net = self.actvn(self.conv_2_1(net))
         net = self.conv2_1_bn(net)
         feature_4 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
-        net = self.maxpool(net)  # out 8
+        net = self.maxpool(net)    # out 8
 
         net = self.actvn(self.conv_3(net))
         net = self.actvn(self.conv_3_1(net))
         net = self.conv3_1_bn(net)
         feature_5 = F.grid_sample(net, p, padding_mode='border', align_corners=True)
-        net = self.maxpool(net)  # out 4
+        net = self.maxpool(net)    # out 4
 
         net = self.actvn(self.conv_4(net))
         net = self.actvn(self.conv_4_1(net))
         net = self.conv4_1_bn(net)
-        feature_6 = F.grid_sample(net, p, padding_mode='border', align_corners=True)  # out 2
+        feature_6 = F.grid_sample(net, p, padding_mode='border', align_corners=True)    # out 2
 
         # 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]))  # (B, featues_per_sample, samples_num)
+            features, (shape[0], shape[1] * shape[3], shape[4])
+        )    # (B, featues_per_sample, samples_num)
         # (B, featue_size, samples_num)
         features = torch.cat((features, p_features), dim=1)
 
@@ -167,4 +177,4 @@ class IFGeoNet(nn.Module):
 
         loss = self.l1_loss(prds, tgts)
 
-        return loss
+        return loss

lib/net/NormalNet.py

@@ -35,7 +35,6 @@ class NormalNet(BasePIFuNet):
         4. Classification.
         5. During training, error is calculated on all stacks.
     """
-
     def __init__(self, cfg):
 
         super(NormalNet, self).__init__()
@@ -65,9 +64,11 @@ class NormalNet(BasePIFuNet):
             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"])
+            [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"])
+            [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")
@@ -134,18 +135,20 @@ class NormalNet(BasePIFuNet):
         if 'mrf' in self.F_losses:
             mrf_F_loss = self.mrf_loss(
                 F.interpolate(prd_F, scale_factor=scale_factor, mode='bicubic', align_corners=True),
-                F.interpolate(tgt_F, scale_factor=scale_factor, mode='bicubic', align_corners=True))
+                F.interpolate(tgt_F, scale_factor=scale_factor, mode='bicubic', align_corners=True)
+            )
             total_loss["netF"] += self.F_losses_ratio[self.F_losses.index('mrf')] * mrf_F_loss
             total_loss["mrf_F"] = self.F_losses_ratio[self.F_losses.index('mrf')] * mrf_F_loss
         if 'mrf' in self.B_losses:
             mrf_B_loss = self.mrf_loss(
                 F.interpolate(prd_B, scale_factor=scale_factor, mode='bicubic', align_corners=True),
-                F.interpolate(tgt_B, scale_factor=scale_factor, mode='bicubic', align_corners=True))
+                F.interpolate(tgt_B, scale_factor=scale_factor, mode='bicubic', align_corners=True)
+            )
             total_loss["netB"] += self.B_losses_ratio[self.B_losses.index('mrf')] * mrf_B_loss
             total_loss["mrf_B"] = self.B_losses_ratio[self.B_losses.index('mrf')] * mrf_B_loss
 
         if 'gan' in self.ALL_losses:
-            
+
             total_loss["netD"] = 0.0
 
             pred_fake = self.netD.forward(prd_B)
@@ -154,8 +157,8 @@ class NormalNet(BasePIFuNet):
             loss_D_real = self.gan_loss(pred_real, True)
             loss_G_fake = self.gan_loss(pred_fake, True)
 
-            total_loss["netD"] += 0.5 * (
-                loss_D_fake + loss_D_real) * self.B_losses_ratio[self.B_losses.index('gan')]
+            total_loss["netD"] += 0.5 * (loss_D_fake + loss_D_real
+                                        ) * self.B_losses_ratio[self.B_losses.index('gan')]
             total_loss["D_fake"] = loss_D_fake * self.B_losses_ratio[self.B_losses.index('gan')]
             total_loss["D_real"] = loss_D_real * self.B_losses_ratio[self.B_losses.index('gan')]
 
@@ -167,8 +170,8 @@ class NormalNet(BasePIFuNet):
                 for i in range(2):
                     for j in range(len(pred_fake[i]) - 1):
                         loss_G_GAN_Feat += self.l1_loss(pred_fake[i][j], pred_real[i][j].detach())
-                total_loss["netB"] += loss_G_GAN_Feat * self.B_losses_ratio[self.B_losses.index(
-                    'gan_feat')]
+                total_loss["netB"] += loss_G_GAN_Feat * self.B_losses_ratio[
+                    self.B_losses.index('gan_feat')]
                 total_loss["G_GAN_Feat"] = loss_G_GAN_Feat * self.B_losses_ratio[
                     self.B_losses.index('gan_feat')]
 

lib/net/geometry.py

@@ -19,12 +19,12 @@ import numpy as np
 import numbers
 from torch.nn import functional as F
 from einops.einops import rearrange
-
 """
 Useful geometric operations, e.g. Perspective projection and a differentiable Rodrigues formula
 Parts of the code are taken from https://github.com/MandyMo/pytorch_HMR
 """
 
+
 def quaternion_to_rotation_matrix(quat):
     """Convert quaternion coefficients to rotation matrix.
     Args:
@@ -42,11 +42,13 @@ 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
 
 
@@ -56,7 +58,7 @@ def index(feat, uv):
     :param uv: [B, 2, N] uv coordinates in the image plane, range [0, 1]
     :return: [B, C, N] image features at the uv coordinates
     """
-    uv = uv.transpose(1, 2)  # [B, N, 2]
+    uv = uv.transpose(1, 2)    # [B, N, 2]
 
     (B, N, _) = uv.shape
     C = feat.shape[1]
@@ -64,14 +66,14 @@ def index(feat, uv):
     if uv.shape[-1] == 3:
         # uv = uv[:,:,[2,1,0]]
         # uv = uv * torch.tensor([1.0,-1.0,1.0]).type_as(uv)[None,None,...]
-        uv = uv.unsqueeze(2).unsqueeze(3)  # [B, N, 1, 1, 3]
+        uv = uv.unsqueeze(2).unsqueeze(3)    # [B, N, 1, 1, 3]
     else:
-        uv = uv.unsqueeze(2)  # [B, N, 1, 2]
+        uv = uv.unsqueeze(2)    # [B, N, 1, 2]
 
     # NOTE: for newer PyTorch, it seems that training results are degraded due to implementation diff in F.grid_sample
     # for old versions, simply remove the aligned_corners argument.
-    samples = torch.nn.functional.grid_sample(feat, uv, align_corners=True)  # [B, C, N, 1]
-    return samples.view(B, C, N)  # [B, C, N]
+    samples = torch.nn.functional.grid_sample(feat, uv, align_corners=True)    # [B, C, N, 1]
+    return samples.view(B, C, N)    # [B, C, N]
 
 
 def orthogonal(points, calibrations, transforms=None):
@@ -84,7 +86,7 @@ def orthogonal(points, calibrations, transforms=None):
     """
     rot = calibrations[:, :3, :3]
     trans = calibrations[:, :3, 3:4]
-    pts = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    pts = torch.baddbmm(trans, rot, points)    # [B, 3, N]
     if transforms is not None:
         scale = transforms[:2, :2]
         shift = transforms[:2, 2:3]
@@ -102,7 +104,7 @@ def perspective(points, calibrations, transforms=None):
     """
     rot = calibrations[:, :3, :3]
     trans = calibrations[:, :3, 3:4]
-    homo = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    homo = torch.baddbmm(trans, rot, points)    # [B, 3, N]
     xy = homo[:, :2, :] / homo[:, 2:3, :]
     if transforms is not None:
         scale = transforms[:2, :2]
@@ -187,7 +189,8 @@ def rotation_matrix_to_angle_axis(rotation_matrix):
     if rotation_matrix.shape[1:] == (3, 3):
         rot_mat = rotation_matrix.reshape(-1, 3, 3)
         hom = torch.tensor([0, 0, 1], dtype=torch.float32, device=rotation_matrix.device).reshape(
-            1, 3, 1).expand(rot_mat.shape[0], -1, -1)
+            1, 3, 1
+        ).expand(rot_mat.shape[0], -1, -1)
         rotation_matrix = torch.cat([rot_mat, hom], dim=-1)
 
     quaternion = rotation_matrix_to_quaternion(rotation_matrix)
@@ -222,8 +225,9 @@ def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
         raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(quaternion)))
 
     if not quaternion.shape[-1] == 4:
-        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}".format(
-            quaternion.shape))
+        raise ValueError(
+            "Input must be a tensor of shape Nx4 or 4. Got {}".format(quaternion.shape)
+        )
     # unpack input and compute conversion
     q1: torch.Tensor = quaternion[..., 1]
     q2: torch.Tensor = quaternion[..., 2]
@@ -276,11 +280,13 @@ def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
         raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(rotation_matrix)))
 
     if len(rotation_matrix.shape) > 3:
-        raise ValueError("Input size must be a three dimensional tensor. Got {}".format(
-            rotation_matrix.shape))
+        raise ValueError(
+            "Input size must be a three dimensional tensor. Got {}".format(rotation_matrix.shape)
+        )
     if not rotation_matrix.shape[-2:] == (3, 4):
-        raise ValueError("Input size must be a N x 3 x 4  tensor. Got {}".format(
-            rotation_matrix.shape))
+        raise ValueError(
+            "Input size must be a N x 3 x 4  tensor. Got {}".format(rotation_matrix.shape)
+        )
 
     rmat_t = torch.transpose(rotation_matrix, 1, 2)
 
@@ -347,8 +353,10 @@ def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
     mask_c3 = mask_c3.view(-1, 1).type_as(q3)
 
     q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
-    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + t2_rep * mask_c2  # noqa
-                    + t3_rep * mask_c3)  # noqa
+    q /= torch.sqrt(
+        t0_rep * mask_c0 + t1_rep * mask_c1 + t2_rep * mask_c2    # noqa
+        + t3_rep * mask_c3
+    )    # noqa
     q *= 0.5
     return q
 
@@ -389,6 +397,7 @@ def rot6d_to_rotmat(x):
         mat = torch.stack((b1, b2, b3), dim=-1)
     return mat
 
+
 def rotmat_to_rot6d(x):
     """Convert 3x3 rotation matrix to 6D rotation representation.
     Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
@@ -402,6 +411,7 @@ def rotmat_to_rot6d(x):
     x = x.reshape(batch_size, 6)
     return x
 
+
 def rotmat_to_angle(x):
     """Convert rotation to one-D angle.
     Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
@@ -440,12 +450,9 @@ def projection(pred_joints, pred_camera, retain_z=False):
     return pred_keypoints_2d
 
 
-def perspective_projection(points,
-                           rotation,
-                           translation,
-                           focal_length,
-                           camera_center,
-                           retain_z=False):
+def perspective_projection(
+    points, rotation, translation, focal_length, camera_center, retain_z=False
+):
     """
     This function computes the perspective projection of a set of points.
     Input:
@@ -501,10 +508,12 @@ 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
@@ -558,15 +567,12 @@ def estimate_translation(S, joints_2d, focal_length=5000., img_size=224., use_al
         S_i = S[i]
         joints_i = joints_2d[i]
         conf_i = joints_conf[i]
-        trans[i] = estimate_translation_np(S_i,
-                                           joints_i,
-                                           conf_i,
-                                           focal_length=focal_length[i],
-                                           img_size=img_size[i])
+        trans[i] = estimate_translation_np(
+            S_i, joints_i, conf_i, focal_length=focal_length[i], img_size=img_size[i]
+        )
     return torch.from_numpy(trans).to(device)
 
 
-
 def Rot_y(angle, category="torch", prepend_dim=True, device=None):
     """Rotate around y-axis by angle
     Args:
@@ -574,11 +580,13 @@ 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)
@@ -600,11 +608,13 @@ 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)
@@ -626,11 +636,13 @@ 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)
@@ -672,7 +684,7 @@ def compute_twist_rotation(rotation_matrix, twist_axis):
     twist_rotation = quaternion_to_rotation_matrix(twist_quaternion)
     twist_aa = quaternion_to_angle_axis(twist_quaternion)
 
-    twist_angle = torch.sum(twist_aa, dim=1, keepdim=True) / torch.sum(
-        twist_axis, dim=1, keepdim=True)
+    twist_angle = torch.sum(twist_aa, dim=1,
+                            keepdim=True) / torch.sum(twist_axis, dim=1, keepdim=True)
 
-    return twist_rotation, twist_angle
+    return twist_rotation, twist_angle

lib/net/net_util.py

@@ -71,11 +71,10 @@ def init_weights(net, init_type="normal", init_gain=0.02):
     We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
     work better for some applications. Feel free to try yourself.
     """
-
-    def init_func(m):  # define the initialization function
+    def init_func(m):    # define the initialization function
         classname = m.__class__.__name__
-        if hasattr(m, "weight") and (classname.find("Conv") != -1 or
-                                     classname.find("Linear") != -1):
+        if hasattr(m,
+                   "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):
             if init_type == "normal":
                 init.normal_(m.weight.data, 0.0, init_gain)
             elif init_type == "xavier":
@@ -85,17 +84,19 @@ def init_weights(net, init_type="normal", init_gain=0.02):
             elif init_type == "orthogonal":
                 init.orthogonal_(m.weight.data, gain=init_gain)
             else:
-                raise NotImplementedError("initialization method [%s] is not implemented" %
-                                          init_type)
+                raise NotImplementedError(
+                    "initialization method [%s] is not implemented" % init_type
+                )
             if hasattr(m, "bias") and m.bias is not None:
                 init.constant_(m.bias.data, 0.0)
-        elif (classname.find("BatchNorm2d") !=
-              -1):  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+        elif (
+            classname.find("BatchNorm2d") != -1
+        ):    # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
             init.normal_(m.weight.data, 1.0, init_gain)
             init.constant_(m.bias.data, 0.0)
 
     # print('initialize network with %s' % init_type)
-    net.apply(init_func)  # apply the initialization function <init_func>
+    net.apply(init_func)    # apply the initialization function <init_func>
 
 
 def init_net(net, init_type="xavier", init_gain=0.02, gpu_ids=[]):
@@ -110,7 +111,7 @@ def init_net(net, init_type="xavier", init_gain=0.02, gpu_ids=[]):
     """
     if len(gpu_ids) > 0:
         assert torch.cuda.is_available()
-        net = torch.nn.DataParallel(net)  # multi-GPUs
+        net = torch.nn.DataParallel(net)    # multi-GPUs
     init_weights(net, init_type, init_gain=init_gain)
     return net
 
@@ -127,13 +128,9 @@ def imageSpaceRotation(xy, rot):
     return (disp * xy).sum(dim=1)
 
 
-def cal_gradient_penalty(netD,
-                         real_data,
-                         fake_data,
-                         device,
-                         type="mixed",
-                         constant=1.0,
-                         lambda_gp=10.0):
+def cal_gradient_penalty(
+    netD, real_data, fake_data, device, type="mixed", constant=1.0, lambda_gp=10.0
+):
     """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
 
     Arguments:
@@ -155,9 +152,11 @@ def cal_gradient_penalty(netD,
             interpolatesv = fake_data
         elif type == "mixed":
             alpha = torch.rand(real_data.shape[0], 1)
-            alpha = (alpha.expand(real_data.shape[0],
-                                  real_data.nelement() //
-                                  real_data.shape[0]).contiguous().view(*real_data.shape))
+            alpha = (
+                alpha.expand(real_data.shape[0],
+                             real_data.nelement() //
+                             real_data.shape[0]).contiguous().view(*real_data.shape)
+            )
             alpha = alpha.to(device)
             interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
         else:
@@ -172,9 +171,9 @@ def cal_gradient_penalty(netD,
             retain_graph=True,
             only_inputs=True,
         )
-        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
-        gradient_penalty = ((
-            (gradients + 1e-16).norm(2, dim=1) - constant)**2).mean() * lambda_gp  # added eps
+        gradients = gradients[0].view(real_data.size(0), -1)    # flat the data
+        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant)**
+                            2).mean() * lambda_gp    # added eps
         return gradient_penalty, gradients
     else:
         return 0.0, None
@@ -201,13 +200,11 @@ def get_norm_layer(norm_type="instance"):
 
 
 class Flatten(nn.Module):
-
     def forward(self, input):
         return input.view(input.size(0), -1)
 
 
 class ConvBlock(nn.Module):
-
     def __init__(self, in_planes, out_planes, opt):
         super(ConvBlock, self).__init__()
         [k, s, d, p] = opt.conv3x3
@@ -258,5 +255,3 @@ class ConvBlock(nn.Module):
         out3 += residual
 
         return out3
-
-

lib/net/voxelize.py

@@ -13,7 +13,6 @@ class VoxelizationFunction(Function):
     Definition of differentiable voxelization function
     Currently implemented only for cuda Tensors
     """
-
     @staticmethod
     def forward(
         ctx,
@@ -48,12 +47,15 @@ class VoxelizationFunction(Function):
         smpl_face_code = smpl_face_code.contiguous()
         smpl_tetrahedrons = smpl_tetrahedrons.contiguous()
 
-        occ_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res, ctx.volume_res,
-                                            ctx.volume_res).fill_(0.0)
-        semantic_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res, ctx.volume_res,
-                                                 ctx.volume_res, 3).fill_(0.0)
-        weight_sum_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res, ctx.volume_res,
-                                                   ctx.volume_res).fill_(1e-3)
+        occ_volume = torch.cuda.FloatTensor(
+            ctx.batch_size, ctx.volume_res, ctx.volume_res, ctx.volume_res
+        ).fill_(0.0)
+        semantic_volume = torch.cuda.FloatTensor(
+            ctx.batch_size, ctx.volume_res, ctx.volume_res, ctx.volume_res, 3
+        ).fill_(0.0)
+        weight_sum_volume = torch.cuda.FloatTensor(
+            ctx.batch_size, ctx.volume_res, ctx.volume_res, ctx.volume_res
+        ).fill_(1e-3)
 
         # occ_volume [B, volume_res, volume_res, volume_res]
         # semantic_volume [B, volume_res, volume_res, volume_res, 3]
@@ -80,7 +82,6 @@ class Voxelization(nn.Module):
     """
     Wrapper around the autograd function VoxelizationFunction
     """
-
     def __init__(
         self,
         smpl_vertex_code,
@@ -151,21 +152,25 @@ class Voxelization(nn.Module):
             self.sigma,
             self.smooth_kernel_size,
         )
-        return vol.permute((0, 4, 1, 2, 3))  # (bzyxc --> bcdhw)
+        return vol.permute((0, 4, 1, 2, 3))    # (bzyxc --> bcdhw)
 
     def vertices_to_faces(self, vertices):
         assert vertices.ndimension() == 3
         bs, nv = vertices.shape[:2]
-        face = (self.smpl_face_indices_batch +
-                (torch.arange(bs, dtype=torch.int32).to(self.device) * nv)[:, None, None])
+        face = (
+            self.smpl_face_indices_batch +
+            (torch.arange(bs, dtype=torch.int32).to(self.device) * nv)[:, None, None]
+        )
         vertices_ = vertices.reshape((bs * nv, 3))
         return vertices_[face.long()]
 
     def vertices_to_tetrahedrons(self, vertices):
         assert vertices.ndimension() == 3
         bs, nv = vertices.shape[:2]
-        tets = (self.smpl_tetraderon_indices_batch +
-                (torch.arange(bs, dtype=torch.int32).to(self.device) * nv)[:, None, None])
+        tets = (
+            self.smpl_tetraderon_indices_batch +
+            (torch.arange(bs, dtype=torch.int32).to(self.device) * nv)[:, None, None]
+        )
         vertices_ = vertices.reshape((bs * nv, 3))
         return vertices_[tets.long()]
 
@@ -174,8 +179,9 @@ class Voxelization(nn.Module):
         assert face_verts.shape[2] == 3
         assert face_verts.shape[3] == 3
         bs, nf = face_verts.shape[:2]
-        face_centers = (face_verts[:, :, 0, :] + face_verts[:, :, 1, :] +
-                        face_verts[:, :, 2, :]) / 3.0
+        face_centers = (
+            face_verts[:, :, 0, :] + face_verts[:, :, 1, :] + face_verts[:, :, 2, :]
+        ) / 3.0
         face_centers = face_centers.reshape((bs, nf, 3))
         return face_centers
 

lib/pixielib/models/FLAME.py

@@ -27,7 +27,6 @@ class FLAMETex(nn.Module):
     FLAME texture converted from BFM:
     https://github.com/TimoBolkart/BFM_to_FLAME
     """
-
     def __init__(self, config):
         super(FLAMETex, self).__init__()
         if config.tex_type == "BFM":
@@ -54,8 +53,7 @@ class FLAMETex(nn.Module):
         n_tex = config.n_tex
         num_components = texture_basis.shape[1]
         texture_mean = torch.from_numpy(texture_mean).float()[None, ...]
-        texture_basis = torch.from_numpy(
-            texture_basis[:, :n_tex]).float()[None, ...]
+        texture_basis = torch.from_numpy(texture_basis[:, :n_tex]).float()[None, ...]
         self.register_buffer("texture_mean", texture_mean)
         self.register_buffer("texture_basis", texture_basis)
 
@@ -64,10 +62,8 @@ class FLAMETex(nn.Module):
         texcode: [batchsize, n_tex]
         texture: [bz, 3, 256, 256], range: 0-1
         """
-        texture = self.texture_mean + (self.texture_basis *
-                                       texcode[:, None, :]).sum(-1)
-        texture = texture.reshape(texcode.shape[0], 512, 512,
-                                  3).permute(0, 3, 1, 2)
+        texture = self.texture_mean + (self.texture_basis * texcode[:, None, :]).sum(-1)
+        texture = texture.reshape(texcode.shape[0], 512, 512, 3).permute(0, 3, 1, 2)
         texture = F.interpolate(texture, [256, 256])
         texture = texture[:, [2, 1, 0], :, :]
         return texture
@@ -78,13 +74,13 @@ def texture_flame2smplx(cached_data, flame_texture, smplx_texture):
     TODO: pytorch version ==> grid sample
     """
     if smplx_texture.shape[0] != smplx_texture.shape[1]:
-        print("SMPL-X texture not squared (%d != %d)" %
-              (smplx_texture[0], smplx_texture[1]))
+        print("SMPL-X texture not squared (%d != %d)" % (smplx_texture[0], smplx_texture[1]))
         return
     if smplx_texture.shape[0] != cached_data["target_resolution"]:
         print(
-            "SMPL-X texture size does not match cached image resolution (%d != %d)"
-            % (smplx_texture.shape[0], cached_data["target_resolution"]))
+            "SMPL-X texture size does not match cached image resolution (%d != %d)" %
+            (smplx_texture.shape[0], cached_data["target_resolution"])
+        )
         return
     x_coords = cached_data["x_coords"]
     y_coords = cached_data["y_coords"]
@@ -98,11 +94,13 @@ def texture_flame2smplx(cached_data, flame_texture, smplx_texture):
         flame_texture.shape[0],
     ).astype(int)
     source_tex_coords[:, 1] = np.clip(
-        flame_texture.shape[1] * (source_uv_points[:, 0]), 0.0,
-        flame_texture.shape[1]).astype(int)
+        flame_texture.shape[1] * (source_uv_points[:, 0]), 0.0, flame_texture.shape[1]
+    ).astype(int)
 
     smplx_texture[y_coords[target_pixel_ids].astype(int),
-                  x_coords[target_pixel_ids].astype(int), :, ] = flame_texture[
-                      source_tex_coords[:, 0], source_tex_coords[:, 1]]
+                  x_coords[target_pixel_ids].astype(int), :, ] = flame_texture[source_tex_coords[:,
+                                                                                                 0],
+                                                                               source_tex_coords[:,
+                                                                                                 1]]
 
     return smplx_texture

lib/pixielib/models/SMPLX.py

@@ -209,452 +209,468 @@ 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]
 
@@ -691,13 +707,12 @@ class SMPLX(nn.Module):
     Given smplx parameters, this class generates a differentiable SMPLX function
     which outputs a mesh and 3D joints
     """
-
     def __init__(self, config):
         super(SMPLX, self).__init__()
         # print("creating the SMPLX Decoder")
         ss = np.load(config.smplx_model_path, allow_pickle=True)
         smplx_model = Struct(**ss)
-        
+
         self.dtype = torch.float32
         self.register_buffer(
             "faces_tensor",
@@ -705,8 +720,8 @@ class SMPLX(nn.Module):
         )
         # The vertices of the template model
         self.register_buffer(
-            "v_template",
-            to_tensor(to_np(smplx_model.v_template), dtype=self.dtype))
+            "v_template", to_tensor(to_np(smplx_model.v_template), dtype=self.dtype)
+        )
         # The shape components and expression
         # expression space is the same as FLAME
         shapedirs = to_tensor(to_np(smplx_model.shapedirs), dtype=self.dtype)
@@ -721,21 +736,18 @@ class SMPLX(nn.Module):
         # The pose components
         num_pose_basis = smplx_model.posedirs.shape[-1]
         posedirs = np.reshape(smplx_model.posedirs, [-1, num_pose_basis]).T
-        self.register_buffer("posedirs",
-                             to_tensor(to_np(posedirs), dtype=self.dtype))
+        self.register_buffer("posedirs", to_tensor(to_np(posedirs), dtype=self.dtype))
         self.register_buffer(
-            "J_regressor",
-            to_tensor(to_np(smplx_model.J_regressor), dtype=self.dtype))
+            "J_regressor", to_tensor(to_np(smplx_model.J_regressor), dtype=self.dtype)
+        )
         parents = to_tensor(to_np(smplx_model.kintree_table[0])).long()
         parents[0] = -1
         self.register_buffer("parents", parents)
-        self.register_buffer(
-            "lbs_weights",
-            to_tensor(to_np(smplx_model.weights), dtype=self.dtype))
+        self.register_buffer("lbs_weights", to_tensor(to_np(smplx_model.weights), dtype=self.dtype))
         # for face keypoints
         self.register_buffer(
-            "lmk_faces_idx",
-            torch.tensor(smplx_model.lmk_faces_idx, dtype=torch.long))
+            "lmk_faces_idx", torch.tensor(smplx_model.lmk_faces_idx, dtype=torch.long)
+        )
         self.register_buffer(
             "lmk_bary_coords",
             torch.tensor(smplx_model.lmk_bary_coords, dtype=self.dtype),
@@ -746,24 +758,20 @@ class SMPLX(nn.Module):
         )
         self.register_buffer(
             "dynamic_lmk_bary_coords",
-            torch.tensor(smplx_model.dynamic_lmk_bary_coords,
-                         dtype=self.dtype),
+            torch.tensor(smplx_model.dynamic_lmk_bary_coords, dtype=self.dtype),
         )
         # pelvis to head, to calculate head yaw angle, then find the dynamic landmarks
-        self.register_buffer("head_kin_chain",
-                             torch.tensor(head_kin_chain, dtype=torch.long))
+        self.register_buffer("head_kin_chain", torch.tensor(head_kin_chain, dtype=torch.long))
 
         # -- initialize parameters
         # shape and expression
         self.register_buffer(
             "shape_params",
-            nn.Parameter(torch.zeros([1, config.n_shape], dtype=self.dtype),
-                         requires_grad=False),
+            nn.Parameter(torch.zeros([1, config.n_shape], dtype=self.dtype), requires_grad=False),
         )
         self.register_buffer(
             "expression_params",
-            nn.Parameter(torch.zeros([1, config.n_exp], dtype=self.dtype),
-                         requires_grad=False),
+            nn.Parameter(torch.zeros([1, config.n_exp], dtype=self.dtype), requires_grad=False),
         )
         # pose: represented as rotation matrx [number of joints, 3, 3]
         self.register_buffer(
@@ -824,8 +832,7 @@ class SMPLX(nn.Module):
         )
 
         if config.extra_joint_path:
-            self.extra_joint_selector = JointsFromVerticesSelector(
-                fname=config.extra_joint_path)
+            self.extra_joint_selector = JointsFromVerticesSelector(fname=config.extra_joint_path)
         self.use_joint_regressor = True
         self.keypoint_names = SMPLX_names
         if self.use_joint_regressor:
@@ -843,7 +850,8 @@ class SMPLX(nn.Module):
             self.register_buffer("target_idxs", torch.from_numpy(target))
             self.register_buffer(
                 "extra_joint_regressor",
-                torch.from_numpy(j14_regressor).to(torch.float32))
+                torch.from_numpy(j14_regressor).to(torch.float32)
+            )
             self.part_indices = part_indices
 
     def forward(
@@ -880,23 +888,17 @@ class SMPLX(nn.Module):
         if expression_params is None:
             expression_params = self.expression_params.expand(batch_size, -1)
         if global_pose is None:
-            global_pose = self.global_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            global_pose = self.global_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
         if body_pose is None:
-            body_pose = self.body_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            body_pose = self.body_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
         if jaw_pose is None:
-            jaw_pose = self.jaw_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            jaw_pose = self.jaw_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
         if eye_pose is None:
-            eye_pose = self.eye_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            eye_pose = self.eye_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
         if left_hand_pose is None:
-            left_hand_pose = self.left_hand_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            left_hand_pose = self.left_hand_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
         if right_hand_pose is None:
-            right_hand_pose = self.right_hand_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+            right_hand_pose = self.right_hand_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
 
         shape_components = torch.cat([shape_params, expression_params], dim=1)
         full_pose = torch.cat(
@@ -910,8 +912,7 @@ class SMPLX(nn.Module):
             ],
             dim=1,
         )
-        template_vertices = self.v_template.unsqueeze(0).expand(
-            batch_size, -1, -1)
+        template_vertices = self.v_template.unsqueeze(0).expand(batch_size, -1, -1)
         # smplx
         vertices, joints = lbs(
             shape_components,
@@ -926,10 +927,8 @@ class SMPLX(nn.Module):
             pose2rot=False,
         )
         # face dynamic landmarks
-        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(
-            batch_size, -1)
-        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).expand(
-            batch_size, -1, -1)
+        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1)
+        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).expand(batch_size, -1, -1)
         dyn_lmk_faces_idx, dyn_lmk_bary_coords = find_dynamic_lmk_idx_and_bcoords(
             vertices,
             full_pose,
@@ -939,14 +938,12 @@ class SMPLX(nn.Module):
         )
         lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
         lmk_bary_coords = torch.cat([lmk_bary_coords, dyn_lmk_bary_coords], 1)
-        landmarks = vertices2landmarks(vertices, self.faces_tensor,
-                                       lmk_faces_idx, lmk_bary_coords)
+        landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords)
 
         final_joint_set = [joints, landmarks]
         if hasattr(self, "extra_joint_selector"):
             # Add any extra joints that might be needed
-            extra_joints = self.extra_joint_selector(vertices,
-                                                     self.faces_tensor)
+            extra_joints = self.extra_joint_selector(vertices, self.faces_tensor)
             final_joint_set.append(extra_joints)
         # Create the final joint set
         joints = torch.cat(final_joint_set, dim=1)
@@ -978,16 +975,15 @@ class SMPLX(nn.Module):
             # -> Left elbow -> Left wrist
             kin_chain = [20, 18, 16, 13, 9, 6, 3, 0]
         else:
-            raise NotImplementedError(
-                f"pose_abs2rel does not support: {abs_joint}")
+            raise NotImplementedError(f"pose_abs2rel does not support: {abs_joint}")
 
         batch_size = global_pose.shape[0]
         dtype = global_pose.dtype
         device = global_pose.device
         full_pose = torch.cat([global_pose, body_pose], dim=1)
-        rel_rot_mat = (torch.eye(3, device=device,
-                                 dtype=dtype).unsqueeze_(dim=0).repeat(
-                                     batch_size, 1, 1))
+        rel_rot_mat = (
+            torch.eye(3, device=device, dtype=dtype).unsqueeze_(dim=0).repeat(batch_size, 1, 1)
+        )
         for idx in kin_chain[1:]:
             rel_rot_mat = torch.bmm(full_pose[:, idx], rel_rot_mat)
 
@@ -1027,11 +1023,8 @@ class SMPLX(nn.Module):
             # -> Left elbow -> Left wrist
             kin_chain = [20, 18, 16, 13, 9, 6, 3, 0]
         else:
-            raise NotImplementedError(
-                f"pose_rel2abs does not support: {abs_joint}")
-        rel_rot_mat = torch.eye(3,
-                                device=full_pose.device,
-                                dtype=full_pose.dtype).unsqueeze_(dim=0)
+            raise NotImplementedError(f"pose_rel2abs does not support: {abs_joint}")
+        rel_rot_mat = torch.eye(3, device=full_pose.device, dtype=full_pose.dtype).unsqueeze_(dim=0)
         for idx in kin_chain:
             rel_rot_mat = torch.matmul(full_pose[:, idx], rel_rot_mat)
         abs_pose = rel_rot_mat[:, None, :, :]

lib/pixielib/models/encoders.py

@@ -5,14 +5,13 @@ import torch.nn.functional as F
 
 
 class ResnetEncoder(nn.Module):
-
     def __init__(self, append_layers=None):
         super(ResnetEncoder, self).__init__()
         from . import resnet
 
         # feature_size = 2048
         self.feature_dim = 2048
-        self.encoder = resnet.load_ResNet50Model()  # out: 2048
+        self.encoder = resnet.load_ResNet50Model()    # out: 2048
         # regressor
         self.append_layers = append_layers
 
@@ -25,7 +24,6 @@ class ResnetEncoder(nn.Module):
 
 
 class MLP(nn.Module):
-
     def __init__(self, channels=[2048, 1024, 1], last_op=None):
         super(MLP, self).__init__()
         layers = []
@@ -45,13 +43,12 @@ class MLP(nn.Module):
 
 
 class HRNEncoder(nn.Module):
-
     def __init__(self, append_layers=None):
         super(HRNEncoder, self).__init__()
         from . import hrnet
 
         self.feature_dim = 2048
-        self.encoder = hrnet.load_HRNet(pretrained=True)  # out: 2048
+        self.encoder = hrnet.load_HRNet(pretrained=True)    # out: 2048
         # regressor
         self.append_layers = append_layers
 

lib/pixielib/models/hrnet.py

@@ -15,38 +15,42 @@ def load_HRNet(pretrained=False):
     hr_net_cfg_dict = {
         "use_old_impl": False,
         "pretrained_layers": ["*"],
-        "stage1": {
-            "num_modules": 1,
-            "num_branches": 1,
-            "num_blocks": [4],
-            "num_channels": [64],
-            "block": "BOTTLENECK",
-            "fuse_method": "SUM",
-        },
-        "stage2": {
-            "num_modules": 1,
-            "num_branches": 2,
-            "num_blocks": [4, 4],
-            "num_channels": [48, 96],
-            "block": "BASIC",
-            "fuse_method": "SUM",
-        },
-        "stage3": {
-            "num_modules": 4,
-            "num_branches": 3,
-            "num_blocks": [4, 4, 4],
-            "num_channels": [48, 96, 192],
-            "block": "BASIC",
-            "fuse_method": "SUM",
-        },
-        "stage4": {
-            "num_modules": 3,
-            "num_branches": 4,
-            "num_blocks": [4, 4, 4, 4],
-            "num_channels": [48, 96, 192, 384],
-            "block": "BASIC",
-            "fuse_method": "SUM",
-        },
+        "stage1":
+            {
+                "num_modules": 1,
+                "num_branches": 1,
+                "num_blocks": [4],
+                "num_channels": [64],
+                "block": "BOTTLENECK",
+                "fuse_method": "SUM",
+            },
+        "stage2":
+            {
+                "num_modules": 1,
+                "num_branches": 2,
+                "num_blocks": [4, 4],
+                "num_channels": [48, 96],
+                "block": "BASIC",
+                "fuse_method": "SUM",
+            },
+        "stage3":
+            {
+                "num_modules": 4,
+                "num_branches": 3,
+                "num_blocks": [4, 4, 4],
+                "num_channels": [48, 96, 192],
+                "block": "BASIC",
+                "fuse_method": "SUM",
+            },
+        "stage4":
+            {
+                "num_modules": 3,
+                "num_branches": 4,
+                "num_blocks": [4, 4, 4, 4],
+                "num_channels": [48, 96, 192, 384],
+                "block": "BASIC",
+                "fuse_method": "SUM",
+            },
     }
     hr_net_cfg = hr_net_cfg_dict
     model = HighResolutionNet(hr_net_cfg)
@@ -55,7 +59,6 @@ def load_HRNet(pretrained=False):
 
 
 class HighResolutionModule(nn.Module):
-
     def __init__(
         self,
         num_branches,
@@ -67,8 +70,7 @@ class HighResolutionModule(nn.Module):
         multi_scale_output=True,
     ):
         super(HighResolutionModule, self).__init__()
-        self._check_branches(num_branches, blocks, num_blocks, num_inchannels,
-                             num_channels)
+        self._check_branches(num_branches, blocks, num_blocks, num_inchannels, num_channels)
 
         self.num_inchannels = num_inchannels
         self.fuse_method = fuse_method
@@ -76,37 +78,33 @@ class HighResolutionModule(nn.Module):
 
         self.multi_scale_output = multi_scale_output
 
-        self.branches = self._make_branches(num_branches, blocks, num_blocks,
-                                            num_channels)
+        self.branches = self._make_branches(num_branches, blocks, num_blocks, num_channels)
         self.fuse_layers = self._make_fuse_layers()
         self.relu = nn.ReLU(True)
 
-    def _check_branches(self, num_branches, blocks, num_blocks, num_inchannels,
-                        num_channels):
+    def _check_branches(self, num_branches, blocks, num_blocks, num_inchannels, num_channels):
         if num_branches != len(num_blocks):
-            error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format(
-                num_branches, len(num_blocks))
+            error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format(num_branches, len(num_blocks))
             raise ValueError(error_msg)
 
         if num_branches != len(num_channels):
             error_msg = "NUM_BRANCHES({}) <> NUM_CHANNELS({})".format(
-                num_branches, len(num_channels))
+                num_branches, len(num_channels)
+            )
             raise ValueError(error_msg)
 
         if num_branches != len(num_inchannels):
             error_msg = "NUM_BRANCHES({}) <> NUM_INCHANNELS({})".format(
-                num_branches, len(num_inchannels))
+                num_branches, len(num_inchannels)
+            )
             raise ValueError(error_msg)
 
-    def _make_one_branch(self,
-                         branch_index,
-                         block,
-                         num_blocks,
-                         num_channels,
-                         stride=1):
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1):
         downsample = None
-        if (stride != 1 or self.num_inchannels[branch_index] !=
-                num_channels[branch_index] * block.expansion):
+        if (
+            stride != 1 or
+            self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion
+        ):
             downsample = nn.Sequential(
                 nn.Conv2d(
                     self.num_inchannels[branch_index],
@@ -115,8 +113,7 @@ class HighResolutionModule(nn.Module):
                     stride=stride,
                     bias=False,
                 ),
-                nn.BatchNorm2d(num_channels[branch_index] * block.expansion,
-                               momentum=BN_MOMENTUM),
+                nn.BatchNorm2d(num_channels[branch_index] * block.expansion, momentum=BN_MOMENTUM),
             )
 
         layers = []
@@ -126,13 +123,11 @@ class HighResolutionModule(nn.Module):
                 num_channels[branch_index],
                 stride,
                 downsample,
-            ))
-        self.num_inchannels[
-            branch_index] = num_channels[branch_index] * block.expansion
+            )
+        )
+        self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
         for i in range(1, num_blocks[branch_index]):
-            layers.append(
-                block(self.num_inchannels[branch_index],
-                      num_channels[branch_index]))
+            layers.append(block(self.num_inchannels[branch_index], num_channels[branch_index]))
 
         return nn.Sequential(*layers)
 
@@ -140,8 +135,7 @@ class HighResolutionModule(nn.Module):
         branches = []
 
         for i in range(num_branches):
-            branches.append(
-                self._make_one_branch(i, block, num_blocks, num_channels))
+            branches.append(self._make_one_branch(i, block, num_blocks, num_channels))
 
         return nn.ModuleList(branches)
 
@@ -167,9 +161,9 @@ class HighResolutionModule(nn.Module):
                                 bias=False,
                             ),
                             nn.BatchNorm2d(num_inchannels[i]),
-                            nn.Upsample(scale_factor=2**(j - i),
-                                        mode="nearest"),
-                        ))
+                            nn.Upsample(scale_factor=2**(j - i), mode="nearest"),
+                        )
+                    )
                 elif j == i:
                     fuse_layer.append(None)
                 else:
@@ -188,7 +182,8 @@ class HighResolutionModule(nn.Module):
                                         bias=False,
                                     ),
                                     nn.BatchNorm2d(num_outchannels_conv3x3),
-                                ))
+                                )
+                            )
                         else:
                             num_outchannels_conv3x3 = num_inchannels[j]
                             conv3x3s.append(
@@ -203,7 +198,8 @@ class HighResolutionModule(nn.Module):
                                     ),
                                     nn.BatchNorm2d(num_outchannels_conv3x3),
                                     nn.ReLU(True),
-                                ))
+                                )
+                            )
                     fuse_layer.append(nn.Sequential(*conv3x3s))
             fuse_layers.append(nn.ModuleList(fuse_layer))
 
@@ -237,7 +233,6 @@ blocks_dict = {"BASIC": BasicBlock, "BOTTLENECK": Bottleneck}
 
 
 class HighResolutionNet(nn.Module):
-
     def __init__(self, cfg, **kwargs):
         self.inplanes = 64
         super(HighResolutionNet, self).__init__()
@@ -245,19 +240,9 @@ class HighResolutionNet(nn.Module):
         self.use_old_impl = use_old_impl
 
         # stem net
-        self.conv1 = nn.Conv2d(3,
-                               64,
-                               kernel_size=3,
-                               stride=2,
-                               padding=1,
-                               bias=False)
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
         self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
-        self.conv2 = nn.Conv2d(64,
-                               64,
-                               kernel_size=3,
-                               stride=2,
-                               padding=1,
-                               bias=False)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
         self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
         self.relu = nn.ReLU(inplace=True)
 
@@ -271,41 +256,29 @@ class HighResolutionNet(nn.Module):
         self.stage2_cfg = cfg.get("stage2", {})
         num_channels = self.stage2_cfg.get("num_channels", (32, 64))
         block = blocks_dict[self.stage2_cfg.get("block")]
-        num_channels = [
-            num_channels[i] * block.expansion for i in range(len(num_channels))
-        ]
+        num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
         stage2_num_channels = num_channels
-        self.transition1 = self._make_transition_layer([stage1_out_channel],
-                                                       num_channels)
-        self.stage2, pre_stage_channels = self._make_stage(
-            self.stage2_cfg, num_channels)
+        self.transition1 = self._make_transition_layer([stage1_out_channel], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(self.stage2_cfg, num_channels)
 
         self.stage3_cfg = cfg.get("stage3")
         num_channels = self.stage3_cfg["num_channels"]
         block = blocks_dict[self.stage3_cfg["block"]]
-        num_channels = [
-            num_channels[i] * block.expansion for i in range(len(num_channels))
-        ]
+        num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
         stage3_num_channels = num_channels
-        self.transition2 = self._make_transition_layer(pre_stage_channels,
-                                                       num_channels)
-        self.stage3, pre_stage_channels = self._make_stage(
-            self.stage3_cfg, num_channels)
+        self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
+        self.stage3, pre_stage_channels = self._make_stage(self.stage3_cfg, num_channels)
 
         self.stage4_cfg = cfg.get("stage4")
         num_channels = self.stage4_cfg["num_channels"]
         block = blocks_dict[self.stage4_cfg["block"]]
-        num_channels = [
-            num_channels[i] * block.expansion for i in range(len(num_channels))
-        ]
-        self.transition3 = self._make_transition_layer(pre_stage_channels,
-                                                       num_channels)
+        num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
+        self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
         stage_4_out_channels = num_channels
 
         self.stage4, pre_stage_channels = self._make_stage(
-            self.stage4_cfg,
-            num_channels,
-            multi_scale_output=not self.use_old_impl)
+            self.stage4_cfg, num_channels, multi_scale_output=not self.use_old_impl
+        )
         stage4_num_channels = num_channels
 
         self.output_channels_dim = pre_stage_channels
@@ -316,35 +289,34 @@ class HighResolutionNet(nn.Module):
         self.avg_pooling = nn.AdaptiveAvgPool2d(1)
 
         if use_old_impl:
-            in_dims = (2**2 * stage2_num_channels[-1] +
-                       2**1 * stage3_num_channels[-1] +
-                       stage_4_out_channels[-1])
+            in_dims = (
+                2**2 * stage2_num_channels[-1] + 2**1 * stage3_num_channels[-1] +
+                stage_4_out_channels[-1]
+            )
         else:
             # TODO: Replace with parameters
             in_dims = 4 * 384
             self.subsample_4 = self._make_subsample_layer(
-                in_channels=stage4_num_channels[0], num_layers=3)
+                in_channels=stage4_num_channels[0], num_layers=3
+            )
 
         self.subsample_3 = self._make_subsample_layer(
-            in_channels=stage2_num_channels[-1], num_layers=2)
+            in_channels=stage2_num_channels[-1], num_layers=2
+        )
         self.subsample_2 = self._make_subsample_layer(
-            in_channels=stage3_num_channels[-1], num_layers=1)
-        self.conv_layers = self._make_conv_layer(in_channels=in_dims,
-                                                 num_layers=5)
+            in_channels=stage3_num_channels[-1], num_layers=1
+        )
+        self.conv_layers = self._make_conv_layer(in_channels=in_dims, num_layers=5)
 
     def get_output_dim(self):
-        base_output = {
-            f"layer{idx + 1}": val
-            for idx, val in enumerate(self.output_channels_dim)
-        }
+        base_output = {f"layer{idx + 1}": val for idx, val in enumerate(self.output_channels_dim)}
         output = base_output.copy()
         for key in base_output:
             output[f"{key}_avg_pooling"] = output[key]
         output["concat"] = 2048
         return output
 
-    def _make_transition_layer(self, num_channels_pre_layer,
-                               num_channels_cur_layer):
+    def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer):
         num_branches_cur = len(num_channels_cur_layer)
         num_branches_pre = len(num_channels_pre_layer)
 
@@ -364,26 +336,24 @@ class HighResolutionNet(nn.Module):
                             ),
                             nn.BatchNorm2d(num_channels_cur_layer[i]),
                             nn.ReLU(inplace=True),
-                        ))
+                        )
+                    )
                 else:
                     transition_layers.append(None)
             else:
                 conv3x3s = []
                 for j in range(i + 1 - num_branches_pre):
                     inchannels = num_channels_pre_layer[-1]
-                    outchannels = (num_channels_cur_layer[i] if j == i -
-                                   num_branches_pre else inchannels)
+                    outchannels = (
+                        num_channels_cur_layer[i] if j == i - num_branches_pre else inchannels
+                    )
                     conv3x3s.append(
                         nn.Sequential(
-                            nn.Conv2d(inchannels,
-                                      outchannels,
-                                      3,
-                                      2,
-                                      1,
-                                      bias=False),
+                            nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False),
                             nn.BatchNorm2d(outchannels),
                             nn.ReLU(inplace=True),
-                        ))
+                        )
+                    )
                 transition_layers.append(nn.Sequential(*conv3x3s))
 
         return nn.ModuleList(transition_layers)
@@ -410,24 +380,13 @@ class HighResolutionNet(nn.Module):
 
         return nn.Sequential(*layers)
 
-    def _make_conv_layer(self,
-                         in_channels=2048,
-                         num_layers=3,
-                         num_filters=2048,
-                         stride=1):
+    def _make_conv_layer(self, in_channels=2048, num_layers=3, num_filters=2048, stride=1):
 
         layers = []
         for i in range(num_layers):
 
-            downsample = nn.Conv2d(in_channels,
-                                   num_filters,
-                                   stride=1,
-                                   kernel_size=1,
-                                   bias=False)
-            layers.append(
-                Bottleneck(in_channels,
-                           num_filters // 4,
-                           downsample=downsample))
+            downsample = nn.Conv2d(in_channels, num_filters, stride=1, kernel_size=1, bias=False)
+            layers.append(Bottleneck(in_channels, num_filters // 4, downsample=downsample))
             in_channels = num_filters
 
         return nn.Sequential(*layers)
@@ -444,18 +403,15 @@ class HighResolutionNet(nn.Module):
                     kernel_size=3,
                     stride=stride,
                     padding=1,
-                ))
+                )
+            )
             in_channels = 2 * in_channels
             layers.append(nn.BatchNorm2d(in_channels, momentum=BN_MOMENTUM))
             layers.append(nn.ReLU(inplace=True))
 
         return nn.Sequential(*layers)
 
-    def _make_stage(self,
-                    layer_config,
-                    num_inchannels,
-                    multi_scale_output=True,
-                    log=False):
+    def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True, log=False):
         num_modules = layer_config["num_modules"]
         num_branches = layer_config["num_branches"]
         num_blocks = layer_config["num_blocks"]
@@ -480,7 +436,8 @@ class HighResolutionNet(nn.Module):
                     num_channels,
                     fuse_method,
                     reset_multi_scale_output,
-                ))
+                )
+            )
             modules[-1].log = log
             num_inchannels = modules[-1].get_num_inchannels()
 
@@ -580,15 +537,14 @@ class HighResolutionNet(nn.Module):
     def load_weights(self, pretrained=""):
         pretrained = osp.expandvars(pretrained)
         if osp.isfile(pretrained):
-            pretrained_state_dict = torch.load(
-                pretrained, map_location=torch.device("cpu"))
+            pretrained_state_dict = torch.load(pretrained, map_location=torch.device("cpu"))
 
             need_init_state_dict = {}
             for name, m in pretrained_state_dict.items():
-                if (name.split(".")[0] in self.pretrained_layers
-                        or self.pretrained_layers[0] == "*"):
+                if (
+                    name.split(".")[0] in self.pretrained_layers or self.pretrained_layers[0] == "*"
+                ):
                     need_init_state_dict[name] = m
-            missing, unexpected = self.load_state_dict(need_init_state_dict,
-                                                       strict=False)
+            missing, unexpected = self.load_state_dict(need_init_state_dict, strict=False)
         elif pretrained:
             raise ValueError("{} is not exist!".format(pretrained))

lib/pixielib/models/lbs.py

@@ -30,8 +30,7 @@ def rot_mat_to_euler(rot_mats):
     # Calculates rotation matrix to euler angles
     # Careful for extreme cases of eular angles like [0.0, pi, 0.0]
 
-    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
-                    rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
+    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] + rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
     return torch.atan2(-rot_mats[:, 2, 0], sy)
 
 
@@ -86,15 +85,13 @@ def find_dynamic_lmk_idx_and_bcoords(
     #     aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)
     rot_mats = torch.index_select(pose, 1, head_kin_chain)
 
-    rel_rot_mat = torch.eye(3, device=vertices.device,
-                            dtype=dtype).unsqueeze_(dim=0)
+    rel_rot_mat = torch.eye(3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0)
     for idx in range(len(head_kin_chain)):
         # rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
         rel_rot_mat = torch.matmul(rot_mats[:, idx], rel_rot_mat)
 
-    y_rot_angle = torch.round(
-        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
-                    max=39)).to(dtype=torch.long)
+    y_rot_angle = torch.round(torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
+                                          max=39)).to(dtype=torch.long)
     # print(y_rot_angle[0])
     neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
     mask = y_rot_angle.lt(-39).to(dtype=torch.long)
@@ -102,8 +99,7 @@ def find_dynamic_lmk_idx_and_bcoords(
     y_rot_angle = neg_mask * neg_vals + (1 - neg_mask) * y_rot_angle
     # print(y_rot_angle[0])
 
-    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx, 0,
-                                           y_rot_angle)
+    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx, 0, y_rot_angle)
     dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords, 0, y_rot_angle)
 
     return dyn_lmk_faces_idx, dyn_lmk_b_coords
@@ -135,11 +131,11 @@ def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
     batch_size, num_verts = vertices.shape[:2]
     device = vertices.device
 
-    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
-        batch_size, -1, 3)
+    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(batch_size, -1, 3)
 
-    lmk_faces += (torch.arange(batch_size, dtype=torch.long,
-                               device=device).view(-1, 1, 1) * num_verts)
+    lmk_faces += (
+        torch.arange(batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts
+    )
 
     lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(batch_size, -1, 3, 3)
 
@@ -211,13 +207,11 @@ def lbs(
     # N x J x 3 x 3
     ident = torch.eye(3, dtype=dtype, device=device)
     if pose2rot:
-        rot_mats = batch_rodrigues(pose.view(-1, 3),
-                                   dtype=dtype).view([batch_size, -1, 3, 3])
+        rot_mats = batch_rodrigues(pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])
 
         pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
         # (N x P) x (P, V * 3) -> N x V x 3
-        pose_offsets = torch.matmul(pose_feature,
-                                    posedirs).view(batch_size, -1, 3)
+        pose_offsets = torch.matmul(pose_feature, posedirs).view(batch_size, -1, 3)
     else:
         pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
         rot_mats = pose.view(batch_size, -1, 3, 3)
@@ -234,12 +228,9 @@ def lbs(
     W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
     # (N x V x (J + 1)) x (N x (J + 1) x 16)
     num_joints = J_regressor.shape[0]
-    T = torch.matmul(W, A.view(batch_size, num_joints,
-                               16)).view(batch_size, -1, 4, 4)
+    T = torch.matmul(W, A.view(batch_size, num_joints, 16)).view(batch_size, -1, 4, 4)
 
-    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
-                               dtype=dtype,
-                               device=device)
+    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1], dtype=dtype, device=device)
     v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
     v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
 
@@ -318,8 +309,7 @@ def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
     K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
 
     zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
-    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros],
-                  dim=1).view((batch_size, 3, 3))
+    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view((batch_size, 3, 3))
 
     ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
     rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
@@ -335,9 +325,7 @@ def transform_mat(R, t):
         - T: Bx4x4 Transformation matrix
     """
     # No padding left or right, only add an extra row
-    return torch.cat([F.pad(R, [0, 0, 0, 1]),
-                      F.pad(t, [0, 0, 0, 1], value=1)],
-                     dim=2)
+    return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
 
 
 def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
@@ -370,15 +358,13 @@ def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
     rel_joints[:, 1:] -= joints[:, parents[1:]]
 
     transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3),
-                                   rel_joints.reshape(-1, 3, 1)).reshape(
-                                       -1, joints.shape[1], 4, 4)
+                                   rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)
 
     transform_chain = [transforms_mat[:, 0]]
     for i in range(1, parents.shape[0]):
         # Subtract the joint location at the rest pose
         # No need for rotation, since it's identity when at rest
-        curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:,
-                                                                            i])
+        curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:, i])
         transform_chain.append(curr_res)
 
     transforms = torch.stack(transform_chain, dim=1)
@@ -392,21 +378,22 @@ def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
     joints_homogen = F.pad(joints, [0, 0, 0, 1])
 
     rel_transforms = transforms - F.pad(
-        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])
+        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]
+    )
 
     return posed_joints, rel_transforms
 
 
 class JointsFromVerticesSelector(nn.Module):
-
     def __init__(self, fname):
         """Selects extra joints from vertices"""
         super(JointsFromVerticesSelector, self).__init__()
 
         err_msg = ("Either pass a filename or triangle face ids, names and"
                    " barycentrics")
-        assert fname is not None or (face_ids is not None and bcs is not None
-                                     and names is not None), err_msg
+        assert fname is not None or (
+            face_ids is not None and bcs is not None and names is not None
+        ), err_msg
         if fname is not None:
             fname = os.path.expanduser(os.path.expandvars(fname))
             with open(fname, "r") as f:
@@ -422,13 +409,11 @@ class JointsFromVerticesSelector(nn.Module):
         assert len(bcs) == len(
             face_ids
         ), "The number of barycentric coordinates must be equal to the faces"
-        assert len(names) == len(
-            face_ids), "The number of names must be equal to the number of "
+        assert len(names) == len(face_ids), "The number of names must be equal to the number of "
 
         self.names = names
         self.register_buffer("bcs", torch.tensor(bcs, dtype=torch.float32))
-        self.register_buffer("face_ids",
-                             torch.tensor(face_ids, dtype=torch.long))
+        self.register_buffer("face_ids", torch.tensor(face_ids, dtype=torch.long))
 
     def extra_joint_names(self):
         """Returns the names of the extra joints"""
@@ -439,8 +424,7 @@ class JointsFromVerticesSelector(nn.Module):
             return []
         vertex_ids = faces[self.face_ids].reshape(-1)
         # Should be BxNx3x3
-        triangles = torch.index_select(vertices, 1, vertex_ids).reshape(
-            -1, len(self.bcs), 3, 3)
+        triangles = torch.index_select(vertices, 1, vertex_ids).reshape(-1, len(self.bcs), 3, 3)
         return (triangles * self.bcs[None, :, :, None]).sum(dim=2)
 
 
@@ -463,7 +447,6 @@ def to_np(array, dtype=np.float32):
 
 
 class Struct(object):
-
     def __init__(self, **kwargs):
         for key, val in kwargs.items():
             setattr(self, key, val)

lib/pixielib/models/moderators.py

@@ -12,11 +12,7 @@ import torch.nn.functional as F
 
 
 class TempSoftmaxFusion(nn.Module):
-
-    def __init__(self,
-                 channels=[2048 * 2, 1024, 1],
-                 detach_inputs=False,
-                 detach_feature=False):
+    def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False):
         super(TempSoftmaxFusion, self).__init__()
         self.detach_inputs = detach_inputs
         self.detach_feature = detach_feature
@@ -63,11 +59,7 @@ class TempSoftmaxFusion(nn.Module):
 
 
 class GumbelSoftmaxFusion(nn.Module):
-
-    def __init__(self,
-                 channels=[2048 * 2, 1024, 1],
-                 detach_inputs=False,
-                 detach_feature=False):
+    def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False):
         super(GumbelSoftmaxFusion, self).__init__()
         self.detach_inputs = detach_inputs
         self.detach_feature = detach_feature

lib/pixielib/models/resnet.py

@@ -22,16 +22,10 @@ from torchvision import models
 
 
 class ResNet(nn.Module):
-
     def __init__(self, block, layers, num_classes=1000):
         self.inplanes = 64
         super(ResNet, self).__init__()
-        self.conv1 = nn.Conv2d(3,
-                               64,
-                               kernel_size=7,
-                               stride=2,
-                               padding=3,
-                               bias=False)
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
@@ -98,12 +92,7 @@ class Bottleneck(nn.Module):
         super(Bottleneck, self).__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
         self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes,
-                               planes,
-                               kernel_size=3,
-                               stride=stride,
-                               padding=1,
-                               bias=False)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
         self.bn2 = nn.BatchNorm2d(planes)
         self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
         self.bn3 = nn.BatchNorm2d(planes * 4)
@@ -136,12 +125,7 @@ class Bottleneck(nn.Module):
 
 def conv3x3(in_planes, out_planes, stride=1):
     """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes,
-                     out_planes,
-                     kernel_size=3,
-                     stride=stride,
-                     padding=1,
-                     bias=False)
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
 
 
 class BasicBlock(nn.Module):
@@ -196,8 +180,7 @@ def load_ResNet50Model():
     model = ResNet(Bottleneck, [3, 4, 6, 3])
     copy_parameter_from_resnet(
         model,
-        torchvision.models.resnet50(
-            weights=models.ResNet50_Weights.DEFAULT).state_dict(),
+        torchvision.models.resnet50(weights=models.ResNet50_Weights.DEFAULT).state_dict(),
     )
     return model
 
@@ -206,8 +189,7 @@ def load_ResNet101Model():
     model = ResNet(Bottleneck, [3, 4, 23, 3])
     copy_parameter_from_resnet(
         model,
-        torchvision.models.resnet101(
-            weights=models.ResNet101_Weights.DEFAULT).state_dict(),
+        torchvision.models.resnet101(weights=models.ResNet101_Weights.DEFAULT).state_dict(),
     )
     return model
 
@@ -216,8 +198,7 @@ def load_ResNet152Model():
     model = ResNet(Bottleneck, [3, 8, 36, 3])
     copy_parameter_from_resnet(
         model,
-        torchvision.models.resnet152(
-            weights=models.ResNet152_Weights.DEFAULT).state_dict(),
+        torchvision.models.resnet152(weights=models.ResNet152_Weights.DEFAULT).state_dict(),
     )
     return model
 
@@ -229,7 +210,6 @@ def load_ResNet152Model():
 
 class DoubleConv(nn.Module):
     """(convolution => [BN] => ReLU) * 2"""
-
     def __init__(self, in_channels, out_channels):
         super().__init__()
         self.double_conv = nn.Sequential(
@@ -247,11 +227,9 @@ class DoubleConv(nn.Module):
 
 class Down(nn.Module):
     """Downscaling with maxpool then double conv"""
-
     def __init__(self, in_channels, out_channels):
         super().__init__()
-        self.maxpool_conv = nn.Sequential(
-            nn.MaxPool2d(2), DoubleConv(in_channels, out_channels))
+        self.maxpool_conv = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_channels, out_channels))
 
     def forward(self, x):
         return self.maxpool_conv(x)
@@ -259,20 +237,16 @@ class Down(nn.Module):
 
 class Up(nn.Module):
     """Upscaling then double conv"""
-
     def __init__(self, in_channels, out_channels, bilinear=True):
         super().__init__()
 
         # if bilinear, use the normal convolutions to reduce the number of channels
         if bilinear:
-            self.up = nn.Upsample(scale_factor=2,
-                                  mode="bilinear",
-                                  align_corners=True)
+            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
         else:
-            self.up = nn.ConvTranspose2d(in_channels // 2,
-                                         in_channels // 2,
-                                         kernel_size=2,
-                                         stride=2)
+            self.up = nn.ConvTranspose2d(
+                in_channels // 2, in_channels // 2, kernel_size=2, stride=2
+            )
 
         self.conv = DoubleConv(in_channels, out_channels)
 
@@ -282,9 +256,7 @@ class Up(nn.Module):
         diffY = x2.size()[2] - x1.size()[2]
         diffX = x2.size()[3] - x1.size()[3]
 
-        x1 = F.pad(
-            x1,
-            [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])
         # if you have padding issues, see
         # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
         # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
@@ -293,7 +265,6 @@ class Up(nn.Module):
 
 
 class OutConv(nn.Module):
-
     def __init__(self, in_channels, out_channels):
         super(OutConv, self).__init__()
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
@@ -303,7 +274,6 @@ class OutConv(nn.Module):
 
 
 class UNet(nn.Module):
-
     def __init__(self, n_channels, n_classes, bilinear=True):
         super(UNet, self).__init__()
         self.n_channels = n_channels

lib/pixielib/pixie.py

@@ -33,7 +33,6 @@ from .utils.config import cfg
 
 
 class PIXIE(object):
-
     def __init__(self, config=None, device="cuda:0"):
         if config is None:
             self.cfg = cfg
@@ -45,10 +44,7 @@ class PIXIE(object):
         self.param_list_dict = {}
         for lst in self.cfg.params.keys():
             param_list = cfg.params.get(lst)
-            self.param_list_dict[lst] = {
-                i: cfg.model.get("n_" + i)
-                for i in param_list
-            }
+            self.param_list_dict[lst] = {i: cfg.model.get("n_" + i) for i in param_list}
 
         # Build the models
         self._create_model()
@@ -97,24 +93,19 @@ class PIXIE(object):
         self.Regressor = {}
         for key in self.cfg.network.regressor.keys():
             n_output = sum(self.param_list_dict[f"{key}_list"].values())
-            channels = ([2048] + self.cfg.network.regressor.get(key).channels +
-                        [n_output])
+            channels = ([2048] + self.cfg.network.regressor.get(key).channels + [n_output])
             if self.cfg.network.regressor.get(key).type == "mlp":
                 self.Regressor[key] = MLP(channels=channels).to(self.device)
-            self.model_dict[f"Regressor_{key}"] = self.Regressor[
-                key].state_dict()
+            self.model_dict[f"Regressor_{key}"] = self.Regressor[key].state_dict()
 
         # Build the extractors
         # to extract separate head/left hand/right hand feature from body feature
         self.Extractor = {}
         for key in self.cfg.network.extractor.keys():
-            channels = [
-                2048
-            ] + self.cfg.network.extractor.get(key).channels + [2048]
+            channels = [2048] + self.cfg.network.extractor.get(key).channels + [2048]
             if self.cfg.network.extractor.get(key).type == "mlp":
                 self.Extractor[key] = MLP(channels=channels).to(self.device)
-            self.model_dict[f"Extractor_{key}"] = self.Extractor[
-                key].state_dict()
+            self.model_dict[f"Extractor_{key}"] = self.Extractor[key].state_dict()
 
         # Build the moderators
         self.Moderator = {}
@@ -122,15 +113,13 @@ class PIXIE(object):
             share_part = key.split("_")[0]
             detach_inputs = self.cfg.network.moderator.get(key).detach_inputs
             detach_feature = self.cfg.network.moderator.get(key).detach_feature
-            channels = [2048 * 2
-                        ] + self.cfg.network.moderator.get(key).channels + [2]
+            channels = [2048 * 2] + self.cfg.network.moderator.get(key).channels + [2]
             self.Moderator[key] = TempSoftmaxFusion(
                 detach_inputs=detach_inputs,
                 detach_feature=detach_feature,
                 channels=channels,
             ).to(self.device)
-            self.model_dict[f"Moderator_{key}"] = self.Moderator[
-                key].state_dict()
+            self.model_dict[f"Moderator_{key}"] = self.Moderator[key].state_dict()
 
         # Build the SMPL-X body model, which we also use to represent faces and
         # hands, using the relevant parts only
@@ -147,9 +136,7 @@ class PIXIE(object):
             print(f"pixie trained model path: {model_path} does not exist!")
             exit()
         # eval mode
-        for module in [
-                self.Encoder, self.Regressor, self.Moderator, self.Extractor
-        ]:
+        for module in [self.Encoder, self.Regressor, self.Moderator, self.Extractor]:
             for net in module.values():
                 net.eval()
 
@@ -185,14 +172,14 @@ class PIXIE(object):
         # crop
         cropper_key = "hand" if "hand" in part_key else part_key
         points_scale = image.shape[-2:]
-        cropped_image, tform = self.Cropper[cropper_key].crop(
-            image, points_for_crop, points_scale)
+        cropped_image, tform = self.Cropper[cropper_key].crop(image, points_for_crop, points_scale)
         # transform points(must be normalized to [-1.1]) accordingly
         cropped_points_dict = {}
         for points_key in points_dict.keys():
             points = points_dict[points_key]
             cropped_points = self.Cropper[cropper_key].transform_points(
-                points, tform, points_scale, normalize=True)
+                points, tform, points_scale, normalize=True
+            )
             cropped_points_dict[points_key] = cropped_points
         return cropped_image, cropped_points_dict
 
@@ -244,8 +231,7 @@ class PIXIE(object):
                 # then predict share parameters
                 feature[key][f"{key}_share"] = feature[key][key]
                 share_dict = self.decompose_code(
-                    self.Regressor[f"{part}_share"](
-                        feature[key][f"{part}_share"]),
+                    self.Regressor[f"{part}_share"](feature[key][f"{part}_share"]),
                     self.param_list_dict[f"{part}_share_list"],
                 )
                 # compose parameters
@@ -257,13 +243,16 @@ class PIXIE(object):
                 f_body = feature["body"]["body"]
                 # extract part feature
                 for part_name in ["head", "left_hand", "right_hand"]:
-                    feature["body"][f"{part_name}_share"] = self.Extractor[
-                        f"{part_name}_share"](f_body)
+                    feature["body"][f"{part_name}_share"] = self.Extractor[f"{part_name}_share"](
+                        f_body
+                    )
 
                 # -- check if part crops are given, if not, crop parts by coarse body estimation
-                if ("head_image" not in data[key].keys()
-                        or "left_hand_image" not in data[key].keys()
-                        or "right_hand_image" not in data[key].keys()):
+                if (
+                    "head_image" not in data[key].keys() or
+                    "left_hand_image" not in data[key].keys() or
+                    "right_hand_image" not in data[key].keys()
+                ):
                     # - run without fusion to get coarse estimation, for cropping parts
                     # body only
                     body_dict = self.decompose_code(
@@ -272,29 +261,26 @@ class PIXIE(object):
                     )
                     # head share
                     head_share_dict = self.decompose_code(
-                        self.Regressor["head" + "_share"](
-                            feature[key]["head" + "_share"]),
+                        self.Regressor["head" + "_share"](feature[key]["head" + "_share"]),
                         self.param_list_dict["head" + "_share_list"],
                     )
                     # right hand share
                     right_hand_share_dict = self.decompose_code(
-                        self.Regressor["hand" + "_share"](
-                            feature[key]["right_hand" + "_share"]),
+                        self.Regressor["hand" + "_share"](feature[key]["right_hand" + "_share"]),
                         self.param_list_dict["hand" + "_share_list"],
                     )
                     # left hand share
                     left_hand_share_dict = self.decompose_code(
-                        self.Regressor["hand" + "_share"](
-                            feature[key]["left_hand" + "_share"]),
+                        self.Regressor["hand" + "_share"](feature[key]["left_hand" + "_share"]),
                         self.param_list_dict["hand" + "_share_list"],
                     )
                     # change the dict name from right to left
-                    left_hand_share_dict[
-                        "left_hand_pose"] = left_hand_share_dict.pop(
-                            "right_hand_pose")
-                    left_hand_share_dict[
-                        "left_wrist_pose"] = left_hand_share_dict.pop(
-                            "right_wrist_pose")
+                    left_hand_share_dict["left_hand_pose"] = left_hand_share_dict.pop(
+                        "right_hand_pose"
+                    )
+                    left_hand_share_dict["left_wrist_pose"] = left_hand_share_dict.pop(
+                        "right_wrist_pose"
+                    )
                     param_dict[key] = {
                         **body_dict,
                         **head_share_dict,
@@ -304,21 +290,18 @@ class PIXIE(object):
                     if body_only:
                         param_dict["moderator_weight"] = None
                         return param_dict
-                    prediction_body_only = self.decode(param_dict[key],
-                                                       param_type="body")
+                    prediction_body_only = self.decode(param_dict[key], param_type="body")
                     # crop
                     for part_name in ["head", "left_hand", "right_hand"]:
                         part = part_name.split("_")[-1]
                         points_dict = {
-                            "smplx_kpt":
-                            prediction_body_only["smplx_kpt"],
-                            "trans_verts":
-                            prediction_body_only["transformed_vertices"],
+                            "smplx_kpt": prediction_body_only["smplx_kpt"],
+                            "trans_verts": prediction_body_only["transformed_vertices"],
                         }
-                        image_hd = torchvision.transforms.Resize(1024)(
-                            data["body"]["image"])
+                        image_hd = torchvision.transforms.Resize(1024)(data["body"]["image"])
                         cropped_image, cropped_joints_dict = self.part_from_body(
-                            image_hd, part_name, points_dict)
+                            image_hd, part_name, points_dict
+                        )
                         data[key][part_name + "_image"] = cropped_image
 
                 # -- encode features from part crops, then fuse feature using the weight from moderator
@@ -338,16 +321,12 @@ class PIXIE(object):
                         self.Regressor[f"{part}_share"](f_part),
                         self.param_list_dict[f"{part}_share_list"],
                     )
-                    param_dict["body_" + part_name] = {
-                        **part_dict,
-                        **part_share_dict
-                    }
+                    param_dict["body_" + part_name] = {**part_dict, **part_share_dict}
 
                     # moderator to assign weight, then integrate features
-                    f_body_out, f_part_out, f_weight = self.Moderator[
-                        f"{part}_share"](feature["body"][f"{part_name}_share"],
-                                         f_part,
-                                         work=True)
+                    f_body_out, f_part_out, f_weight = self.Moderator[f"{part}_share"](
+                        feature["body"][f"{part_name}_share"], f_part, work=True
+                    )
                     if copy_and_paste:
                         # copy and paste strategy always trusts the results from part
                         feature["body"][f"{part_name}_share"] = f_part
@@ -355,8 +334,9 @@ class PIXIE(object):
                         # for hand, if part weight > 0.7 (very confident, then fully trust part)
                         part_w = f_weight[:, [1]]
                         part_w[part_w > 0.7] = 1.0
-                        f_body_out = (feature["body"][f"{part_name}_share"] *
-                                      (1.0 - part_w) + f_part * part_w)
+                        f_body_out = (
+                            feature["body"][f"{part_name}_share"] * (1.0 - part_w) + f_part * part_w
+                        )
                         feature["body"][f"{part_name}_share"] = f_body_out
                     else:
                         feature["body"][f"{part_name}_share"] = f_body_out
@@ -367,29 +347,24 @@ class PIXIE(object):
                 # -- predict parameters from fused body feature
                 # head share
                 head_share_dict = self.decompose_code(
-                    self.Regressor["head" + "_share"](feature[key]["head" +
-                                                                   "_share"]),
+                    self.Regressor["head" + "_share"](feature[key]["head" + "_share"]),
                     self.param_list_dict["head" + "_share_list"],
                 )
                 # right hand share
                 right_hand_share_dict = self.decompose_code(
-                    self.Regressor["hand" + "_share"](
-                        feature[key]["right_hand" + "_share"]),
+                    self.Regressor["hand" + "_share"](feature[key]["right_hand" + "_share"]),
                     self.param_list_dict["hand" + "_share_list"],
                 )
                 # left hand share
                 left_hand_share_dict = self.decompose_code(
-                    self.Regressor["hand" + "_share"](
-                        feature[key]["left_hand" + "_share"]),
+                    self.Regressor["hand" + "_share"](feature[key]["left_hand" + "_share"]),
                     self.param_list_dict["hand" + "_share_list"],
                 )
                 # change the dict name from right to left
-                left_hand_share_dict[
-                    "left_hand_pose"] = left_hand_share_dict.pop(
-                        "right_hand_pose")
-                left_hand_share_dict[
-                    "left_wrist_pose"] = left_hand_share_dict.pop(
-                        "right_wrist_pose")
+                left_hand_share_dict["left_hand_pose"] = left_hand_share_dict.pop("right_hand_pose")
+                left_hand_share_dict["left_wrist_pose"] = left_hand_share_dict.pop(
+                    "right_wrist_pose"
+                )
                 param_dict["body"] = {
                     **body_dict,
                     **head_share_dict,
@@ -403,10 +378,10 @@ class PIXIE(object):
                 if keep_local:
                     # for local change that will not affect whole body and produce unnatral pose, trust part
                     param_dict[key]["exp"] = param_dict["body_head"]["exp"]
-                    param_dict[key]["right_hand_pose"] = param_dict[
-                        "body_right_hand"]["right_hand_pose"]
-                    param_dict[key]["left_hand_pose"] = param_dict[
-                        "body_left_hand"]["right_hand_pose"]
+                    param_dict[key]["right_hand_pose"] = param_dict["body_right_hand"][
+                        "right_hand_pose"]
+                    param_dict[key]["left_hand_pose"] = param_dict["body_left_hand"][
+                        "right_hand_pose"]
 
         return param_dict
 
@@ -426,75 +401,70 @@ class PIXIE(object):
             if "pose" in key and "jaw" not in key:
                 param_dict[key] = converter.batch_cont2matrix(param_dict[key])
         if param_type == "body" or param_type == "head":
-            param_dict["jaw_pose"] = converter.batch_euler2matrix(
-                param_dict["jaw_pose"])[:, None, :, :]
+            param_dict["jaw_pose"] = converter.batch_euler2matrix(param_dict["jaw_pose"]
+                                                                 )[:, None, :, :]
 
         # complement params if it's not in given param dict
         if param_type == "head":
             batch_size = param_dict["shape"].shape[0]
             param_dict["abs_head_pose"] = param_dict["head_pose"].clone()
             param_dict["global_pose"] = param_dict["head_pose"]
-            param_dict["partbody_pose"] = self.smplx.body_pose.unsqueeze(
-                0).expand(
-                    batch_size, -1, -1,
-                    -1)[:, :self.param_list_dict["body_list"]["partbody_pose"]]
+            param_dict["partbody_pose"] = self.smplx.body_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )[:, :self.param_list_dict["body_list"]["partbody_pose"]]
             param_dict["neck_pose"] = self.smplx.neck_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
-            param_dict["left_wrist_pose"] = self.smplx.neck_pose.unsqueeze(
-                0).expand(batch_size, -1, -1, -1)
-            param_dict["left_hand_pose"] = self.smplx.left_hand_pose.unsqueeze(
-                0).expand(batch_size, -1, -1, -1)
-            param_dict["right_wrist_pose"] = self.smplx.neck_pose.unsqueeze(
-                0).expand(batch_size, -1, -1, -1)
-            param_dict[
-                "right_hand_pose"] = self.smplx.right_hand_pose.unsqueeze(
-                    0).expand(batch_size, -1, -1, -1)
+                batch_size, -1, -1, -1
+            )
+            param_dict["left_wrist_pose"] = self.smplx.neck_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
+            param_dict["left_hand_pose"] = self.smplx.left_hand_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
+            param_dict["right_wrist_pose"] = self.smplx.neck_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
+            param_dict["right_hand_pose"] = self.smplx.right_hand_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
         elif param_type == "hand":
             batch_size = param_dict["right_hand_pose"].shape[0]
-            param_dict["abs_right_wrist_pose"] = param_dict[
-                "right_wrist_pose"].clone()
+            param_dict["abs_right_wrist_pose"] = param_dict["right_wrist_pose"].clone()
             dtype = param_dict["right_hand_pose"].dtype
             device = param_dict["right_hand_pose"].device
-            x_180_pose = (torch.eye(3, dtype=dtype,
-                                    device=device).unsqueeze(0).repeat(
-                                        1, 1, 1))
+            x_180_pose = (torch.eye(3, dtype=dtype, device=device).unsqueeze(0).repeat(1, 1, 1))
             x_180_pose[0, 2, 2] = -1.0
             x_180_pose[0, 1, 1] = -1.0
-            param_dict["global_pose"] = x_180_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
-            param_dict["shape"] = self.smplx.shape_params.expand(
-                batch_size, -1)
-            param_dict["exp"] = self.smplx.expression_params.expand(
-                batch_size, -1)
+            param_dict["global_pose"] = x_180_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
+            param_dict["shape"] = self.smplx.shape_params.expand(batch_size, -1)
+            param_dict["exp"] = self.smplx.expression_params.expand(batch_size, -1)
             param_dict["head_pose"] = self.smplx.head_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
+                batch_size, -1, -1, -1
+            )
             param_dict["neck_pose"] = self.smplx.neck_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
-            param_dict["jaw_pose"] = self.smplx.jaw_pose.unsqueeze(0).expand(
-                batch_size, -1, -1, -1)
-            param_dict["partbody_pose"] = self.smplx.body_pose.unsqueeze(
-                0).expand(
-                    batch_size, -1, -1,
-                    -1)[:, :self.param_list_dict["body_list"]["partbody_pose"]]
-            param_dict["left_wrist_pose"] = self.smplx.neck_pose.unsqueeze(
-                0).expand(batch_size, -1, -1, -1)
-            param_dict["left_hand_pose"] = self.smplx.left_hand_pose.unsqueeze(
-                0).expand(batch_size, -1, -1, -1)
+                batch_size, -1, -1, -1
+            )
+            param_dict["jaw_pose"] = self.smplx.jaw_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
+            param_dict["partbody_pose"] = self.smplx.body_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )[:, :self.param_list_dict["body_list"]["partbody_pose"]]
+            param_dict["left_wrist_pose"] = self.smplx.neck_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
+            param_dict["left_hand_pose"] = self.smplx.left_hand_pose.unsqueeze(0).expand(
+                batch_size, -1, -1, -1
+            )
         elif param_type == "body":
             # the predcition from the head and hand share regressor is always absolute pose
             batch_size = param_dict["shape"].shape[0]
             param_dict["abs_head_pose"] = param_dict["head_pose"].clone()
-            param_dict["abs_right_wrist_pose"] = param_dict[
-                "right_wrist_pose"].clone()
-            param_dict["abs_left_wrist_pose"] = param_dict[
-                "left_wrist_pose"].clone()
+            param_dict["abs_right_wrist_pose"] = param_dict["right_wrist_pose"].clone()
+            param_dict["abs_left_wrist_pose"] = param_dict["left_wrist_pose"].clone()
             # the body-hand share regressor is working for right hand
             # so we assume body network get the flipped feature for the left hand. then get the parameters
             # then we need to flip it back to left, which matches the input left hand
-            param_dict["left_wrist_pose"] = util.flip_pose(
-                param_dict["left_wrist_pose"])
-            param_dict["left_hand_pose"] = util.flip_pose(
-                param_dict["left_hand_pose"])
+            param_dict["left_wrist_pose"] = util.flip_pose(param_dict["left_wrist_pose"])
+            param_dict["left_hand_pose"] = util.flip_pose(param_dict["left_hand_pose"])
         else:
             exit()
 
@@ -508,8 +478,7 @@ class PIXIE(object):
         Returns:
             predictions: smplx predictions
         """
-        if "jaw_pose" in param_dict.keys() and len(
-                param_dict["jaw_pose"].shape) == 2:
+        if "jaw_pose" in param_dict.keys() and len(param_dict["jaw_pose"].shape) == 2:
             self.convert_pose(param_dict, param_type)
         elif param_dict["right_wrist_pose"].shape[-1] == 6:
             self.convert_pose(param_dict, param_type)
@@ -532,9 +501,8 @@ class PIXIE(object):
         # change absolute head&hand pose to relative pose according to rest body pose
         if param_type == "head" or param_type == "body":
             param_dict["body_pose"] = self.smplx.pose_abs2rel(
-                param_dict["global_pose"],
-                param_dict["body_pose"],
-                abs_joint="head")
+                param_dict["global_pose"], param_dict["body_pose"], abs_joint="head"
+            )
         if param_type == "hand" or param_type == "body":
             param_dict["body_pose"] = self.smplx.pose_abs2rel(
                 param_dict["global_pose"],
@@ -550,7 +518,7 @@ class PIXIE(object):
         if self.cfg.model.check_pose:
             # check if pose is natural (relative rotation), if not, set relative to 0 (especially for head pose)
             # xyz: pitch(positive for looking down), yaw(positive for looking left), roll(rolling chin to left)
-            for pose_ind in [14]:  # head [15-1, 20-1, 21-1]:
+            for pose_ind in [14]:    # head [15-1, 20-1, 21-1]:
                 curr_pose = param_dict["body_pose"][:, pose_ind]
                 euler_pose = converter._compute_euler_from_matrix(curr_pose)
                 for i, max_angle in enumerate([20, 70, 10]):
@@ -560,9 +528,7 @@ class PIXIE(object):
                         min=-max_angle * np.pi / 180,
                         max=max_angle * np.pi / 180,
                     )] = 0.0
-                param_dict[
-                    "body_pose"][:, pose_ind] = converter.batch_euler2matrix(
-                        euler_pose)
+                param_dict["body_pose"][:, pose_ind] = converter.batch_euler2matrix(euler_pose)
 
         # SMPLX
         verts, landmarks, joints = self.smplx(
@@ -594,8 +560,8 @@ class PIXIE(object):
 
         # change the order of face keypoints, to be the same as "standard" 68 keypoints
         prediction["face_kpt"] = torch.cat(
-            [prediction["face_kpt"][:, -17:], prediction["face_kpt"][:, :-17]],
-            dim=1)
+            [prediction["face_kpt"][:, -17:], prediction["face_kpt"][:, :-17]], dim=1
+        )
 
         prediction.update(param_dict)
 

lib/pixielib/utils/array_cropper.py

@@ -23,15 +23,14 @@ def points2bbox(points, points_scale=None):
     bottom = np.max(points[:, 1])
     size = max(right - left, bottom - top)
     # + old_size*0.1])
-    center = np.array(
-        [right - (right - left) / 2.0, bottom - (bottom - top) / 2.0])
+    center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0])
     return center, size
     # translate center
 
 
 def augment_bbox(center, bbox_size, scale=[1.0, 1.0], trans_scale=0.0):
     trans_scale = (np.random.rand(2) * 2 - 1) * trans_scale
-    center = center + trans_scale * bbox_size  # 0.5
+    center = center + trans_scale * bbox_size    # 0.5
     scale = np.random.rand() * (scale[1] - scale[0]) + scale[0]
     size = int(bbox_size * scale)
     return center, size
@@ -48,27 +47,25 @@ def crop_array(image, center, bboxsize, crop_size):
         tform: 3x3 affine matrix
     """
     # points: top-left, top-right, bottom-right
-    src_pts = np.array([
-        [center[0] - bboxsize / 2, center[1] - bboxsize / 2],
-        [center[0] + bboxsize / 2, center[1] - bboxsize / 2],
-        [center[0] + bboxsize / 2, center[1] + bboxsize / 2],
-    ])
-    DST_PTS = np.array([[0, 0], [crop_size - 1, 0],
-                        [crop_size - 1, crop_size - 1]])
+    src_pts = np.array(
+        [
+            [center[0] - bboxsize / 2, center[1] - bboxsize / 2],
+            [center[0] + bboxsize / 2, center[1] - bboxsize / 2],
+            [center[0] + bboxsize / 2, center[1] + bboxsize / 2],
+        ]
+    )
+    DST_PTS = np.array([[0, 0], [crop_size - 1, 0], [crop_size - 1, crop_size - 1]])
 
     # estimate transformation between points
     tform = estimate_transform("similarity", src_pts, DST_PTS)
 
     # warp images
-    cropped_image = warp(image,
-                         tform.inverse,
-                         output_shape=(crop_size, crop_size))
+    cropped_image = warp(image, tform.inverse, output_shape=(crop_size, crop_size))
 
     return cropped_image, tform.params.T
 
 
 class Cropper(object):
-
     def __init__(self, crop_size, scale=[1, 1], trans_scale=0.0):
         self.crop_size = crop_size
         self.scale = scale
@@ -78,11 +75,9 @@ class Cropper(object):
         # points to bbox
         center, bbox_size = points2bbox(points, points_scale)
         # argument bbox.
-        center, bbox_size = augment_bbox(center,
-                                         bbox_size,
-                                         scale=self.scale,
-                                         trans_scale=self.trans_scale)
+        center, bbox_size = augment_bbox(
+            center, bbox_size, scale=self.scale, trans_scale=self.trans_scale
+        )
         # crop
-        cropped_image, tform = crop_array(image, center, bbox_size,
-                                          self.crop_size)
+        cropped_image, tform = crop_array(image, center, bbox_size, self.crop_size)
         return cropped_image, tform