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