Perform inference on a random val image and save the result
(args, epoch, val_dl, model, hwf, half_res, device, num_samples=1, **render_kwargs_test)
| 298 | return total_loss_mean, total_psnr_mean |
| 299 | |
| 300 | def save_val_result_7Scenes(args, epoch, val_dl, model, hwf, half_res, device, num_samples=1, **render_kwargs_test): |
| 301 | ''' Perform inference on a random val image and save the result ''' |
| 302 | half_res=False # save half res image to reduce computational speed |
| 303 | model.eval() |
| 304 | i = 0 |
| 305 | for batch in val_dl: |
| 306 | if args.NeRFH: |
| 307 | data, pose, img_idx = batch |
| 308 | else: |
| 309 | data, pose = batch |
| 310 | if i >= num_samples: |
| 311 | return |
| 312 | H, W, focal = hwf |
| 313 | target_ = deepcopy(data) |
| 314 | img_idx = img_idx.to(device) |
| 315 | |
| 316 | pose_ = inference_pose_regression(args, data, device, model) |
| 317 | device_cpu = torch.device('cpu') |
| 318 | pose_ = pose_.to(device_cpu) # put predict_pose back to cpu |
| 319 | |
| 320 | pose_nerf = pose_.clone() |
| 321 | if args.NeRFH: |
| 322 | # rescale the predicted pose to nerf scales |
| 323 | pose_nerf = fix_coord_supp(args, pose_nerf) |
| 324 | |
| 325 | # every new tensor from onward is in GPU |
| 326 | torch.set_default_tensor_type('torch.cuda.FloatTensor') |
| 327 | with torch.no_grad(): |
| 328 | if half_res: |
| 329 | rgb, disp, acc, extras = render(H//2, W//2, focal/2, chunk=args.chunk, c2w=pose_nerf[0].to(device), img_idx=img_idx, **render_kwargs_test) |
| 330 | else: |
| 331 | rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0].to(device), img_idx=img_idx, **render_kwargs_test) |
| 332 | |
| 333 | ### Set Save Dir ### |
| 334 | if num_samples <=1: |
| 335 | out_folder = os.path.join(args.basedir, args.model_name, 'val_imgs') |
| 336 | else: |
| 337 | out_folder = os.path.join(args.basedir, args.model_name, 'val_imgs_batches') |
| 338 | if not os.path.isdir(out_folder): |
| 339 | os.mkdir(out_folder) |
| 340 | |
| 341 | if half_res: |
| 342 | target_img = F.interpolate(target_, scale_factor=0.5, mode='area').permute(0,2,3,1).reshape(H//2,W//2,3) |
| 343 | rgb_img = rgb.reshape(H//2,W//2,3) |
| 344 | else: |
| 345 | target_img = target_.permute(0,2,3,1).reshape(H,W,3) |
| 346 | rgb_img = rgb.reshape(H,W,3) |
| 347 | target_img_to_save = to8b(target_img.to(device_cpu).detach().numpy()) |
| 348 | rgb_img_to_save = to8b(rgb_img.to(device_cpu).detach().numpy()) |
| 349 | # save NeRF Rendered RGB Img |
| 350 | import imageio |
| 351 | if num_samples <= 1: |
| 352 | imageio.imwrite(os.path.join(out_folder, '{0:04d}_gt.png'.format(epoch)), target_img_to_save) |
| 353 | imageio.imwrite(os.path.join(out_folder, '{0:04d}.png'.format(epoch)), rgb_img_to_save) |
| 354 | else: |
| 355 | imageio.imwrite(os.path.join(out_folder, '{0:04d}_gt.png'.format(i)), target_img_to_save) |
| 356 | imageio.imwrite(os.path.join(out_folder, '{0:04d}.png'.format(i)), rgb_img_to_save) |
| 357 |
no test coverage detected