Function render_nerfw_imgs

script/feature/misc.py:203–252

render nerfw imgs, save unscaled pose and results

(args, dl, hwf, device, render_kwargs_test, world_setup_dict)

    return

def render_nerfw_imgs(args, dl, hwf, device, render_kwargs_test, world_setup_dict):
    ''' render nerfw imgs, save unscaled pose and results'''
    H, W, focal = hwf
    target_list = []
    rgb_list = []
    pose_list = []
    img_idx_list = []
    # profiling code
    # time0 = time.time()
    
    # inference nerfw and save rgb, target, pose
    for batch_idx, (target, pose, img_idx) in enumerate(dl):
        if batch_idx % 10 == 0:
            print("renders {}/total {}".format(batch_idx, len(dl.dataset)))

        target = target[0].permute(1,2,0).to(device) # (240,360,3)
        pose = pose.reshape(3,4) # reshape to 3x4 rot matrix

        img_idx = img_idx.to(device)
        pose_nerf = pose.clone()

        # rescale the predicted pose to nerf scales
        pose_nerf = fix_coord_supp(args, pose_nerf[None,...], world_setup_dict)

        # generate nerf image
        with torch.no_grad():
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
            if args.tinyimg:
                rgb, _, _, _ = render(int(H//args.tinyscale), int(W//args.tinyscale), focal/args.tinyscale, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=True, img_idx=img_idx, **render_kwargs_test)
                # convert rgb to B,C,H,W
                rgb = rgb[None,...].permute(0,3,1,2)
                # upsample rgb to hwf size
                rgb = torch.nn.Upsample(size=(H, W), mode='bicubic')(rgb)
                # convert rgb back to H,W,C format
                rgb = rgb[0].permute(1,2,0)

            else:
                rgb, _, _, _ = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=True, img_idx=img_idx, **render_kwargs_test)
            torch.set_default_tensor_type('torch.FloatTensor')

        target_list.append(target.cpu())
        rgb_list.append(rgb.cpu())
        pose_list.append(pose.cpu())
        img_idx_list.append(img_idx.cpu())

    targets = torch.stack(target_list).detach()
    rgbs = torch.stack(rgb_list).detach()
    poses = torch.stack(pose_list).detach()
    img_idxs = torch.stack(img_idx_list).detach()
    return targets, rgbs, poses, img_idxs

def render_virtual_imgs(args, pose_perturb, img_idxs, hwf, device, render_kwargs_test, world_setup_dict):
    ''' render nerfw imgs, save unscaled pose and results'''