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)