def _visualize(self, images, sample, output, state, **kwargs): img_id, shape, view, width, name = state if 'colors' in output and output['colors'] is not None: images['{}_color/{}:HWC'.format(name, img_id)] = { 'img': output['colors'][shape, view] } if 'depths' in output and output['depths'] is not None: min_depth, max_depth = output['depths'].min( ), output['depths'].max() images['{}_depth/{}:HWC'.format(name, img_id)] = { 'img': output['depths'][shape, view], 'min_val': min_depth, 'max_val': max_depth } normals = compute_normal_map( output['ray_start'][shape, view].float(), output['ray_dir'][shape, view].float(), output['depths'][shape, view].float(), sample['extrinsics'][shape, view].float().inverse(), width) images['{}_normal/{}:HWC'.format(name, img_id)] = { 'img': normals, 'min_val': -1, 'max_val': 1 } return images
def _visualize(self, images, sample, output, state, **kwargs): img_id, shape, view, width, name = state images = super()._visualize(images, sample, output, state, **kwargs) if 'voxel_edges' in output and output['voxel_edges'] is not None: # voxel hitting visualization images['{}_voxel/{}:HWC'.format(name, img_id)] = { 'img': output['voxel_edges'][shape, view].float(), 'min_val': 0, 'max_val': 1, 'weight': compute_normal_map( sample['ray_start'][shape, view].float(), sample['ray_dir'][shape, view].float(), output['voxel_depth'][shape, view].float(), sample['extrinsics'][shape, view].float().inverse(), width, proj=True) } if 'feat_n2' in output and output['feat_n2'] is not None: images['{}_featn2/{}:HWC'.format(name, img_id)] = { 'img': output['feat_n2'][shape, view].float(), 'min_val': 0, 'max_val': 1 } return images
def _visualize(self, images, sample, output, state, **kwargs): img_id, shape, view, width, name = state if 'colors' in output and output['colors'] is not None: images['{}_color/{}:HWC'.format(name, img_id)] = { 'img': output['colors'][shape, view], 'min_val': float(self.args.min_color) } if 'depths' in output and output['depths'] is not None: min_depth, max_depth = output['depths'].min( ), output['depths'].max() if getattr(self.args, "near", None) is not None: min_depth = self.args.near max_depth = self.args.far images['{}_depth/{}:HWC'.format(name, img_id)] = { 'img': output['depths'][shape, view], 'min_val': min_depth, 'max_val': max_depth } normals = compute_normal_map( sample['ray_start'][shape, view].float(), sample['ray_dir'][shape, view].float(), output['depths'][shape, view].float(), sample['extrinsics'][shape, view].float().inverse(), width) images['{}_normal/{}:HWC'.format(name, img_id)] = { 'img': normals, 'min_val': -1, 'max_val': 1 } # generate point clouds from depth images['{}_point/{}'.format(name, img_id)] = { 'img': torch.cat([ ray(sample['ray_start'][shape, view].float(), sample['ray_dir'][shape, view].float(), output['depths'][shape, view].unsqueeze(-1).float()), (output['colors'][shape, view] - self.args.min_color) / (1 - self.args.min_color) ], 1), # XYZRGB 'raw': True } if 'z' in output and output['z'] is not None: images['{}_z/{}:HWC'.format(name, img_id)] = { 'img': output['z'][shape, view], 'min_val': 0, 'max_val': 1 } if 'normal' in output and output['normal'] is not None: images['{}_predn/{}:HWC'.format(name, img_id)] = { 'img': output['normal'][shape, view], 'min_val': -1, 'max_val': 1 } return images
def add_eval_scores(self, logging_output, sample, output, criterion, scores=['ssim', 'psnr', 'lpips'], outdir=None): predicts, targets = output['colors'], sample['colors'] ssims, psnrs, lpips, rmses = [], [], [], [] for s in range(predicts.size(0)): for v in range(predicts.size(1)): width = int(sample['size'][s, v][1]) p = recover_image(predicts[s, v], width=width, min_val=float(self.args.min_color)) t = recover_image(targets[s, v], width=width, min_val=float(self.args.min_color)) pn, tn = p.numpy(), t.numpy() p, t = p.to(predicts.device), t.to(targets.device) if 'ssim' in scores: ssims += [skimage.metrics.structural_similarity(pn, tn, multichannel=True, data_range=1)] if 'psnr' in scores: psnrs += [skimage.metrics.peak_signal_noise_ratio(pn, tn, data_range=1)] if 'lpips' in scores and hasattr(criterion, 'lpips'): with torch.no_grad(): lpips += [criterion.lpips( 2 * p.unsqueeze(-1).permute(3,2,0,1) - 1, 2 * t.unsqueeze(-1).permute(3,2,0,1) - 1).item()] if 'depths' in sample: td = sample['depths'][sample['depths'] > 0] pd = output['depths'][sample['depths'] > 0] rmses += [torch.sqrt(((td - pd) ** 2).mean()).item()] if outdir is not None: def imsave(filename, image): imageio.imsave(os.path.join(outdir, filename), (image * 255).astype('uint8')) figname = '-{:03d}_{:03d}.png'.format(sample['id'][s], sample['view'][s, v]) imsave('output' + figname, pn) imsave('target' + figname, tn) imsave('normal' + figname, recover_image(compute_normal_map( sample['ray_start'][s, v].float(), sample['ray_dir'][s, v].float(), output['depths'][s, v].float(), sample['extrinsics'][s, v].float().inverse(), width=width), min_val=-1, max_val=1, width=width).numpy()) if 'featn2' in output: imsave('featn2' + figname, output['featn2'][s, v].cpu().numpy()) if 'voxel' in output: imsave('voxel' + figname, output['voxel'][s, v].cpu().numpy()) if len(ssims) > 0: logging_output['ssim_loss'] = np.mean(ssims) if len(psnrs) > 0: logging_output['psnr_loss'] = np.mean(psnrs) if len(lpips) > 0: logging_output['lpips_loss'] = np.mean(lpips) if len(rmses) > 0: logging_output['rmses_loss'] = np.mean(rmses)