asa
unknown
python
a year ago
4.3 kB
6
Indexable
import torch import trimesh from scipy.spatial.transform import Rotation as R import numpy as np 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/2) # Debug: Print shapes and some vertices for inspection print("F_verts shape:", F_verts.shape) print("B_verts shape:", B_verts.shape) print("Sample F_verts:", F_verts[:5]) print("Sample B_verts:", B_verts[:5]) # Rotate the back vertices 180 degrees around the y-axis rotation_matrix_y = R.from_euler('y', 180, degrees=True).as_matrix() B_verts_rotated = np.dot(B_verts, rotation_matrix_y.T) # Debug: Print rotated vertices print("Sample B_verts_rotated:", B_verts_rotated[:5]) # Additional rotation if needed rotation_matrix_x = R.from_euler('x', -25, degrees=True).as_matrix() B_verts_rotated_x = np.dot(B_verts_rotated, rotation_matrix_x.T) # Debug: Print rotated vertices after second rotation print("Sample B_verts_rotated_x:", B_verts_rotated_x[:5]) F_B_verts = torch.cat((F_verts, torch.tensor(B_verts_rotated_x)), 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( torch.tensor(B_verts_rotated_x).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} 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"))
Editor is loading...
Leave a Comment