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