import logging import os import sys import cv2 import numpy as np from absl import app import gin from internal import configs from internal import datasets from internal import models from internal import utils from internal import coord from internal import checkpoints import torch import accelerate from tqdm import tqdm from torch.utils._pytree import tree_map import torch.nn.functional as F from skimage import measure import trimesh import pymeshlab as pml configs.define_common_flags() @torch.no_grad() def evaluate_density(model, accelerator: accelerate.Accelerator, points, config: configs.Config, std_value=0.0): """ Evaluate a signed distance function (SDF) for a batch of points. Args: sdf: A callable function that takes a tensor of size (N, 3) containing 3D points and returns a tensor of size (N,) with the SDF values. points: A torch tensor containing 3D points. Returns: A torch tensor with the SDF values evaluated at the given points. """ z = [] for _, pnts in enumerate(tqdm(torch.split(points, config.render_chunk_size, dim=0), desc="Evaluating density", leave=False, disable=not accelerator.is_main_process)): rays_remaining = pnts.shape[0] % accelerator.num_processes if rays_remaining != 0: padding = accelerator.num_processes - rays_remaining pnts = torch.cat([pnts, torch.zeros_like(pnts[-padding:])], dim=0) else: padding = 0 rays_per_host = pnts.shape[0] // accelerator.num_processes start, stop = accelerator.process_index * rays_per_host, \ (accelerator.process_index + 1) * rays_per_host chunk_means = pnts[start:stop] chunk_stds = torch.full_like(chunk_means[..., 0], std_value) raw_density = model.nerf_mlp.predict_density(chunk_means[:, None], chunk_stds[:, None], no_warp=True)[0] density = F.softplus(raw_density + model.nerf_mlp.density_bias) density = accelerator.gather(density) if padding > 0: density = density[: -padding] z.append(density) z = torch.cat(z, dim=0) return z @torch.no_grad() def evaluate_color(model, accelerator: accelerate.Accelerator, points, config: configs.Config, std_value=0.0): """ Evaluate a signed distance function (SDF) for a batch of points. Args: sdf: A callable function that takes a tensor of size (N, 3) containing 3D points and returns a tensor of size (N,) with the SDF values. points: A torch tensor containing 3D points. Returns: A torch tensor with the SDF values evaluated at the given points. """ z = [] for _, pnts in enumerate(tqdm(torch.split(points, config.render_chunk_size, dim=0), desc="Evaluating color", disable=not accelerator.is_main_process)): rays_remaining = pnts.shape[0] % accelerator.num_processes if rays_remaining != 0: padding = accelerator.num_processes - rays_remaining pnts = torch.cat([pnts, torch.zeros_like(pnts[-padding:])], dim=0) else: padding = 0 rays_per_host = pnts.shape[0] // accelerator.num_processes start, stop = accelerator.process_index * rays_per_host, \ (accelerator.process_index + 1) * rays_per_host chunk_means = pnts[start:stop] chunk_stds = torch.full_like(chunk_means[..., 0], std_value) chunk_viewdirs = torch.zeros_like(chunk_means) ray_results = model.nerf_mlp(False, chunk_means[:, None, None], chunk_stds[:, None, None], chunk_viewdirs) rgb = ray_results['rgb'][:, 0] rgb = accelerator.gather(rgb) if padding > 0: rgb = rgb[: -padding] z.append(rgb) z = torch.cat(z, dim=0) return z @torch.no_grad() def evaluate_color_projection(model, accelerator: accelerate.Accelerator, vertices, faces, config: configs.Config): normals = auto_normals(vertices, faces.long()) viewdirs = -normals origins = vertices - 0.005 * viewdirs vc = [] chunk = config.render_chunk_size model.num_levels = 1 model.opaque_background = True for i in tqdm(range(0, origins.shape[0], chunk), desc="Evaluating color projection", disable=not accelerator.is_main_process): cur_chunk = min(chunk, origins.shape[0] - i) rays_remaining = cur_chunk % accelerator.num_processes rays_per_host = cur_chunk // accelerator.num_processes if rays_remaining != 0: padding = accelerator.num_processes - rays_remaining rays_per_host += 1 else: padding = 0 start = i + accelerator.process_index * rays_per_host stop = start + rays_per_host batch = { 'origins': origins[start:stop], 'directions': viewdirs[start:stop], 'viewdirs': viewdirs[start:stop], 'cam_dirs': viewdirs[start:stop], 'radii': torch.full_like(origins[start:stop, ..., :1], 0.000723), 'near': torch.full_like(origins[start:stop, ..., :1], 0), 'far': torch.full_like(origins[start:stop, ..., :1], 0.01), } batch = accelerator.pad_across_processes(batch) with accelerator.autocast(): renderings, ray_history = model( False, batch, compute_extras=False, train_frac=1) rgb = renderings[-1]['rgb'] acc = renderings[-1]['acc'] rgb /= acc.clamp_min(1e-5)[..., None] rgb = rgb.clamp(0, 1) rgb = accelerator.gather(rgb) rgb[torch.isnan(rgb) | torch.isinf(rgb)] = 1 if padding > 0: rgb = rgb[: -padding] vc.append(rgb) vc = torch.cat(vc, dim=0) return vc def auto_normals(verts, faces): i0 = faces[:, 0] i1 = faces[:, 1] i2 = faces[:, 2] v0 = verts[i0, :] v1 = verts[i1, :] v2 = verts[i2, :] face_normals = torch.cross(v1 - v0, v2 - v0) # Splat face normals to vertices v_nrm = torch.zeros_like(verts) v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) # Normalize, replace zero (degenerated) normals with some default value v_nrm = torch.where((v_nrm ** 2).sum(dim=-1, keepdims=True) > 1e-20, v_nrm, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=verts.device)) v_nrm = F.normalize(v_nrm, dim=-1) return v_nrm def clean_mesh(verts, faces, v_pct=1, min_f=8, min_d=5, repair=True, remesh=True, remesh_size=0.01, logger=None, main_process=True): # verts: [N, 3] # faces: [N, 3] tbar = tqdm(total=9, desc='Clean mesh', leave=False, disable=not main_process) _ori_vert_shape = verts.shape _ori_face_shape = faces.shape m = pml.Mesh(verts, faces) ms = pml.MeshSet() ms.add_mesh(m, 'mesh') # will copy! # filters tbar.set_description('Remove unreferenced vertices') ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces tbar.update() if v_pct > 0: tbar.set_description('Remove unreferenced vertices') ms.meshing_merge_close_vertices(threshold=pml.Percentage(v_pct)) # 1/10000 of bounding box diagonal tbar.update() tbar.set_description('Remove duplicate faces') ms.meshing_remove_duplicate_faces() # faces defined by the same verts tbar.update() tbar.set_description('Remove null faces') ms.meshing_remove_null_faces() # faces with area == 0 tbar.update() if min_d > 0: tbar.set_description('Remove connected component by diameter') ms.meshing_remove_connected_component_by_diameter(mincomponentdiag=pml.Percentage(min_d)) tbar.update() if min_f > 0: tbar.set_description('Remove connected component by face number') ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f) tbar.update() if repair: # tbar.set_description('Remove t vertices') # ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True) tbar.set_description('Repair non manifold edges') ms.meshing_repair_non_manifold_edges(method=0) tbar.update() tbar.set_description('Repair non manifold vertices') ms.meshing_repair_non_manifold_vertices(vertdispratio=0) tbar.update() else: tbar.update(2) if remesh: # tbar.set_description('Coord taubin smoothing') # ms.apply_coord_taubin_smoothing() tbar.set_description('Isotropic explicit remeshing') ms.meshing_isotropic_explicit_remeshing(iterations=3, targetlen=pml.AbsoluteValue(remesh_size)) tbar.update() # extract mesh m = ms.current_mesh() verts = m.vertex_matrix() faces = m.face_matrix() if logger is not None: logger.info(f'Mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}') return verts, faces def decimate_mesh(verts, faces, target, backend='pymeshlab', remesh=False, optimalplacement=True, logger=None): # optimalplacement: default is True, but for flat mesh must turn False to prevent spike artifect. _ori_vert_shape = verts.shape _ori_face_shape = faces.shape if backend == 'pyfqmr': import pyfqmr solver = pyfqmr.Simplify() solver.setMesh(verts, faces) solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False) verts, faces, normals = solver.getMesh() else: m = pml.Mesh(verts, faces) ms = pml.MeshSet() ms.add_mesh(m, 'mesh') # will copy! # filters # ms.meshing_decimation_clustering(threshold=pml.Percentage(1)) ms.meshing_decimation_quadric_edge_collapse(targetfacenum=int(target), optimalplacement=optimalplacement) if remesh: # ms.apply_coord_taubin_smoothing() ms.meshing_isotropic_explicit_remeshing(iterations=3, targetlen=pml.Percentage(1)) # extract mesh m = ms.current_mesh() verts = m.vertex_matrix() faces = m.face_matrix() if logger is not None: logger.info(f'Mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}') return verts, faces def main(unused_argv): config = configs.load_config() config.compute_visibility = True config.exp_path = os.path.join("exp", config.exp_name) config.mesh_path = os.path.join("exp", config.exp_name, "mesh") config.checkpoint_dir = os.path.join(config.exp_path, 'checkpoints') os.makedirs(config.mesh_path, exist_ok=True) # accelerator for DDP accelerator = accelerate.Accelerator() device = accelerator.device # setup logger logging.basicConfig( format="%(asctime)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, handlers=[logging.StreamHandler(sys.stdout), logging.FileHandler(os.path.join(config.exp_path, 'log_extract.txt'))], level=logging.INFO, ) sys.excepthook = utils.handle_exception logger = accelerate.logging.get_logger(__name__) logger.info(config) logger.info(accelerator.state, main_process_only=False) config.world_size = accelerator.num_processes config.global_rank = accelerator.process_index accelerate.utils.set_seed(config.seed, device_specific=True) # setup model and optimizer model = models.Model(config=config) model = accelerator.prepare(model) step = checkpoints.restore_checkpoint(config.checkpoint_dir, accelerator, logger) model.eval() module = accelerator.unwrap_model(model) visibility_path = os.path.join(config.mesh_path, 'visibility_mask_{:.1f}.pt'.format(config.mesh_radius)) visibility_resolution = config.visibility_resolution if not os.path.exists(visibility_path): logger.info('Generate visibility mask...') # load dataset dataset = datasets.load_dataset('train', config.data_dir, config) dataloader = torch.utils.data.DataLoader(np.arange(len(dataset)), num_workers=4, shuffle=True, batch_size=1, collate_fn=dataset.collate_fn, persistent_workers=True, ) visibility_mask = torch.ones( (1, 1, visibility_resolution, visibility_resolution, visibility_resolution), requires_grad=True ).to(device) visibility_mask.retain_grad() tbar = tqdm(dataloader, desc='Generating visibility grid', disable=not accelerator.is_main_process) for index, batch in enumerate(tbar): batch = accelerate.utils.send_to_device(batch, accelerator.device) rendering = models.render_image(model, accelerator, batch, False, 1, config, verbose=False, return_weights=True) coords = rendering['coord'].reshape(-1, 3) weights = rendering['weights'].reshape(-1) valid_points = coords[weights > config.valid_weight_thresh] valid_points /= config.mesh_radius # update mask based on ray samples with torch.enable_grad(): out = torch.nn.functional.grid_sample(visibility_mask, valid_points[None, None, None], align_corners=True) out.sum().backward() tbar.set_postfix({"visibility_mask": (visibility_mask.grad > 0.0001).float().mean().item()}) # if index == 10: # break visibility_mask = (visibility_mask.grad > 0.0001).float() if accelerator.is_main_process: torch.save(visibility_mask.detach().cpu(), visibility_path) else: logger.info('Load visibility mask from {}'.format(visibility_path)) visibility_mask = torch.load(visibility_path, map_location=device) space = config.mesh_radius * 2 / (config.visibility_resolution - 1) logger.info("Extract mesh from visibility mask...") visibility_mask_np = visibility_mask[0, 0].permute(2, 1, 0).detach().cpu().numpy() verts, faces, normals, values = measure.marching_cubes( volume=-visibility_mask_np, level=-0.5, spacing=(space, space, space)) verts -= config.mesh_radius if config.extract_visibility: meshexport = trimesh.Trimesh(verts, faces) meshexport.export(os.path.join(config.mesh_path, "visibility_mask_{}.ply".format(config.mesh_radius)), "ply") logger.info("Extract visibility mask done.") # Initialize variables crop_n = 512 grid_min = verts.min(axis=0) grid_max = verts.max(axis=0) space = ((grid_max - grid_min).prod() / config.mesh_voxels) ** (1 / 3) world_size = ((grid_max - grid_min) / space).astype(np.int32) Nx, Ny, Nz = np.maximum(1, world_size // crop_n) crop_n_x, crop_n_y, crop_n_z = world_size // [Nx, Ny, Nz] xs = np.linspace(grid_min[0], grid_max[0], Nx + 1) ys = np.linspace(grid_min[1], grid_max[1], Ny + 1) zs = np.linspace(grid_min[2], grid_max[2], Nz + 1) # Initialize meshes list meshes = [] # Iterate over the grid for i in range(Nx): for j in range(Ny): for k in range(Nz): logger.info(f"Process grid cell ({i + 1}/{Nx}, {j + 1}/{Ny}, {k + 1}/{Nz})...") # Calculate grid cell boundaries x_min, x_max = xs[i], xs[i + 1] y_min, y_max = ys[j], ys[j + 1] z_min, z_max = zs[k], zs[k + 1] # Create point grid x = np.linspace(x_min, x_max, crop_n_x) y = np.linspace(y_min, y_max, crop_n_y) z = np.linspace(z_min, z_max, crop_n_z) xx, yy, zz = np.meshgrid(x, y, z, indexing="ij") points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float, device=device) # Construct point pyramids points_tmp = points.reshape(crop_n_x, crop_n_y, crop_n_z, 3)[None] points_tmp /= config.mesh_radius current_mask = torch.nn.functional.grid_sample(visibility_mask, points_tmp, align_corners=True) current_mask = (current_mask > 0.0).cpu().numpy()[0, 0] pts_density = evaluate_density(module, accelerator, points, config, std_value=config.std_value) # bound the vertices points_world = coord.inv_contract(2 * points) pts_density[points_world.norm(dim=-1) > config.mesh_max_radius] = 0.0 z = pts_density.detach().cpu().numpy() # Skip if no surface found valid_z = z.reshape(crop_n_x, crop_n_y, crop_n_z)[current_mask] if valid_z.shape[0] <= 0 or ( np.min(valid_z) > config.isosurface_threshold or np.max( valid_z) < config.isosurface_threshold ): continue if not (np.min(z) > config.isosurface_threshold or np.max(z) < config.isosurface_threshold): # Extract mesh logger.info('Extract mesh...') z = z.astype(np.float32) verts, faces, _, _ = measure.marching_cubes( volume=-z.reshape(crop_n_x, crop_n_y, crop_n_z), level=-config.isosurface_threshold, spacing=( (x_max - x_min) / (crop_n_x - 1), (y_max - y_min) / (crop_n_y - 1), (z_max - z_min) / (crop_n_z - 1), ), mask=current_mask, ) verts = verts + np.array([x_min, y_min, z_min]) meshcrop = trimesh.Trimesh(verts, faces) logger.info('Extract vertices: {}, faces: {}'.format(meshcrop.vertices.shape[0], meshcrop.faces.shape[0])) meshes.append(meshcrop) # Save mesh logger.info('Concatenate mesh...') combined_mesh = trimesh.util.concatenate(meshes) # from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/renderer.py # clean logger.info('Clean mesh...') vertices = combined_mesh.vertices.astype(np.float32) faces = combined_mesh.faces.astype(np.int32) vertices, faces = clean_mesh(vertices, faces, remesh=False, remesh_size=0.01, logger=logger, main_process=accelerator.is_main_process) v = torch.from_numpy(vertices).contiguous().float().to(device) v = coord.inv_contract(2 * v) vertices = v.detach().cpu().numpy() f = torch.from_numpy(faces).contiguous().int().to(device) # decimation if config.decimate_target > 0 and faces.shape[0] > config.decimate_target: logger.info('Decimate mesh...') vertices, triangles = decimate_mesh(vertices, faces, config.decimate_target, logger=logger) # import ipdb; ipdb.set_trace() if config.vertex_color: # batched inference to avoid OOM logger.info('Evaluate mesh vertex color...') if config.vertex_projection: rgbs = evaluate_color_projection(module, accelerator, v, f, config) else: rgbs = evaluate_color(module, accelerator, v, config, std_value=config.std_value) rgbs = (rgbs * 255).detach().cpu().numpy().astype(np.uint8) if accelerator.is_main_process: logger.info('Export mesh (vertex color)...') mesh = trimesh.Trimesh(vertices, faces, vertex_colors=rgbs, process=False) # important, process=True leads to seg fault... mesh.export(os.path.join(config.mesh_path, 'mesh_{}.ply'.format(config.mesh_radius))) logger.info('Finish extracting mesh.') return def _export(v, f, h0=2048, w0=2048, ssaa=1, name=''): logger.info('Export mesh (atlas)...') # v, f: torch Tensor device = v.device v_np = v.cpu().numpy() # [N, 3] f_np = f.cpu().numpy() # [M, 3] # unwrap uvs import xatlas import nvdiffrast.torch as dr from sklearn.neighbors import NearestNeighbors from scipy.ndimage import binary_dilation, binary_erosion logger.info(f'Running xatlas to unwrap UVs for mesh: v={v_np.shape} f={f_np.shape}') atlas = xatlas.Atlas() atlas.add_mesh(v_np, f_np) chart_options = xatlas.ChartOptions() chart_options.max_iterations = 4 # for faster unwrap... atlas.generate(chart_options=chart_options) vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2] # vmapping, ft_np, vt_np = xatlas.parametrize(v_np, f_np) # [N], [M, 3], [N, 2] vt = torch.from_numpy(vt_np.astype(np.float32)).float().to(device) ft = torch.from_numpy(ft_np.astype(np.int64)).int().to(device) # render uv maps uv = vt * 2.0 - 1.0 # uvs to range [-1, 1] uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4] if ssaa > 1: h = int(h0 * ssaa) w = int(w0 * ssaa) else: h, w = h0, w0 if h <= 2048 and w <= 2048: glctx = dr.RasterizeCudaContext() else: glctx = dr.RasterizeGLContext() rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), ft, (h, w)) # [1, h, w, 4] xyzs, _ = dr.interpolate(v.unsqueeze(0), rast, f) # [1, h, w, 3] mask, _ = dr.interpolate(torch.ones_like(v[:, :1]).unsqueeze(0), rast, f) # [1, h, w, 1] # masked query xyzs = xyzs.view(-1, 3) mask = (mask > 0).view(-1) feats = torch.zeros(h * w, 3, device=device, dtype=torch.float32) if mask.any(): xyzs = xyzs[mask] # [M, 3] # batched inference to avoid OOM all_feats = evaluate_color(module, accelerator, xyzs, config, std_value=config.std_value) feats[mask] = all_feats feats = feats.view(h, w, -1) mask = mask.view(h, w) # quantize [0.0, 1.0] to [0, 255] feats = feats.cpu().numpy() feats = (feats * 255).astype(np.uint8) ### NN search as an antialiasing ... mask = mask.cpu().numpy() inpaint_region = binary_dilation(mask, iterations=3) inpaint_region[mask] = 0 search_region = mask.copy() not_search_region = binary_erosion(search_region, iterations=2) search_region[not_search_region] = 0 search_coords = np.stack(np.nonzero(search_region), axis=-1) inpaint_coords = np.stack(np.nonzero(inpaint_region), axis=-1) knn = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(search_coords) _, indices = knn.kneighbors(inpaint_coords) feats[tuple(inpaint_coords.T)] = feats[tuple(search_coords[indices[:, 0]].T)] feats = cv2.cvtColor(feats, cv2.COLOR_RGB2BGR) # do ssaa after the NN search, in numpy if ssaa > 1: feats = cv2.resize(feats, (w0, h0), interpolation=cv2.INTER_LINEAR) cv2.imwrite(os.path.join(config.mesh_path, f'{name}albedo.png'), feats) # save obj (v, vt, f /) obj_file = os.path.join(config.mesh_path, f'{name}mesh.obj') mtl_file = os.path.join(config.mesh_path, f'{name}mesh.mtl') logger.info(f'writing obj mesh to {obj_file}') with open(obj_file, "w") as fp: fp.write(f'mtllib {name}mesh.mtl \n') logger.info(f'writing vertices {v_np.shape}') for v in v_np: fp.write(f'v {v[0]} {v[1]} {v[2]} \n') logger.info(f'writing vertices texture coords {vt_np.shape}') for v in vt_np: fp.write(f'vt {v[0]} {1 - v[1]} \n') logger.info(f'writing faces {f_np.shape}') fp.write(f'usemtl mat0 \n') for i in range(len(f_np)): fp.write( f"f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1} {f_np[i, 1] + 1}/{ft_np[i, 1] + 1} {f_np[i, 2] + 1}/{ft_np[i, 2] + 1} \n") with open(mtl_file, "w") as fp: fp.write(f'newmtl mat0 \n') fp.write(f'Ka 1.000000 1.000000 1.000000 \n') fp.write(f'Kd 1.000000 1.000000 1.000000 \n') fp.write(f'Ks 0.000000 0.000000 0.000000 \n') fp.write(f'Tr 1.000000 \n') fp.write(f'illum 1 \n') fp.write(f'Ns 0.000000 \n') fp.write(f'map_Kd {name}albedo.png \n') # could be extremely slow _export(v, f) logger.info('Finish extracting mesh.') if __name__ == '__main__': with gin.config_scope('bake'): app.run(main)