(args, train_dl, val_dl, hwf, start, model, device, render_kwargs_test)
| 31 | device = torch.device('cuda:0') # this is really controlled in train.sh |
| 32 | |
| 33 | def render_test(args, train_dl, val_dl, hwf, start, model, device, render_kwargs_test): |
| 34 | model.eval() |
| 35 | |
| 36 | # ### Eval Training set result |
| 37 | if args.render_video_train: |
| 38 | images_train = [] |
| 39 | poses_train = [] |
| 40 | # views from train set |
| 41 | for img, pose in train_dl: |
| 42 | predict_pose = inference_pose_regression(args, img, device, model) |
| 43 | device_cpu = torch.device('cpu') |
| 44 | predict_pose = predict_pose.to(device_cpu) # put predict pose back to cpu |
| 45 | |
| 46 | img_val = img.permute(0,2,3,1) # (1,240,320,3) |
| 47 | pose_val = torch.zeros(1,4,4) |
| 48 | pose_val[0,:3,:4] = predict_pose.reshape(3,4)[:3,:4] # (1,3,4)) |
| 49 | pose_val[0,3,3] = 1. |
| 50 | images_train.append(img_val) |
| 51 | poses_train.append(pose_val) |
| 52 | |
| 53 | images_train = torch.cat(images_train, dim=0).numpy() |
| 54 | poses_train = torch.cat(poses_train, dim=0) |
| 55 | print('train poses shape', poses_train.shape) |
| 56 | torch.set_default_tensor_type('torch.cuda.FloatTensor') |
| 57 | with torch.no_grad(): |
| 58 | rgbs, disps = render_path(poses_train.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_train, savedir=None) |
| 59 | torch.set_default_tensor_type('torch.FloatTensor') |
| 60 | print('Saving trainset as video', rgbs.shape, disps.shape) |
| 61 | moviebase = os.path.join(args.basedir, args.model_name, '{}_trainset_{:06d}_'.format(args.model_name, start)) |
| 62 | imageio.mimwrite(moviebase + 'train_rgb.mp4', to8b(rgbs), fps=15, quality=8) |
| 63 | imageio.mimwrite(moviebase + 'train_disp.mp4', to8b(disps / np.max(disps)), fps=15, quality=8) |
| 64 | |
| 65 | ### Eval Validation set result |
| 66 | if args.render_video_test: |
| 67 | images_val = [] |
| 68 | poses_val = [] |
| 69 | # views from val set |
| 70 | for img, pose in val_dl: |
| 71 | predict_pose = inference_pose_regression(args, img, device, model) |
| 72 | device_cpu = torch.device('cpu') |
| 73 | predict_pose = predict_pose.to(device_cpu) # put predict pose back to cpu |
| 74 | |
| 75 | img_val = img.permute(0,2,3,1) # (1,240,360,3) |
| 76 | pose_val = torch.zeros(1,4,4) |
| 77 | pose_val[0,:3,:4] = predict_pose.reshape(3,4)[:3,:4] # (1,3,4)) |
| 78 | pose_val[0,3,3] = 1. |
| 79 | images_val.append(img_val) |
| 80 | poses_val.append(pose_val) |
| 81 | |
| 82 | images_val = torch.cat(images_val, dim=0).numpy() |
| 83 | poses_val = torch.cat(poses_val, dim=0) |
| 84 | print('test poses shape', poses_val.shape) |
| 85 | torch.set_default_tensor_type('torch.cuda.FloatTensor') |
| 86 | with torch.no_grad(): |
| 87 | rgbs, disps = render_path(poses_val.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_val, savedir=None) |
| 88 | torch.set_default_tensor_type('torch.FloatTensor') |
| 89 | print('Saving testset as video', rgbs.shape, disps.shape) |
| 90 | moviebase = os.path.join(args.basedir, args.model_name, '{}_test_{:06d}_'.format(args.model_name, start)) |
no test coverage detected