os.makedirs(obj_out_dir, exist_ok=True)
for i in range(n_gen_views):
out_file = os.path.join(
obj_out_dir, "{:06}.png".format(novel_view_idxs[i].item())
)
imageio.imwrite(out_file, (all_rgb[i] * 255).astype(np.uint8))
if args.write_depth:
out_depth_file = os.path.join(
obj_out_dir, "{:06}_depth.exr".format(novel_view_idxs[i].item())
)
out_depth_norm_file = os.path.join(
obj_out_dir,
"{:06}_depth_norm.png".format(novel_view_idxs[i].item()),
)
depth_cmap_norm = util.cmap(all_depth[i])
cv2.imwrite(out_depth_file, all_depth[i])
imageio.imwrite(out_depth_norm_file, depth_cmap_norm)
curr_ssim = 0.0
curr_psnr = 0.0
if not args.no_compare_gt:
images_0to1 = images * 0.5 + 0.5 # (NV, 3, H, W)
images_gt = images_0to1[target_view_mask]
rgb_gt_all = (
images_gt.permute(0, 2, 3, 1).contiguous().numpy()
) # (NV-NS, H, W, 3)
for view_idx in range(n_gen_views):
ssim = skimage.measure.compare_ssim(
all_rgb[view_idx],
rgb_gt_all[view_idx],
def vis_step(self, data, global_step, idx=None):
if "images" not in data:
return {}
if idx is None:
batch_idx = np.random.randint(0, data["images"].shape[0])
else:
print(idx)
batch_idx = idx
images = data["images"][batch_idx].to(device=device) # (NV, 3, H, W)
poses = data["poses"][batch_idx].to(device=device) # (NV, 4, 4)
focal = data["focal"][batch_idx:batch_idx + 1] # (1)
c = data.get("c")
if c is not None:
c = c[batch_idx:batch_idx + 1] # (1)
NV, _, H, W = images.shape
cam_rays = util.gen_rays(poses,
W,
H,
focal,
self.z_near,
self.z_far,
c=c) # (NV, H, W, 8)
images_0to1 = images * 0.5 + 0.5 # (NV, 3, H, W)
curr_nviews = nviews[torch.randint(0, len(nviews), (1, )).item()]
views_src = np.sort(np.random.choice(NV, curr_nviews, replace=False))
view_dest = np.random.randint(0, NV - curr_nviews)
for vs in range(curr_nviews):
view_dest += view_dest >= views_src[vs]
views_src = torch.from_numpy(views_src)
# set renderer net to eval mode
renderer.eval()
source_views = (images_0to1[views_src].permute(
0, 2, 3, 1).cpu().numpy().reshape(-1, H, W, 3))
gt = images_0to1[view_dest].permute(1, 2,
0).cpu().numpy().reshape(H, W, 3)
with torch.no_grad():
test_rays = cam_rays[view_dest] # (H, W, 8)
test_images = images[views_src] # (NS, 3, H, W)
net.encode(
test_images.unsqueeze(0),
poses[views_src].unsqueeze(0),
focal.to(device=device),
c=c.to(device=device) if c is not None else None,
)
test_rays = test_rays.reshape(1, H * W, -1)
render_dict = DotMap(render_par(test_rays, want_weights=True))
coarse = render_dict.coarse
fine = render_dict.fine
using_fine = len(fine) > 0
alpha_coarse_np = coarse.weights[0].sum(
dim=-1).cpu().numpy().reshape(H, W)
rgb_coarse_np = coarse.rgb[0].cpu().numpy().reshape(H, W, 3)
depth_coarse_np = coarse.depth[0].cpu().numpy().reshape(H, W)
if using_fine:
alpha_fine_np = fine.weights[0].sum(
dim=1).cpu().numpy().reshape(H, W)
depth_fine_np = fine.depth[0].cpu().numpy().reshape(H, W)
rgb_fine_np = fine.rgb[0].cpu().numpy().reshape(H, W, 3)
print("c rgb min {} max {}".format(rgb_coarse_np.min(),
rgb_coarse_np.max()))
print("c alpha min {}, max {}".format(alpha_coarse_np.min(),
alpha_coarse_np.max()))
alpha_coarse_cmap = util.cmap(alpha_coarse_np) / 255
depth_coarse_cmap = util.cmap(depth_coarse_np) / 255
vis_list = [
*source_views,
gt,
depth_coarse_cmap,
rgb_coarse_np,
alpha_coarse_cmap,
]
vis_coarse = np.hstack(vis_list)
vis = vis_coarse
if using_fine:
print("f rgb min {} max {}".format(rgb_fine_np.min(),
rgb_fine_np.max()))
print("f alpha min {}, max {}".format(alpha_fine_np.min(),
alpha_fine_np.max()))
depth_fine_cmap = util.cmap(depth_fine_np) / 255
alpha_fine_cmap = util.cmap(alpha_fine_np) / 255
vis_list = [
*source_views,
gt,
depth_fine_cmap,
rgb_fine_np,
alpha_fine_cmap,
]
vis_fine = np.hstack(vis_list)
vis = np.vstack((vis_coarse, vis_fine))
rgb_psnr = rgb_fine_np
else:
rgb_psnr = rgb_coarse_np
psnr = util.psnr(rgb_psnr, gt)
vals = {"psnr": psnr}
print("psnr", psnr)
# set the renderer network back to train mode
renderer.train()
return vis, vals