def visualize(self, data): ''' Performs a visualization step for the data. Args: data (dict): data dictionary ''' device = self.device batch_size = data['points'].size(0) inputs = data.get('inputs', torch.empty(batch_size, 0)).to(device) shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device) p = p.expand(batch_size, *p.size()) kwargs = {} with torch.no_grad(): p_r = self.model(p, inputs, sample=self.eval_sample, **kwargs) occ_hat = p_r.probs.view(batch_size, *shape) voxels_out = (occ_hat >= self.threshold).cpu().numpy() for i in trange(batch_size): input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i) vis.visualize_data(inputs[i].cpu(), self.input_type, input_img_path) vis.visualize_voxels(voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
def export_points(mesh, modelname, loc, scale, args): if not mesh.is_watertight: print('Warning: mesh %s is not watertight!' 'Cannot sample points.' % modelname) return filename = os.path.join(args.points_folder, modelname + '.npz') if not args.overwrite and os.path.exists(filename): print('Points already exist: %s' % filename) return # uniform grid total_l = 1. + args.points_padding res = args.points_resolution points_uniform = make_3d_grid((-total_l/2.0 + total_l/(res*2),)*3, (total_l/2.0 - total_l/(res*2),)*3, (res,)*3).numpy() points_uniform_occupancies = check_mesh_contains(mesh, points_uniform).astype(np.float32) points = points_uniform occupancies = points_uniform_occupancies # Compress if args.float16: dtype = np.float16 else: dtype = np.float32 points = points.astype(dtype) occupancies = occupancies.astype(dtype) print('Writing points: %s' % filename) np.savez(filename, points=points, occupancies=occupancies, loc=loc, scale=scale)
def visualize(self, data): """ Performs a visualization step for the data. Args: data (dict): data dictionary """ batch_size = data["points"].shape[0] inputs = data.get("inputs", tf.zeros([batch_size, 0])) shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape) # CHECK p = tf.broadcast_to(p, [batch_size, *p.shape]) kwargs = {} p_r = self.model(p, inputs, sample=self.eval_sample, training=False, **kwargs) occ_hat = tf.reshape(p_r.probs_parameter(), [batch_size, *shape]) voxels_out = (occ_hat >= self.threshold).numpy() for i in trange(batch_size): input_img_path = os.path.join(self.vis_dir, "%03d_in.png" % i) vis.visualize_data(inputs[i], self.input_type, input_img_path) vis.visualize_voxels(voxels_out[i], os.path.join(self.vis_dir, "%03d.png" % i))
def generate_from_latent(self, z, c_first=None, Rt=None, K=None, stats_dict={}, **kwargs): ''' Generates mesh from latent. Args: z (tensor): latent code z c (tensor): latent conditioned code c stats_dict (dict): stats dictionary ''' threshold = np.log(self.threshold) - np.log(1. - self.threshold) t0 = time.time() # Compute bounding box size box_size = 1 + self.padding # Shortcut if self.upsampling_steps == 0: nx = self.resolution0 pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3, (nx, ) * 3) values = self.eval_points(pointsf, z, c_first, **kwargs).cpu().numpy() value_grid = values.reshape(nx, nx, nx) else: mesh_extractor = MISE(self.resolution0, self.upsampling_steps, threshold) points = mesh_extractor.query() while points.shape[0] != 0: # Query points pointsf = torch.FloatTensor(points).to(self.device) # Normalize to bounding box pointsf = pointsf / mesh_extractor.resolution pointsf = box_size * (pointsf - 0.5) # Evaluate model and update values = self.eval_points(pointsf, z, c_first, Rt, K, **kwargs).cpu().numpy() values = values.astype(np.float64) mesh_extractor.update(points, values) points = mesh_extractor.query() value_grid = mesh_extractor.to_dense() # Extract mesh stats_dict['time (eval points)'] = time.time() - t0 mesh = self.extract_mesh(value_grid, z, c_first, Rt, K, stats_dict=stats_dict) return mesh
def visualize(self, data, it=0, vis_type='mesh'): ''' Visualized the data. Args: data (dict): data dictionary it (int): training iteration vis_type (string): visualization type ''' if self.multi_gpu: print( "Sorry, visualizations currently not implemented when using \ multi GPU training.") return 0 device = self.device inputs = data.get('inputs', torch.empty(1, 0)).to(device) batch_size = inputs.shape[0] c = self.model.encode_inputs(inputs) if vis_type == 'voxel': shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device) p = p.unsqueeze(0).repeat(batch_size, 1, 1) with torch.no_grad(): p_r = self.model.decode(p, c=c).probs voxels_out = (p_r >= self.threshold).cpu().numpy() voxels_out = voxels_out.reshape(batch_size, 32, 32, 32) for i in range(batch_size): out_file = os.path.join(self.vis_dir, '%03d.png' % i) vis.visualize_voxels(voxels_out[i], out_file) elif vis_type == 'pointcloud': p = torch.rand(batch_size, 60000, 3).to(device) - 0.5 with torch.no_grad(): occ = self.model.decode(p, c=c).probs mask = occ > self.threshold for i in range(batch_size): pi = p[i][mask[i]].cpu() out_file = os.path.join(self.vis_dir, '%03d.png' % i) vis.visualize_pointcloud(pi, out_file=out_file) elif vis_type == 'mesh': try: mesh_list = self.generator.generate_meshes( data, return_stats=False) for i, mesh in tqdm(enumerate(mesh_list)): if self.overwrite_visualization: ending = '' else: ending = '_%010d' % it mesh_out_file = os.path.join( self.vis_dir, '%03d%s.ply' % (i, ending)) mesh.export(mesh_out_file) except Exception as e: print("Exception occurred during visualization: ", e) else: print('The visualization type %s is not valid!' % vis_type)
def voxelize_interior(mesh, resolution): shape = (resolution, ) * 3 bb_min = (0.5, ) * 3 bb_max = (resolution - 0.5, ) * 3 # Create points. Add noise to break symmetry points = make_3d_grid(bb_min, bb_max, shape=shape).numpy() points = points + 0.1 * (np.random.rand(*points.shape) - 0.5) points = (points / resolution - 0.5) occ = check_mesh_contains(mesh, points) occ = occ.reshape(shape) return occ
def visualize(self, data, it=0., epoch_it=0.): ''' Performs a visualization step for the data. Args: data (dict): data dictionary ''' device = self.device batch_size = data['points'].size(0) inputs = data.get('inputs', torch.empty(batch_size, 0)).to(device) angles = data.get('angles').to(device) shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device) p = p.expand(batch_size, *p.size()) kwargs = {} with torch.no_grad(): _, _, sgn, _, _ = self.model(p * self.pnet_point_scale, inputs, sample=self.eval_sample, angles=angles, **kwargs) if self.is_sdf: if self.is_logits_by_min: logits = (sgn.min(1)[0] <= 0).float() else: raise NotImplementedError else: if self.is_logits_by_max: logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0]) elif self.is_logits_by_sign_filter: positive = torch.relu(sgn).sum(1) negative = torch.relu(-sgn).sum(1) logits = torch.where(positive >= negative, positive, -negative) else: logits = convert_tsd_range_to_zero_to_one(sgn).sum(1) occ_hat = logits.view(batch_size, *shape) voxels_out = (occ_hat >= self.threshold).cpu().numpy() input_images = [] voxels_images = [] for i in trange(batch_size): input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i) vis.visualize_data(inputs[i].cpu(), self.input_type, input_img_path) vis.visualize_voxels(voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
def visualize(self, data): ''' Performs a visualization step for the data. Args: data (dict): data dictionary ''' device = self.device batch_size = data['points'].size(0) inputs = data.get('inputs').to(device) #gt_depth_maps = data.get('inputs.depth').to(device) gt_mask = data.get('inputs.mask').to(device).byte() shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device) p = p.expand(batch_size, *p.size()) kwargs = {} with torch.no_grad(): pr_depth_maps = self.model.predict_depth_map(inputs) background_setting(pr_depth_maps, gt_mask) p_r = self.model.forward_halfway(p, pr_depth_maps, sample=self.eval_sample, **kwargs) occ_hat = p_r.probs.view(batch_size, *shape) voxels_out = (occ_hat >= self.threshold).cpu().numpy() for i in trange(batch_size): input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i) vis.visualize_data(inputs[i].cpu(), 'img', input_img_path) vis.visualize_voxels(voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i)) depth_map_path = os.path.join(self.vis_dir, '%03d_pr_depth.png' % i) depth_map = pr_depth_maps[i].cpu() depth_map = depth_to_L(depth_map, gt_mask[i].cpu()) vis.visualize_data(depth_map, 'img', depth_map_path)
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) occ = data.get('points.occ').to(device) inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) occ_iou = data.get('points_iou.occ').to(device) angles = data.get('angles').to(device) kwargs = {} """ with torch.no_grad(): elbo, rec_error, kl = self.model.compute_elbo( points, occ, inputs, **kwargs) eval_dict['loss'] = -elbo.mean().item() eval_dict['rec_error'] = rec_error.mean().item() eval_dict['kl'] = kl.mean().item() """ # Compute iou batch_size = points.size(0) with torch.no_grad(): _, _, sgn, _, _ = self.model(points_iou * self.pnet_point_scale, inputs, sample=self.eval_sample, angles=angles, **kwargs) if self.is_sdf: if self.is_logits_by_min: logits = (sgn.min(1)[0] <= 0).float() else: raise NotImplementedError else: if self.is_eval_logits_by_max: logits = (sgn >= 0.).float().max(1)[0] elif self.is_logits_by_max: logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0]) elif self.is_logits_by_sign_filter: positive = torch.relu(sgn).sum(1) negative = torch.relu(-sgn).sum(1) logits = torch.where(positive >= negative, positive, -negative) else: logits = convert_tsd_range_to_zero_to_one(sgn).sum(1) occ_iou_np = (occ_iou >= self.threshold).cpu().numpy() occ_iou_hat_np = (logits >= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou eval_dict['overlap'] = (torch.where(sgn >= 0, torch.ones_like(sgn), torch.zeros_like(sgn)).sum(1) > 1).sum().detach().cpu().numpy().item() eval_dict['overlap_mean'] = ( torch.where(sgn >= 0, torch.ones_like(sgn), torch.zeros_like(sgn)).sum(1) > 1).float().mean().detach().cpu().numpy().item() # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): _, _, sgn, _, _ = self.model(points_voxels * self.pnet_point_scale, inputs, sample=self.eval_sample, angles=angles, **kwargs) if self.is_sdf: if self.is_logits_by_min: logits = (sgn.min(1)[0] <= 0).float() else: raise NotImplementedError else: if self.is_eval_logits_by_max: logits = (sgn >= 0.).float().max(1)[0] elif self.is_logits_by_max: logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0]) elif self.is_logits_by_sign_filter: positive = torch.relu(sgn).sum(1) negative = torch.relu(-sgn).sum(1) logits = torch.where(positive >= negative, positive, -negative) else: logits = convert_tsd_range_to_zero_to_one(sgn).sum(1) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = (logits >= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) occ = data.get('points.occ').to(device) inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) occ_iou = data.get('points_iou.occ').to(device) world_mat = data.get('inputs.world_mat').to(device) camera_mat = data.get('inputs.camera_mat').to(device) camera_args = common.get_camera_args(data, 'points.loc', 'points.scale', device=self.device) world_mat, camera_mat = camera_args['Rt'], camera_args['K'] kwargs = {} with torch.no_grad(): elbo, rec_error, kl = self.model.compute_elbo( points, occ, inputs, world_mat, camera_mat, **kwargs) eval_dict['loss'] = -elbo.mean().item() eval_dict['rec_error'] = rec_error.mean().item() eval_dict['kl'] = kl.mean().item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out = self.model(points_iou, inputs, world_mat, camera_mat, sample=self.eval_sample, **kwargs) occ_iou_np = (occ_iou >= 0.5).cpu().numpy() occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out = self.model(points_voxels, inputs, world_mat, camera_mat, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = (p_out.probs >= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} points = data.get('points').to(device) sdf = data.get('points.sdf').to(device) inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) sdf_iou = data.get('points_iou.sdf').to(device) kwargs = {} # Compute loss with torch.no_grad(): p_out = self.model(points, inputs, sample=self.eval_sample, **kwargs) loss = get_sdf_loss(p_out.logits, sdf, self.loss_type, ratio=self.sdf_ratio) loss = loss.sum(-1).mean() eval_dict['loss'] = loss.item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out = self.model(points_iou, inputs, sample=self.eval_sample, **kwargs) occ_iou_np = (sdf_iou <= 0.).cpu().numpy() occ_iou_pred_sdf = p_out.logits / self.sdf_ratio occ_iou_hat_np = (occ_iou_pred_sdf <= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid( (-0.5 + 1/64,) * 3, (0.5 - 1/64,) * 3, (32,) * 3) points_voxels = points_voxels.expand( batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out = self.model(points_voxels, inputs, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_pred_sdf = p_out.logits / self.sdf_ratio occ_hat_np = (occ_hat_pred_sdf <= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def visualize(self, data): ''' Performs a visualization step for the data. Args: data (dict): data dictionary ''' device = self.device batch_size = data['points'].size(0) shape = (32, 32, 32) p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device) p = p.expand(batch_size, *p.size()) encoder_inputs, raw_data = compose_inputs( data, mode='val', device=self.device, input_type=self.input_type, use_gt_depth_map=self.use_gt_depth_map, depth_map_mix=self.depth_map_mix, with_img=self.with_img, depth_pointcloud_transfer=self.depth_pointcloud_transfer, local=self.local) kwargs = {} with torch.no_grad(): p_r = self.model.forward_halfway(p, encoder_inputs, sample=self.eval_sample, **kwargs) occ_hat = p_r.probs.view(batch_size, *shape) voxels_out = (occ_hat >= self.threshold).cpu().numpy() # visualize if self.local: encoder_inputs = encoder_inputs[None] if self.input_type == 'depth_pred': gt_mask = raw_data['mask'] if self.with_img: encoder_inputs = encoder_inputs['depth'] for i in trange(batch_size): if self.use_gt_depth_map: input_img_path = os.path.join(self.vis_dir, '%03d_in_gt.png' % i) else: input_img_path = os.path.join(self.vis_dir, '%03d_in_pr.png' % i) depth_map = encoder_inputs[i].cpu() depth_map = depth_to_L(depth_map, gt_mask[i].cpu()) vis.visualize_data(depth_map, 'img', input_img_path) vis.visualize_voxels( voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i)) elif self.input_type == 'depth_pointcloud': for i in trange(batch_size): input_pointcloud_file = os.path.join( self.vis_dir, '%03d_depth_pointcloud.png' % i) pc = encoder_inputs[i].cpu() if self.depth_pointcloud_transfer in ('view', 'view_scale_model'): vis.visualize_pointcloud(pc, out_file=input_pointcloud_file, elev=15, azim=180) else: vis.visualize_pointcloud(pc, out_file=input_pointcloud_file) vis.visualize_voxels( voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) occ = data.get('points.occ').to(device) voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) occ_iou = data.get('points_iou.occ').to(device) kwargs = {} encoder_inputs, _ = compose_inputs( data, mode='val', device=self.device, input_type=self.input_type, use_gt_depth_map=self.use_gt_depth_map, depth_map_mix=self.depth_map_mix, with_img=self.with_img, depth_pointcloud_transfer=self.depth_pointcloud_transfer, local=self.local) with torch.no_grad(): elbo, rec_error, kl = self.model.compute_elbo_halfway( points, occ, encoder_inputs, **kwargs) eval_dict['loss'] = -elbo.mean().item() eval_dict['rec_error'] = rec_error.mean().item() eval_dict['kl'] = kl.mean().item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out = self.model.forward_halfway(points_iou, encoder_inputs, sample=self.eval_sample, **kwargs) occ_iou_np = (occ_iou >= 0.5).cpu().numpy() occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out = self.model.forward_halfway(points_voxels, encoder_inputs, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = (p_out.probs >= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) occ = data.get('points.occ').to(device) inputs = data.get('inputs').to(device) #gt_depth_maps = data.get('inputs.depth').to(device) gt_mask = data.get('inputs.mask').to(device).byte() voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) occ_iou = data.get('points_iou.occ').to(device) kwargs = {} with torch.no_grad(): elbo, rec_error, kl = self.model.compute_elbo( points, occ, inputs, gt_mask, **kwargs) eval_dict['loss'] = -elbo.mean().item() eval_dict['rec_error'] = rec_error.mean().item() eval_dict['kl'] = kl.mean().item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out = self.model(points_iou, inputs, gt_mask, sample=self.eval_sample, **kwargs) occ_iou_np = (occ_iou >= 0.5).cpu().numpy() occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out = self.model(points_voxels, inputs, gt_mask, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = (p_out.probs >= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def eval_step(self, data): """ Performs an evaluation step. Args: data (dict): data dictionary """ threshold = self.threshold eval_dict = {} # Compute elbo points = data.get("points") occ = data.get("points.occ") inputs = data.get("inputs", tf.zeros([points.shape[0], 0])) voxels_occ = data.get("voxels") points_iou = data.get("points_iou") occ_iou = data.get("points_iou.occ") kwargs = {} elbo, rec_error, kl = self.model.compute_elbo(points, occ, inputs, training=False, **kwargs) eval_dict["loss"] = -float(tf.reduce_mean(elbo)) eval_dict["rec_error"] = float(tf.reduce_mean(rec_error)) eval_dict["kl"] = float(tf.reduce_mean(kl)) # Compute iou batch_size = points.shape[0] p_out = self.model(points_iou, inputs, sample=self.eval_sample, training=False, **kwargs) occ_iou_np = (occ_iou >= 0.5).numpy() occ_iou_hat_np = (p_out.probs_parameter() >= threshold).numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict["iou"] = float(iou) # Estimate voxel iou if voxels_occ is not None: points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = tf.broadcast_to(points_voxels, [batch_size, *points_voxels.shape]) p_out = self.model(points_voxels, inputs, sample=self.eval_sample, training=False, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).numpy() occ_hat_np = (p_out.probs_parameter() >= threshold).numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict["iou_voxels"] = float(iou_voxels) return eval_dict
def generate_from_latent(self, z, c=None, stats_dict={}, data=None, **kwargs): ''' Generates mesh from latent. Args: z (tensor): latent code z c (tensor): latent conditioned code c stats_dict (dict): stats dictionary ''' assert data is not None if self.is_explicit_mesh: normal_faces = data.get('angles.normal_face').to(self.device) normal_angles = data.get('angles.normal_angles').to(self.device) if self.is_skip_surface_mask_generation_time: t0 = time.time() output = self.model.decode(None, z, c, angles=normal_angles, only_return_points=True, **kwargs) stats_dict['time (eval points)'] = time.time() - t0 t0 = time.time() if not self.is_just_measuring_time: output = self.model.decode(None, z, c, angles=normal_angles, **kwargs) if not self.is_skip_surface_mask_generation_time: stats_dict['time (eval points)'] = time.time() - t0 normal_vertices, normal_mask, _, _, _ = output t0 = time.time() B, N, P, D = normal_vertices.shape normal_faces_all = torch.cat([(normal_faces + idx * P) for idx in range(N)], axis=1) assert B == 1 mem_t = time.time() verts = normal_vertices.view( N * P, D).to('cpu').detach().numpy() / self.pnet_point_scale faces = normal_faces_all.view(-1, 3).to('cpu').detach().numpy() if self.is_just_measuring_time: visbility = None else: visbility = (normal_mask > 0.5).view( N * P).to('cpu').detach().numpy() skip_t = time.time() - mem_t mesh = trimesh.Trimesh( verts, faces, process=False, vertex_attributes={'vertex_visibility': visbility}) stats_dict['time (copy to trimesh)'] = time.time() - t0 - skip_t else: t0 = time.time() threshold = 0. #threshold = np.log(self.threshold) - np.log(1. - self.threshold) # Compute bounding box size box_size = 1 + self.padding # Shortcut if self.upsampling_steps == 0: nx = self.resolution0 pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3, (nx, ) * 3) values = self.eval_points(pointsf, z, c, data=data, **kwargs).cpu().numpy() value_grid = values.reshape(nx, nx, nx) else: mesh_extractor = MISE(self.resolution0, self.upsampling_steps, threshold) points = mesh_extractor.query() while points.shape[0] != 0: # Query points pointsf = torch.FloatTensor(points).to(self.device) # Normalize to bounding box pointsf = pointsf / mesh_extractor.resolution pointsf = box_size * (pointsf - 0.5) # Evaluate model and update values = self.eval_points(pointsf, z, c, data=data, **kwargs).cpu().numpy() values = values.astype(np.float64) mesh_extractor.update(points, values) points = mesh_extractor.query() value_grid = mesh_extractor.to_dense() # Extract mesh stats_dict['time (eval points)'] = time.time() - t0 mesh = self.extract_mesh(value_grid, z, c, stats_dict=stats_dict) return mesh
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) #occ = data.get('points.occ').to(device) inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) voxels_occ = data.get('voxels') #occ_iou = data.get('points_iou.occ').to(device) kwargs = {} #with torch.no_grad(): # elbo, rec_error, kl = self.model.compute_elbo( # points, occ, inputs, **kwargs) target = data.get('points.point_lab').to(device) # eval_dict['rec_error'] = rec_error.mean().item() # eval_dict['kl'] = kl.mean().item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out, vote = self.model(points, inputs, sample=self.eval_sample, **kwargs) instance_loss = True if instance_loss: centers = data.get('points.centers').to(device) vote_loss = (torch.max(vote.float() - centers.T.float(), torch.tensor([0.]).cuda())**2).sum() logits = p_out.logits if self.loss_type == 'cross_entropy': loss_i = F.cross_entropy(logits, target, reduction='none') occ_iou_np = (target >= 0).cpu().numpy() occ_iou_hat_np = p_out.probs.argmax(dim=1).cpu().numpy() #iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() #occ_iou_hat_np_04 = (p_out.probs >= 0.4).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou #loss = loss_i.sum(-1).mean() loss = loss_i.sum(-1).mean() + vote_loss * 1000 eval_dict['loss'] = loss.cpu().numpy() # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out, vote = self.model(points_voxels, inputs, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = p_out.probs.argmax(dim=1).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict
def generate_from_conditional(self, c, stats_dict={}, **kwargs): ''' Generates mesh from latent. Args: c (tensor): latent conditioned code c stats_dict (dict): stats dictionary ''' threshold = np.log(self.threshold) - np.log(1. - self.threshold) t0 = time.time() # Compute bounding box size box_size = 1 + self.padding # Shortcut if self.upsampling_steps == 0: nx = self.resolution0 pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3, (nx, ) * 3).to(self.device) out_dict = self.eval_points(pointsf, c, **kwargs) if self.tf_type is None: value_minimal_grid = out_dict['values_minimal'].detach().cpu( ).numpy().reshape([nx, nx, nx]) value_cloth_grid = out_dict['values_cloth'].detach().cpu( ).numpy().reshape([nx, nx, nx]) label_grid = out_dict['labels'].detach().cpu().numpy().reshape( [nx, nx, nx]) p_hat_grid = None else: value_minimal_grid = out_dict['values_minimal'].detach().cpu( ).numpy().reshape([nx, nx, nx]) value_cloth_grid = out_dict['values_cloth'].detach().cpu( ).numpy().reshape([nx, nx, nx]) label_grid = out_dict['labels'].detach().cpu().numpy().reshape( [nx, nx, nx]) p_hat_grid = out_dict['p_hats'].detach().cpu().numpy().reshape( [nx, nx, nx, 3]) else: raise ValueError('We do not support MISE for double layer') # Extract mesh stats_dict['time (eval points)'] = time.time() - t0 mesh_all = {} # Generate body under cloth mesh_minimal = self.extract_mesh(value_minimal_grid, label_grid, p_hat_grid, c, stats_dict=stats_dict, **kwargs) for k, v in mesh_minimal.items(): mesh_all['minimal_' + k] = v # Generate body with cloth mesh_cloth = self.extract_mesh(value_cloth_grid, label_grid, p_hat_grid, c, stats_dict=stats_dict, **kwargs) for k, v in mesh_cloth.items(): mesh_all['cloth_' + k] = v return mesh_all
def eval_step(self, data): ''' Performs an evaluation step. Args: data (dict): data dictionary ''' self.model.eval() device = self.device threshold = self.threshold eval_dict = {} # Compute elbo points = data.get('points').to(device) occ = data.get('points.occ').to(device) inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device) voxels_occ = data.get('voxels') points_iou = data.get('points_iou').to(device) occ_iou = data.get('points_iou.occ').to(device) '''TRANSFORMATION TO CAMERA SPACE''' if self.camera_space: transform = data.get('inputs.world_mat').to(device) R = transform[:, :, :3] points = transform_points(points, transform) points_iou = transform_points(points_iou, R) '''END''' kwargs = {} with torch.no_grad(): elbo, rec_error, kl = self.model.compute_elbo( points, occ, inputs, **kwargs) eval_dict['loss'] = -elbo.mean().item() eval_dict['rec_error'] = rec_error.mean().item() eval_dict['kl'] = kl.mean().item() # Compute iou batch_size = points.size(0) with torch.no_grad(): p_out = self.model(points_iou, inputs, sample=self.eval_sample, **kwargs) occ_iou_np = (occ_iou >= 0.5).cpu().numpy() occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy() iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean() eval_dict['iou'] = iou ''' with torch.no_grad(): p_out_r = self.model(points_iou_r, inputs, sample=self.eval_sample, **kwargs) occ_iou_r_hat_np = (p_out_r.probs >= threshold).cpu().numpy() iou_r = compute_iou(occ_iou_np, occ_iou_r_hat_np).mean() eval_dict['iou_r'] = iou_r data['iou_r'] = iou_r data['iou'] = iou data['occ_iou_r_hat_np'] = occ_iou_r_hat_np data['occ_iou_hat_np'] = occ_iou_hat_np SAVE IOU DATA # Create npz and save im_path = str(data.get('inputs.image_path')[0]).split('/') model=im_path[3] im_nr=im_path[5].split('.')[0] out_dir = 'IOU' np.savez(os.path.join(out_dir, model+'_'+im_nr), iou_r=iou_r, occ_iou_hat_np=occ_iou_hat_np, occ_iou_r_hat_np=occ_iou_r_hat_np, iou=iou, occ_iou=occ_iou.cpu().numpy(), transform=transform.cpu().numpy(), points_iou=points_iou.cpu().numpy(), points_iou_r=points_iou_r.cpu().numpy(), model=model, im_nr=im_nr, image=inputs.cpu().numpy() ) END''' # Estimate voxel iou if voxels_occ is not None: voxels_occ = voxels_occ.to(device) points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3, (0.5 - 1 / 64, ) * 3, (32, ) * 3) points_voxels = points_voxels.expand(batch_size, *points_voxels.size()) points_voxels = points_voxels.to(device) with torch.no_grad(): p_out = self.model(points_voxels, inputs, sample=self.eval_sample, **kwargs) voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy() occ_hat_np = (p_out.probs >= threshold).cpu().numpy() iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean() eval_dict['iou_voxels'] = iou_voxels return eval_dict