Spaces:
Runtime error
Runtime error
| # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # | |
| # NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. | |
| from pdb import set_trace as st | |
| import json | |
| import math | |
| import pytorch3d.ops | |
| import numpy as np | |
| import os | |
| import argparse | |
| import multiprocessing as mp | |
| from multiprocessing import Pool | |
| import trimesh | |
| import tqdm | |
| import torch | |
| import nvdiffrast.torch as dr | |
| import kaolin as kal | |
| import glob | |
| import ipdb | |
| parser = argparse.ArgumentParser(description='sample surface points from mesh') | |
| parser.add_argument( | |
| '--n_proc', type=int, default=8, | |
| help='Number of processes to run in parallel' | |
| '(0 means sequential execution).') | |
| parser.add_argument( | |
| '--n_points', type=int, default=4096, | |
| help='Number of points to sample per model.') | |
| parser.add_argument( | |
| '--n_views', type=int, default=100, | |
| help='Number of views per model.') | |
| parser.add_argument( | |
| '--image_height', type=int, default=640, | |
| help='Depth image height.') | |
| parser.add_argument( | |
| '--image_width', type=int, default=640, | |
| help='Depth image width.') | |
| parser.add_argument( | |
| '--focal_length_x', type=float, default=640, | |
| help='Focal length in x direction.') | |
| parser.add_argument( | |
| '--focal_length_y', type=float, default=640, | |
| help='Focal length in y direction.') | |
| parser.add_argument( | |
| '--principal_point_x', type=float, default=320, | |
| help='Principal point location in x direction.') | |
| parser.add_argument( | |
| '--principal_point_y', type=float, default=320, | |
| help='Principal point location in y direction.') | |
| parser.add_argument("--shape_root", type=str, default='/mnt/petrelfs/caoziang/3D_generation/Checkpoint_all/diffusion_shapenet_testmodel27_omni_ablation2/ddpm_5000/test', help="path to the save resules shapenet dataset") | |
| parser.add_argument("--save_root", type=str, default='/mnt/petrelfs/caoziang/3D_generation/Checkpoint_all/diffusion_shapenet_testmodel27_omni_ablation2/ddpm_vis_ab2surface', help="path to the split shapenet dataset") | |
| options = parser.parse_args() | |
| # create array for inverse mapping | |
| coordspx2 = np.stack(np.nonzero(np.ones((options.image_height, options.image_width))), -1).astype(np.float32) | |
| coordspx2 = coordspx2[:, ::-1] | |
| fusion_intrisics = np.array( | |
| [ | |
| [options.focal_length_x, 0, options.principal_point_x], | |
| [0, options.focal_length_y, options.principal_point_y], | |
| [0, 0, 1] | |
| ]) | |
| # glctx = dr.RasterizeGLContext() # EGL/egl.h: No such file or directory | |
| glctx = dr.RasterizeCudaContext() | |
| def CalcLinearZ(depth): | |
| # depth = depth * 2 - 1 | |
| zFar = 100.0 | |
| zNear = 0.1 | |
| linear = zNear / (zFar - depth * (zFar - zNear)) * zFar | |
| return linear | |
| def projection_cv_new(fx, fy, cx, cy, width, height, n=1.0, f=50.0): | |
| return np.array( | |
| [[-2 * fx / width, 0.0, (width - 2 * cx) / width, 0.0], | |
| [0.0, -2 * fy / height, (height - 2 * cy) / height, 0.0], | |
| [0.0, 0.0, (-f - n) / (f - n), -2.0 * f * n / (f - n)], | |
| [0.0, 0.0, -1.0, 0.0]]) | |
| def interpolate(attr, rast, attr_idx, rast_db=None): | |
| return dr.interpolate( | |
| attr.contiguous(), rast, attr_idx, rast_db=rast_db, | |
| diff_attrs=None if rast_db is None else 'all') | |
| def render_nvdiffrast(v_pos, tris, T_bx4x4): | |
| # T_bx4x4 - world to cam | |
| proj = projection_cv_new( | |
| fx=options.focal_length_x, fy=options.focal_length_y, cx=options.principal_point_x, | |
| cy=options.principal_point_y, | |
| width=options.image_width, height=options.image_height, n=0.1, f=100.0) | |
| fix = torch.eye(4, dtype=torch.float32, device='cuda') | |
| fix[2, 2] = -1 | |
| fix[1, 1] = -1 | |
| fix[0, 0] = -1 | |
| fix = fix.unsqueeze(0).repeat(T_bx4x4.shape[0], 1, 1) | |
| proj = torch.tensor(proj, dtype=torch.float32, device='cuda').unsqueeze(0).repeat(T_bx4x4.shape[0], 1, 1) | |
| T_world_cam_bx4x4 = torch.bmm(fix, T_bx4x4) | |
| mvp = torch.bmm(proj, T_world_cam_bx4x4) | |
| v_pos_clip = torch.matmul( | |
| torch.nn.functional.pad(v_pos, pad=(0, 1), mode='constant', value=1.0), | |
| torch.transpose(mvp, 1, 2)) | |
| rast, db = dr.rasterize( | |
| glctx, torch.tensor(v_pos_clip, dtype=torch.float32, device='cuda'), tris.int(), | |
| (options.image_height, options.image_width)) | |
| v_pos_cam = torch.matmul( | |
| torch.nn.functional.pad(v_pos, pad=(0, 1), mode='constant', value=1.0), | |
| torch.transpose(T_world_cam_bx4x4, 1, 2)) | |
| gb_pos_cam, _ = interpolate(v_pos_cam, rast, tris.int()) | |
| depth_maps = gb_pos_cam[..., 2].abs() | |
| return depth_maps | |
| def as_mesh(scene_or_mesh): | |
| """ | |
| Convert a possible scene to a mesh. | |
| If conversion occurs, the returned mesh has only vertex and face data. | |
| """ | |
| if isinstance(scene_or_mesh, trimesh.Scene): | |
| if len(scene_or_mesh.geometry) == 0: | |
| mesh = None # empty scene | |
| else: | |
| # we lose texture information here | |
| mesh = trimesh.util.concatenate( | |
| tuple( | |
| trimesh.Trimesh(vertices=g.vertices, faces=g.faces) | |
| for g in scene_or_mesh.geometry.values())) | |
| else: | |
| assert (isinstance(scene_or_mesh, trimesh.Trimesh)) | |
| mesh = scene_or_mesh | |
| return mesh | |
| def render(mesh_v, mesh_f, Rs): | |
| """ | |
| Render the given mesh using the generated views. | |
| :param base_mesh: mesh to render | |
| :type base_mesh: mesh.Mesh | |
| :param Rs: rotation matrices | |
| :type Rs: [numpy.ndarray] | |
| :return: depth maps | |
| :rtype: numpy.ndarray | |
| """ | |
| T_bx4x4 = torch.zeros((options.n_views, 4, 4), dtype=torch.float32, device='cuda') | |
| T_bx4x4[:, 3, 3] = 1 | |
| T_bx4x4[:, 2, 3] = 1 | |
| T_bx4x4[:, :3, :3] = torch.tensor(Rs, dtype=torch.float32, device='cuda') | |
| depthmaps = render_nvdiffrast( | |
| mesh_v, | |
| mesh_f, T_bx4x4) | |
| return depthmaps | |
| def get_points(): | |
| """ | |
| :param n_points: number of points | |
| :type n_points: int | |
| :return: list of points | |
| :rtype: numpy.ndarray | |
| """ | |
| rnd = 1. | |
| points = [] | |
| offset = 2. / options.n_views | |
| increment = math.pi * (3. - math.sqrt(5.)) | |
| for i in range(options.n_views): | |
| y = ((i * offset) - 1) + (offset / 2) | |
| r = math.sqrt(1 - pow(y, 2)) | |
| phi = ((i + rnd) % options.n_views) * increment | |
| x = math.cos(phi) * r | |
| z = math.sin(phi) * r | |
| points.append([x, y, z]) | |
| return np.array(points) | |
| def get_views(semi_sphere=False): | |
| """ | |
| Generate a set of views to generate depth maps from. | |
| :param n_views: number of views per axis | |
| :type n_views: int | |
| :return: rotation matrices | |
| :rtype: [numpy.ndarray] | |
| """ | |
| Rs = [] | |
| points = get_points() | |
| if semi_sphere: | |
| points[:, 2] = -np.abs(points[:, 2]) - 0.1 | |
| for i in range(points.shape[0]): | |
| longitude = - math.atan2(points[i, 0], points[i, 1]) | |
| latitude = math.atan2(points[i, 2], math.sqrt(points[i, 0] ** 2 + points[i, 1] ** 2)) | |
| R_x = np.array( | |
| [[1, 0, 0], | |
| [0, math.cos(latitude), -math.sin(latitude)], | |
| [0, math.sin(latitude), math.cos(latitude)]]) | |
| R_y = np.array( | |
| [[math.cos(longitude), 0, math.sin(longitude)], | |
| [0, 1, 0], | |
| [-math.sin(longitude), 0, math.cos(longitude)]]) | |
| R = R_x @ R_y | |
| Rs.append(R) | |
| return Rs | |
| def fusion(depthmaps, Rs): | |
| """ | |
| Fuse the rendered depth maps. | |
| :param depthmaps: depth maps | |
| :type depthmaps: numpy.ndarray | |
| :param Rs: rotation matrices corresponding to views | |
| :type Rs: [numpy.ndarray] | |
| :return: (T)SDF | |
| :rtype: numpy.ndarray | |
| """ | |
| # sample points inside mask | |
| sample_per_view = options.n_points // options.n_views | |
| sample_bxn = torch.zeros((options.n_views, sample_per_view), device='cuda', dtype=torch.long) | |
| for i in range(len(Rs)): | |
| mask = depthmaps[i] > 0 | |
| valid_idx = torch.nonzero(mask.reshape(-1)).squeeze(-1) | |
| idx = list(range(valid_idx.shape[0])) | |
| np.random.shuffle(idx) | |
| idx = idx[:sample_per_view] | |
| sample_bxn[i] = torch.tensor(valid_idx[idx]) | |
| depthmaps = torch.gather(depthmaps.reshape(options.n_views, -1), 1, sample_bxn) | |
| inv_Ks_bx3x3 = torch.tensor(np.linalg.inv(fusion_intrisics), dtype=torch.float32, device='cuda').unsqueeze( | |
| 0).repeat(options.n_views, 1, 1) | |
| T_bx4x4 = torch.zeros((options.n_views, 4, 4), dtype=torch.float32, device='cuda') | |
| T_bx4x4[:, 3, 3] = 1 | |
| T_bx4x4[:, 2, 3] = 1 | |
| T_bx4x4[:, :3, :3] = torch.tensor(Rs, dtype=torch.float32, device='cuda') | |
| inv_T_bx4x4 = torch.inverse(T_bx4x4) | |
| tf_coords_bxpx2 = torch.tensor(coordspx2.copy(), dtype=torch.float32, device='cuda').unsqueeze(0).repeat( | |
| options.n_views, 1, 1) | |
| tf_coords_bxpx2 = torch.gather(tf_coords_bxpx2, 1, sample_bxn.unsqueeze(-1).repeat(1, 1, 2)) | |
| tf_coords_bxpx3 = torch.cat([tf_coords_bxpx2, torch.ones_like(tf_coords_bxpx2[..., :1])], -1) | |
| tf_coords_bxpx3 *= depthmaps.reshape(options.n_views, -1, 1) | |
| tf_cam_bxpx3 = torch.bmm(inv_Ks_bx3x3, tf_coords_bxpx3.transpose(1, 2)).transpose(1, 2) | |
| tf_cam_bxpx4 = torch.cat([tf_cam_bxpx3, torch.ones_like(tf_cam_bxpx3[..., :1])], -1) | |
| tf_world_bxpx3 = torch.bmm(inv_T_bx4x4, tf_cam_bxpx4.transpose(1, 2)).transpose(1, 2)[..., :3] | |
| return tf_world_bxpx3.reshape(-1, 3) | |
| def normalize(vertices, faces, normalized_scale=0.9, rotate_x=False, center_xyz=False): | |
| vertices = vertices.cuda() | |
| if center_xyz: # some mesh's center is not origin | |
| vertices = vertices - vertices.mean(0, keepdim=True) | |
| if rotate_x: # rotate along x axis for 90 degrees to match the two coordiantes | |
| rot_mat = torch.eye(n=3, device='cuda') | |
| theta = np.pi / 90 # degree | |
| rot_mat[1,1] = np.cos(theta) | |
| rot_mat[2,2] = np.cos(theta) | |
| rot_mat[1,2] =-np.sin(theta) | |
| rot_mat[2,1] = np.sin(theta) | |
| # ipdb.set_trace() | |
| vertices = rot_mat @ vertices.transpose(0,1) | |
| vertices = vertices.transpose(0,1) | |
| scale = (vertices.max(dim=0)[0] - vertices.min(dim=0)[0]).max() | |
| mesh_v1 = vertices / scale * normalized_scale | |
| mesh_f1 = faces.long().cuda() | |
| return mesh_v1, mesh_f1 | |
| def sample_surface_pts(path): | |
| # ipdb.set_trace() | |
| try: | |
| mesh_path, output_pth, debug = path | |
| # mesh = kal.io.obj.import_mesh(mesh_path) | |
| # ipdb.set_trace() | |
| mesh = trimesh.load(mesh_path) # fail to load ply? | |
| #ipdb.set_trace() | |
| if mesh.vertices.shape[0] == 0: | |
| return | |
| mesh_v = torch.Tensor(mesh.vertices) | |
| # mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=0.9, rotate_x=True) | |
| # mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=1.0, rotate_x=True, center_xyz=True) | |
| mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=1.0, rotate_x=False, center_xyz=True) | |
| # generate camera matrices | |
| # Rs = get_views() | |
| # Rs = get_views(semi_sphere=True) | |
| Rs = get_views(semi_sphere=False) | |
| # get depth images | |
| depths = render(mesh_v, mesh_f, Rs) | |
| # project to world space | |
| try: | |
| pcd = fusion(depths, Rs) | |
| except: | |
| return | |
| # fps again | |
| pcd, fps_idx = pytorch3d.ops.sample_farthest_points( | |
| # torch.from_numpy(pcd).unsqueeze(0).cuda(), K=4096, | |
| # torch.from_numpy(pcd).unsqueeze(0).cuda(), K=4000, | |
| pcd.unsqueeze(0).cuda(), K=4096, | |
| random_start_point=True) # B self.latent_num | |
| pcd = pcd[0] | |
| pcd = pcd.cpu().numpy() | |
| #np.savez(output_pth, pcd=pcd) | |
| #ipdb.set_trace() | |
| #if debug: | |
| pcd = trimesh.points.PointCloud(pcd) | |
| pcd.export(output_pth.replace('.npz', '.obj')) | |
| except Exception as e: | |
| # print('error') | |
| print(e, flush=True) | |
| if __name__ == '__main__': | |
| mp.set_start_method('spawn') | |
| shapenet_root = options.shape_root | |
| save_root = options.save_root | |
| debug = True | |
| cmds = [] | |
| # ! for baseline samples | |
| objv_dataset = '/mnt/sfs-common/yslan/Dataset/Obajverse/chunk-jpeg-normal/bs_16_fixsave3/170K/512/' | |
| dataset_json = os.path.join(objv_dataset, 'dataset.json') | |
| with open(dataset_json, 'r') as f: | |
| dataset_json = json.load(f) | |
| # all_objs = dataset_json['Animals'][::3][:6250] | |
| all_objs = dataset_json['Animals'][::3][1100:2200] | |
| all_objs = all_objs[:600] | |
| os.makedirs(save_root, exist_ok=True) | |
| for obj_folder in tqdm.tqdm(all_objs): | |
| obj_folder = '/'.join(obj_folder.split('/')[1:]) | |
| for idx in [0]: | |
| if 'CRM' in shapenet_root: | |
| # ipdb.set_trace() | |
| mesh_path = glob.glob(os.path.join(shapenet_root, obj_folder, f'{idx}', '*.obj'))[0] | |
| elif 'Lara' in shapenet_root: | |
| try: | |
| mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), f'{idx}.jpg', '*.obj'))[0] | |
| except: | |
| continue | |
| elif 'scale3d' in shapenet_root: | |
| if 'eval_nerf' in shapenet_root: | |
| mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '1', 'mesh', '*.ply'))[0] | |
| elif 'eval_mesh' in shapenet_root: | |
| mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '1', 'mesh', '*.obj'))[0] | |
| # files.append(os.path.join(path, '/'.join(obj_folder.split('/')[:-1]), '1', f'{idx}.png')) | |
| elif 'LRM' in shapenet_root: | |
| mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '0', '*.ply'))[0] | |
| elif 'ln3diff' in shapenet_root: | |
| mesh_path = glob.glob(os.path.join(shapenet_root, obj_folder, '0', 'mesh.obj'))[0] | |
| else: | |
| if os.path.exists((os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.obj'))): | |
| mesh_path = os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.obj') | |
| else: | |
| mesh_path = os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.ply') | |
| save_name = '-'.join(obj_folder.split('/')) | |
| cmds += [(mesh_path, os.path.join(save_root, f'{save_name}_pcd_4096.ply'), debug)] | |
| if options.n_proc == 0: | |
| for filepath in tqdm.tqdm(cmds): | |
| sample_surface_pts(filepath) | |
| else: | |
| with Pool(options.n_proc) as p: | |
| list(tqdm.tqdm(p.imap(sample_surface_pts, cmds), total=len(cmds))) | |