Spaces:
Runtime error
Runtime error
| import torch | |
| import trimesh | |
| from lib.common.BNI_utils import ( | |
| depth_inverse_transform, | |
| double_side_bilateral_normal_integration, | |
| verts_inverse_transform, | |
| ) | |
| class BNI: | |
| def __init__(self, dir_path, name, BNI_dict, cfg, device): | |
| self.scale = 256.0 | |
| self.cfg = cfg | |
| self.name = name | |
| self.normal_front = BNI_dict["normal_F"] | |
| self.normal_back = BNI_dict["normal_B"] | |
| self.mask = BNI_dict["mask"] | |
| self.depth_front = BNI_dict["depth_F"] | |
| self.depth_back = BNI_dict["depth_B"] | |
| self.depth_mask = BNI_dict["depth_mask"] | |
| # hparam: | |
| # k --> smaller, keep continuity | |
| # lambda --> larger, more depth-awareness | |
| self.k = self.cfg['k'] | |
| self.lambda1 = self.cfg['lambda1'] | |
| self.boundary_consist = self.cfg['boundary_consist'] | |
| self.cut_intersection = self.cfg['cut_intersection'] | |
| self.F_B_surface = None | |
| self.F_B_trimesh = None | |
| self.F_depth = None | |
| self.B_depth = None | |
| self.device = device | |
| self.export_dir = dir_path | |
| # code: https://github.com/hoshino042/bilateral_normal_integration | |
| # paper: Bilateral Normal Integration | |
| def extract_surface(self, verbose=True): | |
| bni_result = double_side_bilateral_normal_integration( | |
| normal_front=self.normal_front, | |
| normal_back=self.normal_back, | |
| normal_mask=self.mask, | |
| depth_front=self.depth_front * self.scale, | |
| depth_back=self.depth_back * self.scale, | |
| depth_mask=self.depth_mask, | |
| k=self.k, | |
| lambda_normal_back=1.0, | |
| lambda_depth_front=self.lambda1, | |
| lambda_depth_back=self.lambda1, | |
| lambda_boundary_consistency=self.boundary_consist, | |
| cut_intersection=self.cut_intersection, | |
| ) | |
| F_verts = verts_inverse_transform(bni_result["F_verts"], self.scale) | |
| B_verts = verts_inverse_transform(bni_result["B_verts"], self.scale) | |
| self.F_depth = depth_inverse_transform(bni_result["F_depth"], self.scale) | |
| self.B_depth = depth_inverse_transform(bni_result["B_depth"], self.scale) | |
| 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 | |
| ) | |
| 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.B_trimesh = trimesh.Trimesh( | |
| # B_verts.float(), bni_result["B_faces"].long(), process=False, maintain_order=True | |
| # ) | |
| if __name__ == "__main__": | |
| import os.path as osp | |
| import numpy as np | |
| from tqdm import tqdm | |
| root = "/home/yxiu/Code/ECON/results/examples/BNI" | |
| npy_file = f"{root}/304e9c4798a8c3967de7c74c24ef2e38.npy" | |
| 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 | |
| # 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') | |
| ) | |
| bni_object.extract_surface() | |
| bni_object.F_trimesh.export(osp.join(osp.dirname(npy_file), "F.obj")) | |
| bni_object.B_trimesh.export(osp.join(osp.dirname(npy_file), "B.obj")) | |
| bni_object.F_B_trimesh.export(osp.join(osp.dirname(npy_file), "BNI.obj")) | |