def test_ball_query_output_simple(self): device = get_random_cuda_device() N, P1, P2, K = 5, 8, 16, 4 sphere = ico_sphere(level=2, device=device).extend(N) points_1 = sample_points_from_meshes(sphere, P1) points_2 = sample_points_from_meshes(sphere, P2) * 5.0 radius = 6.0 naive_out = self._ball_query_naive( points_1, points_2, lengths1=None, lengths2=None, K=K, radius=radius ) cuda_out = ball_query(points_1, points_2, K=K, radius=radius) # All points should have N sample neighbors as radius is large # Zero is a valid index but can only be present once (i.e. no zero padding) naive_out_zeros = (naive_out.idx == 0).sum(dim=-1).max() cuda_out_zeros = (cuda_out.idx == 0).sum(dim=-1).max() self.assertTrue(naive_out_zeros == 0 or naive_out_zeros == 1) self.assertTrue(cuda_out_zeros == 0 or cuda_out_zeros == 1) # All points should now have zero sample neighbors as radius is small radius = 0.5 naive_out = self._ball_query_naive( points_1, points_2, lengths1=None, lengths2=None, K=K, radius=radius ) cuda_out = ball_query(points_1, points_2, K=K, radius=radius) naive_out_allzeros = (naive_out.idx == -1).all() cuda_out_allzeros = (cuda_out.idx == -1).sum() self.assertTrue(naive_out_allzeros) self.assertTrue(cuda_out_allzeros)
def get_loss(mesh, trg_mesh, w_chamfer, w_edge, w_normal, w_laplacian, n_points=5000): # We sample 5k points from the surface of each mesh sample_trg = sample_points_from_meshes(trg_mesh, n_points) sample_src = sample_points_from_meshes(mesh, n_points) # We compare the two sets of pointclouds by computing (a) the chamfer loss loss_chamfer, _ = chamfer_distance(sample_trg, sample_src) # and (b) the edge length of the predicted mesh loss_edge = mesh_edge_loss(mesh) # mesh normal consistency loss_normal = mesh_normal_consistency(mesh) # mesh laplacian smoothing loss_laplacian = mesh_laplacian_smoothing(mesh, method="uniform") # Weighted sum of the losses loss = loss_chamfer * w_chamfer + loss_edge * w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian return loss
def save_predictions(model, loader, output_dir): """ This function is used save predicted and gt meshes """ # Note that all eval runs on main process assert comm.is_main_process() if isinstance(model, torch.nn.parallel.DistributedDataParallel): model = model.module device = torch.device("cuda:0") os.makedirs(output_dir, exist_ok=True) for batch in tqdm.tqdm(loader): batch = loader.postprocess(batch, device) model_kwargs = {} module = model.module if hasattr(model, "module") else model if isinstance(module, VoxMeshMultiViewHead): model_kwargs["intrinsics"] = batch["intrinsics"] model_kwargs["extrinsics"] = batch["extrinsics"] if isinstance(module, VoxDepthHead): model_kwargs["masks"] = batch["masks"] if module.mvsnet is None: model_kwargs["depths"] = batch["depths"] model_outputs = model(batch["imgs"], **model_kwargs) # TODO: debug only # save_debug_predictions(batch, model_outputs) # continue pred_mesh = model_outputs["meshes_pred"][-1] gt_mesh = batch["meshes"] pred_mesh = pred_mesh.scale_verts(P2M_SCALE) gt_mesh = gt_mesh.scale_verts(P2M_SCALE) pred_points = sample_points_from_meshes( pred_mesh, NUM_PRED_SURFACE_SAMPLES, return_normals=False ) gt_points = sample_points_from_meshes( gt_mesh, NUM_GT_SURFACE_SAMPLES, return_normals=False ) pred_points = pred_points.cpu().detach().numpy() gt_points = gt_points.cpu().detach().numpy() batch_size = pred_points.shape[0] for batch_idx in range(batch_size): label, label_appendix = batch["id_strs"][batch_idx].split("-")[:2] pred_filename = os.path.join( output_dir, "{}_{}_predict.xyz".format(label, label_appendix) ) gt_filename = os.path.join( output_dir, "{}_{}_ground.xyz".format(label, label_appendix) ) np.savetxt(pred_filename, pred_points[batch_idx]) np.savetxt(gt_filename, gt_points[batch_idx])
def test_all_empty_meshes(self): """ Check sample_points_from_meshes raises an exception if all meshes are invalid. """ device = torch.device("cuda:0") verts1 = torch.tensor([], dtype=torch.float32, device=device) faces1 = torch.tensor([], dtype=torch.int64, device=device) meshes = Meshes(verts=[verts1, verts1, verts1], faces=[faces1, faces1, faces1]) with self.assertRaises(ValueError) as err: sample_points_from_meshes(meshes, num_samples=100, return_normals=True) self.assertTrue("Meshes are empty." in str(err.exception))
def plot_pointcloud(mesh, title=""): # Sample points uniformly from the surface of the mesh. points = sample_points_from_meshes(mesh, 5000) x, y, z = points.clone().detach().cpu().squeeze().unbind(1) """ fig = go.Figure() fig = px.scatter_3d( x = x, y = y, z = z, labels={'x':'x', 'y':'y', 'z':'z'} ) fig.update_layout( title_text = f"PointCloud: {title}" ) fig.show() """ # Matplot 3D scatter plot fig = plt.figure(figsize=(5, 5)) ax = Axes3D(fig) ax.scatter3D(x, z, -y) ax.set_xlabel("x") ax.set_ylabel("z") ax.set_zlabel("y") ax.set_title(title) ax.view_init(190, 30) plt.show()
def postprocess(self, batch, device=None): if device is None: device = torch.device("cuda") imgs, meshes, points, normals, voxels, Ps, id_strs = batch imgs = imgs.to(device) if meshes is not None: meshes = meshes.to(device) if points is not None and normals is not None: points = points.to(device) normals = normals.to(device) else: points, normals = sample_points_from_meshes( meshes, num_samples=self.num_samples, return_normals=True ) if voxels is not None: if torch.is_tensor(voxels): # We used cached voxels on disk, just cast and return voxels = voxels.to(device) else: # We got a list of voxel_coords, and need to compute voxels on-the-fly voxel_coords = voxels Ps = Ps.to(device) voxels = [] for i, cur_voxel_coords in enumerate(voxel_coords): cur_voxel_coords = cur_voxel_coords.to(device) cur_voxels = self._voxelize(cur_voxel_coords, Ps[i]) voxels.append(cur_voxels) voxels = torch.stack(voxels, dim=0) if self.return_id_str: return imgs, meshes, points, normals, voxels, id_strs else: return imgs, meshes, points, normals, voxels
def _sample_meshes(meshes, num_samples): """ Helper to either sample points uniformly from the surface of a mesh (with normals), or take the verts of the mesh as samples. Inputs: - meshes: A MeshList - num_samples: An integer, or the string 'verts' Outputs: - verts: Either a Tensor of shape (N, S, 3) if we take the same number of samples from each mesh; otherwise a list of length N, whose ith element is a Tensor of shape (S_i, 3) - normals: Either a Tensor of shape (N, S, 3) or None if we take verts as samples. """ if num_samples == "verts": normals = None if meshes.equisized: verts = meshes.verts_batch else: verts = meshes.verts_list else: verts, normals = sample_points_from_meshes(meshes, num_samples, return_normals=True) return verts, normals
def load_ShapeNet_pointclouds(data_path, split_path, codes_path, this_device): codes = torch.load(codes_path) #print(codes['latent_codes']['weight'].shape) split = json.load(open(split_path)) object_list = get_objectnames_from_split(split) object_paths = [] i = 0 for object_name in object_list: object_paths.append( os.path.join(data_path, object_name, "models/model_normalized.obj")) start1 = time.perf_counter() print(start1) print("Start loading of objects.") input_meshes = load_objs_as_meshes(object_paths, device=this_device, load_textures=False) loading_time = time.perf_counter() - start1 print("Loading of objects finished. Time necessary to load " + str(len(object_paths)) + " objects is: " + str(loading_time) + " seconds.") print("Start sampling of points.") start2 = time.perf_counter() number_samples = 4096 input_pointclouds = sample_points_from_meshes(input_meshes, number_samples) sampling_time = time.perf_counter() - start2 print("Sampling of points finished. Time necessary to sample " + str(number_samples) + " points from " + str(len(object_paths)) + " objects each is: " + str(sampling_time) + " seconds.") return input_pointclouds, codes['latent_codes']['weight']
def forward(self, src_mesh): loss = 0 # Sample from target meshes target_verts = sample_points_from_meshes(self.target_meshes, 3000) if self.consider_loss("chamfer"): loss_chamfer, _ = chamfer_distance(target_verts, src_mesh.verts_padded()) loss += self.loss_weights["w_chamfer"] * loss_chamfer if self.consider_loss("edge"): loss_edge = mesh_edge_loss( src_mesh) # and (b) the edge length of the predicted mesh loss += self.loss_weights["w_edge"] * loss_edge if self.consider_loss("normal"): loss_normal = mesh_normal_consistency( src_mesh) # mesh normal consistency loss += self.loss_weights["w_normal"] * loss_normal if self.consider_loss("laplacian"): loss_laplacian = mesh_laplacian_smoothing( src_mesh, method="uniform") # mesh laplacian smoothing loss += self.loss_weights["w_laplacian"] * loss_laplacian return loss
def animate_pointcloud(mesh, anim_file, points_to_sample, restore_anim=True, is_mesh=True): if os.path.isfile(anim_file) and restore_anim: return anim_file frames = [] for plot_i in range(24): if is_mesh: points = sample_points_from_meshes(mesh, points_to_sample) x, y, z = points.clone().detach().cpu().squeeze().unbind(1) else: x, y, z = mesh.unbind(1) fig = plt.figure(figsize=(5, 5)) canvas = FigureCanvas(fig) ax = Axes3D(fig) ax.scatter3D(x, z, -y) ax.view_init(elev=190, azim=360 * (plot_i / 24)) plt.axis('off') plt.close() canvas.draw() s, (width, height) = canvas.print_to_buffer() frames.append(np.frombuffer(s, np.uint8).reshape((height, width, 4))) imageio.mimsave(anim_file, frames, 'GIF', fps=8) return anim_file
def get_plot3d_mesh_img(mesh, title="plot mesh", n_sample=500, fig_size=(5, 5), view_points=(190, 30)): points = sample_points_from_meshes(mesh, n_sample) x, y, z = points.clone().detach().cpu().squeeze().unbind(1) fig = plt.figure(figsize=fig_size) ax = Axes3D(fig) ax.scatter3D(x, z, -y) ax.set_xlabel('x') ax.set_ylabel('z') ax.set_zlabel('y') ax.set_title(title) ax.view_init(view_points[0], view_points[1]) buffer = io.BytesIO() # bufferを用意 plt.savefig(buffer, format='png') # bufferに保持 buffer_np = np.frombuffer(buffer.getvalue(), dtype=np.uint8) # bufferからの読み出し buffer_cv = cv2.imdecode(buffer_np, 1) # デコード buffer_cv = buffer_cv[:, :, ::-1] # BGR->RGB #print( "buffer_cv.shape : ", buffer_cv.shape ) img = Image.fromarray(buffer_cv) #print( "img : ", img) return img
def loss(self, data, epoch): pred = self.forward(data) # embed() # loss_coef = max(1/(2**(epoch//10000)), 0.1) # CE_Loss = nn.CrossEntropyLoss() # ce_loss = CE_Loss(pred[0][-1][3], data['y_voxels']) weight = data['base_plane'].float().cuda() CE_Loss = nn.CrossEntropyLoss(reduction='none') ce_loss = CE_Loss(pred[0][-1][3], data['y_voxels'].cuda()) * weight ce_loss = ce_loss.mean() chamfer_loss = torch.tensor(0).float().cuda() edge_loss = torch.tensor(0).float().cuda() laplacian_loss = torch.tensor(0).float().cuda() normal_consistency_loss = torch.tensor(0).float().cuda() for c in range(self.config.num_classes-1): target = data['surface_points'][c].cuda() for k, (vertices, faces, _, _, _) in enumerate(pred[c][1:]): pred_mesh = Meshes(verts=list(vertices), faces=list(faces)) pred_points = sample_points_from_meshes(pred_mesh, 3000) chamfer_loss += chamfer_distance(pred_points, target)[0] laplacian_loss += mesh_laplacian_smoothing(pred_mesh, method="uniform") normal_consistency_loss += mesh_normal_consistency(pred_mesh) edge_loss += mesh_edge_loss(pred_mesh) # vertices, faces, _, _, _ = pred[c][-1] # pred_mesh = Meshes(verts=list(vertices), faces=list(faces)) # pred_points = sample_points_from_meshes(pred_mesh, 3000) # # chamfer_loss += chamfer_distance(pred_points, target)[0]*5 # laplacian_loss += mesh_laplacian_smoothing(pred_mesh, method="uniform")*5 # normal_consistency_loss += mesh_normal_consistency(pred_mesh)*5 # edge_loss += mesh_edge_loss(pred_mesh)*5 # # # chamfer_loss = chamfer_loss/2 # # laplacian_loss = laplacian_loss/2 # # normal_consistency_loss = normal_consistency_loss/2 # # edge_loss = edge_loss/2 loss = 1 * chamfer_loss + 1 * ce_loss + 0.1 * laplacian_loss + 1 * edge_loss + 0.1 * normal_consistency_loss # loss = 1 * chamfer_loss + 0.1 * laplacian_loss + 1 * edge_loss + 0.1 * normal_consistency_loss # loss = 1 * chamfer_loss + 0.1 * laplacian_loss + loss_coef * edge_loss + 0.1 * normal_consistency_loss log = {"loss": loss.detach(), "chamfer_loss": chamfer_loss.detach(), # "loss_coef": loss_coef, "ce_loss": ce_loss.detach(), "normal_consistency_loss": normal_consistency_loss.detach(), "edge_loss": edge_loss.detach(), "laplacian_loss": laplacian_loss.detach()} return loss, log
def test_verts_nan(self): num_verts = 30 num_faces = 50 for device in ["cpu", "cuda:0"]: for invalid in ["nan", "inf"]: verts = torch.rand((num_verts, 3), dtype=torch.float32, device=device) # randomly assign an invalid type verts[torch.randperm(num_verts)[:10]] = float(invalid) faces = torch.randint( num_verts, size=(num_faces, 3), dtype=torch.int64, device=device ) meshes = Meshes(verts=[verts], faces=[faces]) with self.assertRaisesRegex(ValueError, "Meshes contain nan or inf."): sample_points_from_meshes( meshes, num_samples=100, return_normals=True )
def test_texture_sampling_cow(self): # test texture sampling for the cow example by converting # the cow mesh and its texture uv to a pointcloud with texture device = torch.device("cuda:0") obj_dir = get_pytorch3d_dir() / "docs/tutorials/data" obj_filename = obj_dir / "cow_mesh/cow.obj" for text_type in ("uv", "atlas"): # Load mesh + texture if text_type == "uv": mesh = load_objs_as_meshes( [obj_filename], device=device, load_textures=True, texture_wrap=None ) elif text_type == "atlas": mesh = load_objs_as_meshes( [obj_filename], device=device, load_textures=True, create_texture_atlas=True, texture_atlas_size=8, texture_wrap=None, ) points, normals, textures = sample_points_from_meshes( mesh, num_samples=50000, return_normals=True, return_textures=True ) pointclouds = Pointclouds(points, normals=normals, features=textures) for pos in ("front", "back"): # Init rasterizer settings if pos == "back": azim = 0.0 elif pos == "front": azim = 180 R, T = look_at_view_transform(2.7, 0, azim) cameras = FoVPerspectiveCameras(device=device, R=R, T=T) raster_settings = PointsRasterizationSettings( image_size=512, radius=1e-2, points_per_pixel=1 ) rasterizer = PointsRasterizer( cameras=cameras, raster_settings=raster_settings ) compositor = NormWeightedCompositor() renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor) images = renderer(pointclouds) rgb = images[0, ..., :3].squeeze().cpu() if DEBUG: filename = "DEBUG_cow_mesh_to_pointcloud_%s_%s.png" % ( text_type, pos, ) Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( DATA_DIR / filename )
def __init__(self, n_it: int, param_group: str, SMBLD: SMBLDMesh, target_meshes: Meshes, mesh_names=[], name="optimise", loss_weights=None, lr=1e-3, lr_decay=1.0, out_dir="static_fits_output", custom_lrs=None): """ n_its = integer, number of iterations in stage parameters = list of items over which to be optimised get_mesh = function that returns Mesh object for identifying losses name = name of stage lr_decay = factor by which lr decreases at each it""" self.n_it = n_it self.name = name self.out_dir = out_dir self.target_meshes = target_meshes self.mesh_names = mesh_names self.SMBLD = SMBLD self.loss_weights = default_weights.copy() if loss_weights is not None: for k, v in loss_weights.items(): self.loss_weights[k] = v self.losses_to_plot = [] # Store losses for review later if custom_lrs is not None: for attr in custom_lrs: assert hasattr(SMBLD, attr), f"attr '{attr}' not in SMBLD." self.param_group = SMBLDMeshParamGroup(SMBLD, param_group, custom_lrs) self.optimizer = torch.optim.Adam(self.param_group, lr=lr) self.scheduler = torch.optim.lr_scheduler.LambdaLR( self.optimizer, lambda epoch: lr * (lr_decay)**epoch) # Decay of learning rate with torch.no_grad(): self.prev_mesh = self.SMBLD.get_meshes() self.prev_verts, _ = SMBLD.get_verts( ) # Get template verts to use for ARAP method self.n_verts = self.prev_verts.shape[1] # Sample from target meshes - an equal number to the SMBLD mesh self.target_verts = sample_points_from_meshes(self.target_meshes, 3000) self.consider_loss = lambda loss_name: self.loss_weights[ f"w_{loss_name}"] > 0 # function to check if loss is non-zero
def compute_geometric_metrics(pred_mesh_path, gdth_mesh_path): with torch.no_grad(): pred_mesh = trimesh.load(pred_mesh_path) gdth_mesh = trimesh.load(gdth_mesh_path) pred_vertices = pred_mesh.vertices gdth_vertices = gdth_mesh.vertices pred_pts = pred_mesh.sample(10000) gdth_pts = gdth_mesh.sample(10000) pred_vertices -= np.mean(pred_pts, axis=0) gdth_vertices -= np.mean(gdth_pts, axis=0) pred_vertices /= np.max(np.linalg.norm(pred_vertices, axis=1)) gdth_vertices /= np.max(np.linalg.norm(gdth_vertices, axis=1)) # sph = regular_on_sphere_points(1024) # both = np.vstack((np.vstack((pred_vertices, gdth_vertices)), sph)) # np.savetxt('test.xyz',both) pred_vertices = torch.from_numpy(pred_vertices).float() gdth_vertices = torch.from_numpy(gdth_vertices).float() pred_faces = torch.from_numpy(pred_mesh.faces) gdth_faces = torch.from_numpy(gdth_mesh.faces) pred_mesh = Meshes(verts=[pred_vertices], faces=[pred_faces]) gdth_mesh = Meshes(verts=[gdth_vertices], faces=[gdth_faces]) pred_points, pred_normals = sample_points_from_meshes( pred_mesh, num_samples=10000, return_normals=True) gt_points, gt_normals = sample_points_from_meshes(gdth_mesh, num_samples=10000, return_normals=True) metrics = _compute_sampling_metrics(pred_points, pred_normals, gt_points, gt_normals, eps=1e-8) return metrics
def get_deform_verts(target_mesh, points_to_sample=5000, sphere_level=4): device = torch.device("cuda:0") src_mesh = ico_sphere(sphere_level, device) deform_verts = torch.full(src_mesh.verts_packed().shape, 0.0, device=device, requires_grad=True) learning_rate = 0.01 num_iter = 500 w_chamfer = 1.0 w_edge = 0.05 w_normal = 0.0005 w_laplacian = 0.005 optimizer = torch.optim.Adam([deform_verts], lr=learning_rate, betas=(0.5, 0.999)) for _ in range(num_iter): optimizer.zero_grad() new_src_mesh = src_mesh.offset_verts(deform_verts) sample_trg = sample_points_from_meshes(target_mesh, points_to_sample) sample_src = sample_points_from_meshes(new_src_mesh, points_to_sample) loss_chamfer, _ = chamfer_distance(sample_trg, sample_src) loss_edge = mesh_edge_loss(new_src_mesh) loss_normal = mesh_normal_consistency(new_src_mesh) loss_laplacian = mesh_laplacian_smoothing(new_src_mesh, method="uniform") loss = loss_chamfer * w_chamfer + loss_edge * w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian loss.backward() optimizer.step() print( f"{datetime.now()} Loss Chamfer:{loss_chamfer * w_chamfer}, Loss Edge:{loss_edge * w_edge}, Loss Normal:{loss_normal * w_normal}, Loss Laplacian:{loss_laplacian * w_laplacian}" ) return deform_verts
def create_model(cfg, device, mode="train", camera_model=None, **kwargs): ''' Returns model Args: cfg (edict): imported yaml config device (device): pytorch device ''' if cfg.model.type == 'point': decoder = None texture = None use_lighting = (cfg.renderer is not None and not cfg.renderer.get('is_neural_texture', True)) if use_lighting: texture = LightingTexture() else: if 'rgb' not in cfg.model.decoder_kwargs.out_dims: Texture = get_class_from_string(cfg.model.texture_type) cfg.model.texture_kwargs[ 'c_dim'] = cfg.model.decoder_kwargs.out_dims.get('latent', 0) texture_decoder = Texture(**cfg.model.texture_kwargs) else: texture_decoder = decoder logger_py.info("Decoder used as NeuralTexture") texture = NeuralTexture( view_dependent=cfg.model.texture_kwargs.view_dependent, decoder=texture_decoder).to(device=device) logger_py.info("Created NeuralTexture {}".format(texture.__class__)) logger_py.info(texture) Model = get_class_from_string("DSS.models.{}_modeling.Model".format( cfg.model.type)) # if not using decoder, then use non-parameterized point renderer # create icosphere as initial point cloud sphere_mesh = ico_sphere(level=4) sphere_mesh.scale_verts_(0.5) points, normals = sample_points_from_meshes( sphere_mesh, num_samples=int(cfg['model']['model_kwargs']['n_points_per_cloud']), return_normals=True) colors = torch.ones_like(points) renderer = create_renderer(cfg.renderer).to(device) model = Model( points, normals, colors, renderer, device=device, texture=texture, **cfg.model.model_kwargs, ).to(device=device) return model
def run(self): deform_verts = torch.full(self.src.verts_packed().shape, 0.0, device=device, requires_grad=True) optimizer = torch.optim.SGD([deform_verts], lr=1.0, momentum=0.9) Niter = 2000 w_chamfer = 1.0 w_edge = 1.0 w_normal = 0.01 w_laplacian = 0.1 for i in range(Niter): optimizer.zero_grad() new_src_mesh = self.src.offset_verts(deform_verts) sampmle_trg = sample_points_from_meshes(self.target, 5000) sample_src = sample_points_from_meshes(new_src_mesh, 5000) loss_chamfer, _ = chamfer_distance(sampmle_trg, sample_src) loss_edge = mesh_edge_loss(new_src_mesh) loss_normal = mesh_normal_consistency(new_src_mesh) loss_laplacian = mesh_laplacian_smoothing(new_src_mesh, method="uniform") #weighted sum of the losses loss = loss_chamfer*w_chamfer + loss_edge*w_edge + loss_normal * w_normal + loss_laplacian * w_laplacian loss.backward() optimizer.step() print('total_loss = %.6f' % loss) self.backwarded.emit(new_src_mesh.verts_packed())
def plot_pointcloud(mesh, title=""): # Sample points uniformly from the surface of the mesh. points = sample_points_from_meshes(mesh, 5000) x, y, z = points.clone().detach().cpu().squeeze().unbind(1) fig = plt.figure(figsize=(5, 5)) ax = Axes3D(fig) ax.scatter3D(x, z, -y) ax.set_xlabel('x') ax.set_ylabel('z') ax.set_zlabel('y') ax.set_title(title) ax.view_init(190, 30) plt.show()
def save_plot3d_mesh_img( mesh, file_path, title = "plot mesh", n_sample = 500, fig_size = (5,5), view_points = (190,30) ): points = sample_points_from_meshes(mesh, n_sample) x, y, z = points.clone().detach().cpu().squeeze().unbind(1) fig = plt.figure(figsize=fig_size) ax = Axes3D(fig) ax.scatter3D(x, z, -y) ax.set_xlabel('x') ax.set_ylabel('z') ax.set_zlabel('y') ax.set_title(title) ax.view_init(view_points[0], view_points[1]) plt.savefig( file_path, dpi = 200, bbox_inches = 'tight' ) return
def test_texture_sampling(self): device = torch.device("cuda:0") batch_size = 6 # verts verts = torch.rand((batch_size, 6, 3), device=device, dtype=torch.float32) verts[:, :3, 2] = 1.0 verts[:, 3:, 2] = -1.0 # textures texts = torch.rand((batch_size, 6, 3), device=device, dtype=torch.float32) # faces faces = torch.tensor([[0, 1, 2], [3, 4, 5]], device=device, dtype=torch.int64) faces = faces.view(1, 2, 3).expand(batch_size, -1, -1) meshes = Meshes(verts=verts, faces=faces, textures=TexturesVertex(texts)) num_samples = 24 samples, normals, textures = sample_points_from_meshes( meshes, num_samples=num_samples, return_normals=True, return_textures=True) textures_naive = torch.zeros((batch_size, num_samples, 3), dtype=torch.float32, device=device) for n in range(batch_size): for i in range(num_samples): p = samples[n, i] if p[2] > 0.0: # sampled from 1st face v0, v1, v2 = verts[n, 0, :2], verts[n, 1, :2], verts[n, 2, :2] w0, w1, w2 = barycentric_coordinates(p[:2], v0, v1, v2) t0, t1, t2 = texts[n, 0], texts[n, 1], texts[n, 2] else: # sampled from 2nd face v0, v1, v2 = verts[n, 3, :2], verts[n, 4, :2], verts[n, 5, :2] w0, w1, w2 = barycentric_coordinates(p[:2], v0, v1, v2) t0, t1, t2 = texts[n, 3], texts[n, 4], texts[n, 5] tt = w0 * t0 + w1 * t1 + w2 * t2 textures_naive[n, i] = tt self.assertClose(textures, textures_naive)
def postprocess(self, batch, device=None): if device is None: device = torch.device("cuda") non_standard_keys = ["points", "normals", "voxels", "id_strs"] # process standard items processed_batch = { key: (value.to(device) if value is not None else None) for key, value in batch.items() if key not in non_standard_keys } # process non-standard items if batch["points"] is not None and batch["normals"] is not None: processed_batch["points"] = batch["points"].to(device) processed_batch["normals"] = batch["normals"].to(device) else: processed_batch["points"], processed_batch["normals"] = \ sample_points_from_meshes( batch["meshes"], num_samples=self.num_samples, return_normals=True ) if batch["voxels"] is not None: if torch.is_tensor(batch["voxels"]): # We used cached voxels on disk, just cast and return processed_batch["voxels"] = batch["voxels"].to(device) # TODO: need to transform voxel grid to all views raise NotImplementedError( "need to transform voxel grid to all views" ) else: # We got a list of voxel_coords, and need to compute voxels on-the-fly voxel_coords = batch["voxels"] voxels = [] for batch_idx, cur_voxel_coords in enumerate(voxel_coords): cur_voxel_coords = cur_voxel_coords.to(device) voxels_views = [] K = batch["intrinsics"][batch_idx] # find voxel grid in all views coordinate frames for view_idx, transform in \ enumerate(batch["extrinsics"][batch_idx].unbind(0)): P = K.matmul(transform) cur_voxels = self._voxelize(cur_voxel_coords, P) voxels_views.append(cur_voxels) voxels.append(torch.stack(voxels_views, dim=0)) processed_batch["voxels"] = torch.stack(voxels, dim=0) if self.return_id_str: processed_batch["id_strs"] = batch["id_strs"] return processed_batch
def forward(self, voxel_scores, meshes_pred, voxels_gt, meshes_gt): """ Args: meshes_pred: Meshes meshes_gt: Either Meshes, or a tuple (points_gt, normals_gt) Returns: loss (float): Torch scalar giving the total loss, or None if an error occured and we should skip this loss. TODO use an exception instead? losses (dict): A dictionary mapping loss names to Torch scalars giving their (unweighted) values. """ # Sample from meshes_gt if we haven't already if isinstance(meshes_gt, tuple): points_gt, normals_gt = meshes_gt else: points_gt, normals_gt = sample_points_from_meshes( meshes_gt, num_samples=self.gt_num_samples, return_normals=True) total_loss = torch.tensor(0.0).to(points_gt) losses = {} if voxel_scores is not None and voxels_gt is not None and self.voxel_weight > 0: voxels_gt = voxels_gt.float() voxel_loss = F.binary_cross_entropy_with_logits( voxel_scores, voxels_gt) total_loss = total_loss + self.voxel_weight * voxel_loss losses["voxel"] = voxel_loss if isinstance(meshes_pred, Meshes): meshes_pred = [meshes_pred] elif meshes_pred is None: meshes_pred = [] # Now assume meshes_pred is a list if not self.skip_mesh_loss: for i, cur_meshes_pred in enumerate(meshes_pred): cur_out = self._mesh_loss(cur_meshes_pred, points_gt, normals_gt) cur_loss, cur_losses = cur_out if total_loss is None or cur_loss is None: total_loss = None else: total_loss = total_loss + cur_loss / len(meshes_pred) for k, v in cur_losses.items(): losses["%s_%d" % (k, i)] = v return total_loss, losses
def compute_geometric_metrics_points(pred_mesh_path, gdth_mesh_path, rot=None): with torch.no_grad(): pred_mesh = trimesh.load(pred_mesh_path) gt_points = np.load(gdth_mesh_path + '/points.npy') gt_normals = np.load(gdth_mesh_path + '/normals.npy') mask = np.random.randint(0, gt_points.shape[0], 10000) gt_points = gt_points[mask, :] gt_normals = gt_normals[mask, :] if rot is not None: r = R.from_euler(rot[0], rot[1], degrees=True) gt_points = r.apply(gt_points) gt_normals = r.apply(gt_normals) # make to unit scale and shift to origin gt_points -= np.mean(gt_points, axis=0) gt_points /= np.max(np.linalg.norm(gt_points, axis=1)) # pred_vertices = pred_mesh.verts_list()[0] pred_vertices = pred_mesh.vertices pred_pts = pred_mesh.sample(10000) pred_vertices -= np.mean(pred_pts, axis=0) pred_vertices /= np.max(np.linalg.norm(pred_vertices, axis=1)) # sph = regular_on_sphere_points(1024) # both = np.vstack((np.vstack((pred_vertices, gt_points)), sph)) # np.savetxt('test.xyz',both) pred_vertices = torch.from_numpy(pred_vertices).float() pred_faces = torch.from_numpy(pred_mesh.faces) pred_mesh = Meshes(verts=[pred_vertices], faces=[pred_faces]) pred_points, pred_normals = sample_points_from_meshes( pred_mesh, num_samples=10000, return_normals=True) pred_points = pred_points pred_normals = pred_normals gt_points = torch.from_numpy(gt_points).type_as(pred_points).unsqueeze( 0) gt_normals = torch.from_numpy(gt_normals).type_as( pred_points).unsqueeze(0) metrics = _compute_sampling_metrics(pred_points, pred_normals, gt_points, gt_normals, eps=1e-8) return metrics
def validation_epoch_end(self, outputs): log_mean = {"log": {}} for k in outputs[0]["log"].keys(): log_mean["log"][k] = torch.stack([x["log"][k] for x in outputs]).mean() log_mean['val_loss'] = torch.stack([x["val_loss"] for x in outputs]).mean() # Compute chamfer Loss if self.cfg.experiment.chamfer_loss and isinstance( self.val_dataset, SynthesizableDataset): assert self.val_dataset.target_mesh is not None, "To compute the " \ "chamfer loss, a target mesh .obj must be provided in the dataset folder" # Target model to query based on the grid model = self.get_model() # Read the input 3D model vertices, faces, _, _, _, _ = extract_geometry( model, self.device, None) # We construct a Meshes structure for the target mesh input_mesh = create_mesh(vertices, faces) # Sparse sampling target_samples = sample_points_from_meshes( self.val_dataset.target_mesh, self.cfg.experiment.chamfer_sampling_size) input_samples = sample_points_from_meshes( input_mesh, self.cfg.experiment.chamfer_sampling_size) chamfer_loss, _ = chamfer_distance(target_samples, input_samples) log_mean["log"]["validation/chamfer_loss"] = chamfer_loss return log_mean
def _mesh_loss(self, meshes_pred, points_gt, normals_gt): """ Args: meshes_pred: Meshes containing N meshes points_gt: Tensor of shape NxPx3 normals_gt: Tensor of shape NxPx3 Returns: total_loss (float): The sum of all losses specific to meshes losses (dict): All (unweighted) mesh losses in a dictionary """ device = meshes_pred.verts_list()[0].device zero = torch.tensor(0.0).to(device) losses = {"chamfer": zero, "normal": zero, "edge": zero} if self.upsample_pred_mesh: points_pred, normals_pred = sample_points_from_meshes( meshes_pred, num_samples=self.pred_num_samples, return_normals=True) else: points_pred = meshes_pred.verts_list() normals_pred = meshes_pred.verts_normals_list() total_loss = torch.tensor(0.0).to(device) if points_pred is None or points_gt is None: # Sampling failed, so return None total_loss = None which = "predictions" if points_pred is None else "GT" logger.info("WARNING: Sampling %s failed" % (which)) return total_loss, losses losses = {} cham_loss, normal_loss = chamfer_distance(points_pred, points_gt, x_normals=normals_pred, y_normals=normals_gt) total_loss = total_loss + self.chamfer_weight * cham_loss total_loss = total_loss + self.normal_weight * normal_loss losses["chamfer"] = cham_loss losses["normal"] = normal_loss edge_loss = mesh_edge_loss(meshes_pred) total_loss = total_loss + self.edge_weight * edge_loss losses["edge"] = edge_loss return total_loss, losses
def create_data(folder_path='meshes/', nb_of_pointclouds=50, nb_of_points=5000, sphere_level=4, normalize_data=True): device = torch.device("cuda:0") data_path = os.path.join(os.getcwd(), folder_path) src_mesh = ico_sphere(sphere_level, device) for filename in os.listdir(data_path): print(f"{datetime.now()} Starting:{filename}") file_path = os.path.join(data_path, filename) cur_mesh = utils.load_mesh(file_path) cur_deform_verts = deformation.get_deform_verts( cur_mesh, nb_of_points, sphere_level) data_verts = np.expand_dims(cur_deform_verts.detach().cpu().numpy(), axis=0) data_input = None data_output = None for _ in range(nb_of_pointclouds): data_a = sample_points_from_meshes( cur_mesh, nb_of_points).squeeze().cpu().numpy() if normalize_data: data_a = data_a - np.mean(data_a, axis=0) data_a = data_a / np.max(data_a, axis=0) data_a_sort_indices = np.argsort(np.linalg.norm(data_a, axis=1)) data_a = data_a[data_a_sort_indices] data_a = np.expand_dims(data_a, axis=0) data_input = data_a if data_input is None else np.concatenate( (data_input, data_a)) data_output = data_verts if data_output is None else np.concatenate( (data_output, data_verts)) np.save(f'data/{os.path.splitext(filename)[0]}_input.npy', data_input) np.save(f'data/{os.path.splitext(filename)[0]}_output.npy', data_output) deformed_mesh = src_mesh.offset_verts(cur_deform_verts) final_verts, final_faces = deformed_mesh.get_mesh_verts_faces(0) final_obj = os.path.join( 'deformed_meshes/', f'{os.path.splitext(filename)[0]}_deformed.obj') save_obj(final_obj, final_verts, final_faces) print( f"{datetime.now()} Finished:{filename}, Point Cloud Shape:{data_input.shape} Deform Verts Shape:{data_output.shape}" )
def validate_training_AE(validation_generator, model): ''' This function is used to calculate validation loss during training NMF AE ''' print("Validating model......") with torch.no_grad(): total_loss = 0 items = 0 for input,_,_ in validation_generator: input = input.cuda() _, _, pred2, face = model(input) # Point prediction after each deform block and face information (refer figure 4 in paper) mesh_p2 = Meshes(verts = pred2, faces = face) # Construct Differentiable mesh M_p2 pts2 = sample_points_from_meshes(mesh_p2,num_samples=2562) # Differentiably sample random points from mesh surfaces loss,_ = chamfer_distance(pts2, input) total_loss+=loss.item() items+=1 return total_loss/items # Return average validation loss
def test_outputs(self): for add_texture in (True, False): meshes = TestSamplePoints.init_meshes( device=torch.device("cuda:0"), add_texture=add_texture ) out1 = sample_points_from_meshes(meshes, num_samples=100) self.assertTrue(torch.is_tensor(out1)) out2 = sample_points_from_meshes( meshes, num_samples=100, return_normals=True ) self.assertTrue(isinstance(out2, tuple) and len(out2) == 2) if add_texture: out3 = sample_points_from_meshes( meshes, num_samples=100, return_textures=True ) self.assertTrue(isinstance(out3, tuple) and len(out3) == 2) out4 = sample_points_from_meshes( meshes, num_samples=100, return_normals=True, return_textures=True ) self.assertTrue(isinstance(out4, tuple) and len(out4) == 3) else: with self.assertRaisesRegex( ValueError, "Meshes do not contain textures." ): sample_points_from_meshes( meshes, num_samples=100, return_textures=True ) with self.assertRaisesRegex( ValueError, "Meshes do not contain textures." ): sample_points_from_meshes( meshes, num_samples=100, return_normals=True, return_textures=True, )