Spaces:
Paused
Paused
| import numpy as np | |
| from .load_llff import load_llff_data | |
| from .load_blender import load_blender_data | |
| from .load_nsvf import load_nsvf_data | |
| from .load_blendedmvs import load_blendedmvs_data | |
| from .load_tankstemple import load_tankstemple_data | |
| from .load_deepvoxels import load_dv_data | |
| from .load_co3d import load_co3d_data | |
| from .load_nerfpp import load_nerfpp_data | |
| from .load_replica import load_replica_data | |
| from .load_lerf import load_lerf_data | |
| def load_data(args): | |
| K, depths = None, None | |
| near_clip = None | |
| if args.dataset_type == 'llff': | |
| images, depths, poses, bds, render_poses, i_test = load_llff_data( | |
| args.datadir, args.factor, args.width, args.height, | |
| recenter=True, bd_factor=args.bd_factor, | |
| spherify=args.spherify, | |
| load_depths=args.load_depths, | |
| movie_render_kwargs=args.movie_render_kwargs, args=args) | |
| hwf = poses[0,:3,-1] | |
| poses = poses[:,:3,:4] | |
| print('Loaded llff', images.shape, render_poses.shape, hwf, args.datadir) | |
| if not isinstance(i_test, list): | |
| i_test = [i_test] | |
| if args.llffhold > 0: | |
| print('Auto LLFF holdout,', args.llffhold) | |
| i_test = np.arange(images.shape[0])[::args.llffhold] | |
| # i_test = [1, 2] | |
| # i_test = [] | |
| i_val = i_test | |
| i_train = np.array([i for i in np.arange(int(images.shape[0])) if | |
| (i not in i_test and i not in i_val)]) | |
| print('DEFINING BOUNDS') | |
| if args.ndc: | |
| near = 0. | |
| far = 1. | |
| else: | |
| near_clip = max(np.ndarray.min(bds) * .9, 0) | |
| _far = max(np.ndarray.max(bds) * 1., 0) | |
| near = 0 | |
| far = inward_nearfar_heuristic(poses[i_train, :3, 3])[1] | |
| print('near_clip', near_clip) | |
| print('original far', _far) | |
| print('NEAR FAR', near, far) | |
| if depths == 0: | |
| depths = np.zeros_like(images[..., :1]) | |
| elif args.dataset_type == 'blender': | |
| images, poses, render_poses, hwf, i_split = load_blender_data(args.datadir, args.half_res, args.testskip, args=args) | |
| print('Loaded blender', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near, far = 2., 6. | |
| if images.shape[-1] == 4: | |
| if args.white_bkgd: | |
| images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:]) | |
| else: | |
| images = images[...,:3]*images[...,-1:] | |
| elif args.dataset_type == 'blendedmvs': | |
| images, poses, render_poses, hwf, K, i_split = load_blendedmvs_data(args.datadir) | |
| print('Loaded blendedmvs', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3]) | |
| assert images.shape[-1] == 3 | |
| elif args.dataset_type == 'tankstemple': | |
| images, poses, render_poses, hwf, K, i_split = load_tankstemple_data( | |
| args.datadir, movie_render_kwargs=args.movie_render_kwargs) | |
| print('Loaded tankstemple', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| # i_test = [0] | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3], ratio=0) | |
| near_clip = near | |
| if images.shape[-1] == 4: | |
| if args.white_bkgd: | |
| images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:]) | |
| else: | |
| images = images[...,:3]*images[...,-1:] | |
| elif args.dataset_type == 'nsvf': | |
| images, poses, render_poses, hwf, i_split = load_nsvf_data(args.datadir) | |
| print('Loaded nsvf', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3]) | |
| near_clip = near | |
| if images.shape[-1] == 4: | |
| if args.white_bkgd: | |
| images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:]) | |
| else: | |
| images = images[...,:3]*images[...,-1:] | |
| elif args.dataset_type == 'deepvoxels': | |
| images, poses, render_poses, hwf, i_split = load_dv_data(scene=args.scene, basedir=args.datadir, testskip=args.testskip) | |
| print('Loaded deepvoxels', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| hemi_R = np.mean(np.linalg.norm(poses[:,:3,-1], axis=-1)) | |
| near = hemi_R - 1 | |
| far = hemi_R + 1 | |
| assert args.white_bkgd | |
| assert images.shape[-1] == 3 | |
| elif args.dataset_type == 'co3d': | |
| # each image can be in different shapes and intrinsics | |
| images, masks, poses, render_poses, hwf, K, i_split = load_co3d_data(args) | |
| print('Loaded co3d', args.datadir, args.annot_path, args.sequence_name) | |
| i_train, i_val, i_test = i_split | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3], ratio=0) | |
| for i in range(len(images)): | |
| if args.white_bkgd: | |
| images[i] = images[i] * masks[i][...,None] + (1.-masks[i][...,None]) | |
| else: | |
| images[i] = images[i] * masks[i][...,None] | |
| elif args.dataset_type == 'nerfpp': | |
| images, poses, render_poses, hwf, K, i_split = load_nerfpp_data(args.datadir) | |
| print('Loaded nerf_pp', images.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near_clip, far = inward_nearfar_heuristic(poses[i_train, :3, 3], ratio=0.02) | |
| near = 0 | |
| elif args.dataset_type == 'replica': | |
| images, poses, render_poses, hwf, i_split = load_replica_data(args.datadir, args.half_res, args.testskip, args=args, \ | |
| spherify=args.spherify,movie_render_kwargs=args.movie_render_kwargs) | |
| print('Loaded replica', images.shape, render_poses.shape, hwf, args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3], ratio=0) | |
| near_clip = near | |
| print('NEAR FAR', near, far) | |
| if images.shape[-1] == 4: | |
| if args.white_bkgd: | |
| images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:]) | |
| else: | |
| images = images[...,:3]*images[...,-1:] | |
| elif args.dataset_type == 'lerf': | |
| images, poses, render_poses, hwf, K, i_split = load_lerf_data(args.datadir, args.factor,movie_render_kwargs=args.movie_render_kwargs) | |
| print('Loaded lerf', images.shape, render_poses.shape, hwf[:2], args.datadir) | |
| i_train, i_val, i_test = i_split | |
| near, far = inward_nearfar_heuristic(poses[i_train, :3, 3], ratio=0) | |
| near_clip = near | |
| print('NEAR FAR', near, far) | |
| if images.shape[-1] == 4: | |
| if args.white_bkgd: | |
| images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:]) | |
| else: | |
| images = images[...,:3]*images[...,-1:] | |
| else: | |
| raise NotImplementedError(f'Unknown dataset type {args.dataset_type} exiting') | |
| # Cast intrinsics to right types | |
| H, W, focal = hwf | |
| H, W = int(H), int(W) | |
| hwf = [H, W, focal] | |
| HW = np.array([im.shape[:2] for im in images]) | |
| irregular_shape = (images.dtype is np.dtype('object')) | |
| if K is None: | |
| K = np.array([ | |
| [focal, 0, 0.5*W], | |
| [0, focal, 0.5*H], | |
| [0, 0, 1] | |
| ]) | |
| if len(K.shape) == 2: | |
| Ks = K[None].repeat(len(poses), axis=0) | |
| else: | |
| Ks = K | |
| render_poses = render_poses[...,:4] | |
| data_dict = dict( | |
| hwf=hwf, HW=HW, Ks=Ks, | |
| near=near, far=far, near_clip=near_clip, | |
| i_train=i_train, i_val=i_val, i_test=i_test, | |
| poses=poses, render_poses=render_poses, | |
| images=images, depths=depths, | |
| irregular_shape=irregular_shape | |
| ) | |
| return data_dict | |
| def inward_nearfar_heuristic(cam_o, ratio=0.05): | |
| dist = np.linalg.norm(cam_o[:,None] - cam_o, axis=-1) | |
| far = dist.max() # could be too small to exist the scene bbox | |
| # it is only used to determined scene bbox | |
| # lib/dvgo use 1e9 as far | |
| near = far * ratio | |
| return near, far | |