import numpy as np import os, imageio import torch import scipy from tqdm import tqdm ########## Slightly modified version of LLFF data loading code ########## see https://github.com/Fyusion/LLFF for original def imread(f): if f.endswith('png'): return imageio.imread(f, ignoregamma=True) else: return imageio.imread(f) def depthread(path): with open(path, "rb") as fid: width, height, channels = np.genfromtxt(fid, delimiter="&", max_rows=1, usecols=(0, 1, 2), dtype=int) fid.seek(0) num_delimiter = 0 byte = fid.read(1) while True: if byte == b"&": num_delimiter += 1 if num_delimiter >= 3: break byte = fid.read(1) array = np.fromfile(fid, np.float32) array = array.reshape((width, height, channels), order="F") return np.transpose(array, (1, 0, 2)).squeeze() def _minify(basedir, factors=[], resolutions=[]): needtoload = False for r in factors: imgdir = os.path.join(basedir, 'images_{}'.format(r)) if not os.path.exists(imgdir): needtoload = True for r in resolutions: imgdir = os.path.join(basedir, 'images_{}x{}'.format(r[1], r[0])) if not os.path.exists(imgdir): needtoload = True if not needtoload: return from shutil import copy from subprocess import check_output imgdir = os.path.join(basedir, 'images') imgs = [os.path.join(imgdir, f) for f in sorted(os.listdir(imgdir))] imgs = [f for f in imgs if any([f.endswith(ex) for ex in ['JPG', 'jpg', 'png', 'jpeg', 'PNG']])] imgdir_orig = imgdir wd = os.getcwd() for r in factors + resolutions: if isinstance(r, int): name = 'images_{}'.format(r) resizearg = '{}%'.format(100./r) else: name = 'images_{}x{}'.format(r[1], r[0]) resizearg = '{}x{}'.format(r[1], r[0]) imgdir = os.path.join(basedir, name) if os.path.exists(imgdir): continue print('Minifying', r, basedir) os.makedirs(imgdir) check_output('cp {}/* {}'.format(imgdir_orig, imgdir), shell=True) ext = imgs[0].split('.')[-1] args = ' '.join(['mogrify', '-resize', resizearg, '-format', 'png', '*.{}'.format(ext)]) print(args) os.chdir(imgdir) check_output(args, shell=True) os.chdir(wd) if ext != 'png': check_output('rm {}/*.{}'.format(imgdir, ext), shell=True) print('Removed duplicates') print('Done') def _load_data(basedir, factor=None, width=None, height=None, load_imgs=True, load_depths=False, args=None): poses_arr = np.load(os.path.join(basedir, 'poses_bounds.npy')) if poses_arr.shape[1] == 17: poses = poses_arr[:, :-2].reshape([-1, 3, 5]).transpose([1,2,0]) elif poses_arr.shape[1] == 14: poses = poses_arr[:, :-2].reshape([-1, 3, 4]).transpose([1,2,0]) else: raise NotImplementedError bds = poses_arr[:, -2:].transpose([1,0]) img0 = [os.path.join(basedir, 'images', f) for f in sorted(os.listdir(os.path.join(basedir, 'images'))) \ if f.endswith('JPG') or f.endswith('jpg') or f.endswith('png')][0] sh = imageio.imread(img0).shape sfx = '' if height is not None and width is not None: _minify(basedir, resolutions=[[height, width]]) sfx = '_{}x{}'.format(width, height) elif factor is not None and factor != 1: sfx = '_{}'.format(factor) _minify(basedir, factors=[factor]) factor = factor elif height is not None: factor = sh[0] / float(height) width = int(sh[1] / factor) _minify(basedir, resolutions=[[height, width]]) sfx = '_{}x{}'.format(width, height) elif width is not None: factor = sh[1] / float(width) height = int(sh[0] / factor) _minify(basedir, resolutions=[[height, width]]) sfx = '_{}x{}'.format(width, height) else: factor = 1 imgdir = os.path.join(basedir, 'images' + sfx) print(f'Loading images from {imgdir}') if not os.path.exists(imgdir): print( imgdir, 'does not exist, returning' ) return imgfiles = [os.path.join(imgdir, f) for f in sorted(os.listdir(imgdir)) if f.endswith('JPG') or f.endswith('jpg') or f.endswith('png')] if poses.shape[-1] != len(imgfiles): print() print( 'Mismatch between imgs {} and poses {} !!!!'.format(len(imgfiles), poses.shape[-1]) ) names = set(name[:-4] for name in np.load(os.path.join(basedir, 'poses_names.npy'))) assert len(names) == poses.shape[-1] print('Below failed files are skip due to SfM failure:') new_imgfiles = [] for i in imgfiles: fname = os.path.split(i)[1][:-4] if fname in names: new_imgfiles.append(i) else: print('==>', i) imgfiles = new_imgfiles if len(imgfiles) < 3: print('Too few images...') import sys; sys.exit() sh = imageio.imread(imgfiles[0]).shape if poses.shape[1] == 4: poses = np.concatenate([poses, np.zeros_like(poses[:,[0]])], 1) poses[2, 4, :] = np.load(os.path.join(basedir, 'hwf_cxcy.npy'))[2] poses[:2, 4, :] = np.array(sh[:2]).reshape([2, 1]) poses[2, 4, :] = poses[2, 4, :] * 1./factor if not load_imgs: return poses, bds imgs = imgs = [imread(f)[...,:3]/255. for f in tqdm(imgfiles)] imgs = np.stack(imgs, -1) height, width = imgs.shape[:2] print('Loaded image data', imgs.shape, poses[:,-1,0]) if not load_depths: return poses, bds, imgs depthdir = os.path.join(basedir, 'stereo', 'depth_maps') assert os.path.exists(depthdir), f'Dir not found: {depthdir}' depthfiles = [os.path.join(depthdir, f) for f in sorted(os.listdir(depthdir)) if f.endswith('.geometric.bin')] assert poses.shape[-1] == len(depthfiles), 'Mismatch between imgs {} and poses {} !!!!'.format(len(depthfiles), poses.shape[-1]) depths = [depthread(f) for f in depthfiles] depths = np.stack(depths, -1) print('Loaded depth data', depths.shape) return poses, bds, imgs, depths def normalize(x): return x / np.linalg.norm(x) def viewmatrix(z, up, pos): vec2 = normalize(z) vec1_avg = up vec0 = normalize(np.cross(vec1_avg, vec2)) vec1 = normalize(np.cross(vec2, vec0)) m = np.stack([vec0, vec1, vec2, pos], 1) return m def ptstocam(pts, c2w): tt = np.matmul(c2w[:3,:3].T, (pts-c2w[:3,3])[...,np.newaxis])[...,0] return tt def poses_avg(poses): hwf = poses[0, :3, -1:] center = poses[:, :3, 3].mean(0) vec2 = normalize(poses[:, :3, 2].sum(0)) up = poses[:, :3, 1].sum(0) c2w = np.concatenate([viewmatrix(vec2, up, center), hwf], 1) return c2w def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, rots, N): render_poses = [] rads = np.array(list(rads) + [1.]) hwf = c2w[:,4:5] for theta in np.linspace(0., 2. * np.pi * rots, N+1)[:-1]: c = np.dot(c2w[:3,:4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta*zrate)*zdelta, 1.]) * rads) z = normalize(c - np.dot(c2w[:3,:4], np.array([0,0,-focal, 1.]))) render_poses.append(np.concatenate([viewmatrix(z, up, c), hwf], 1)) return render_poses def recenter_poses(poses): poses_ = poses+0 bottom = np.reshape([0,0,0,1.], [1,4]) c2w = poses_avg(poses) c2w = np.concatenate([c2w[:3,:4], bottom], -2) bottom = np.tile(np.reshape(bottom, [1,1,4]), [poses.shape[0],1,1]) poses = np.concatenate([poses[:,:3,:4], bottom], -2) poses = np.linalg.inv(c2w) @ poses poses_[:,:3,:4] = poses[:,:3,:4] poses = poses_ return poses def rerotate_poses(poses): # [B, 3, 5] poses = np.copy(poses) centroid = poses[:,:3,3].mean(0) poses[:,:3,3] = poses[:,:3,3] - centroid # Find the minimum pca vector with minimum eigen value x = poses[:,:,3] mu = x.mean(0) cov = np.cov((x-mu).T) ev , eig = np.linalg.eig(cov) cams_up = eig[:,np.argmin(ev)] # [3, ] if cams_up[1] < 0: cams_up = -cams_up # Find rotation matrix that align cams_up with [0,1,0] R = scipy.spatial.transform.Rotation.align_vectors( [[0,1,0]], cams_up[None])[0].as_matrix() # Apply rotation and add back the centroid position poses[:,:3,:3] = R @ poses[:,:3,:3] poses[:,:3,[3]] = R @ poses[:,:3,[3]] poses[:,:3,3] = poses[:,:3,3] + centroid return poses ##################### def spherify_poses(poses, bds, depths): p34_to_44 = lambda p : np.concatenate([p, np.tile(np.reshape(np.eye(4)[-1,:], [1,1,4]), [p.shape[0], 1,1])], 1) rays_d = poses[:,:3,2:3] rays_o = poses[:,:3,3:4] def min_line_dist(rays_o, rays_d): A_i = np.eye(3) - rays_d * np.transpose(rays_d, [0,2,1]) b_i = -A_i @ rays_o pt_mindist = np.squeeze(-np.linalg.inv((np.transpose(A_i, [0,2,1]) @ A_i).mean(0)) @ (b_i).mean(0)) return pt_mindist pt_mindist = min_line_dist(rays_o, rays_d) center = pt_mindist up = (poses[:,:3,3] - center).mean(0) vec0 = normalize(up) vec1 = normalize(np.cross([.1,.2,.3], vec0)) vec2 = normalize(np.cross(vec0, vec1)) pos = center c2w = np.stack([vec1, vec2, vec0, pos], 1) poses_reset = np.linalg.inv(p34_to_44(c2w[None])) @ p34_to_44(poses[:,:3,:4]) radius = np.sqrt(np.mean(np.sum(np.square(poses_reset[:,:3,3]), -1))) sc = 1./radius poses_reset[:,:3,3] *= sc bds *= sc radius *= sc depths *= sc poses_reset = np.concatenate([poses_reset[:,:3,:4], np.broadcast_to(poses[0,:3,-1:], poses_reset[:,:3,-1:].shape)], -1) return poses_reset, radius, bds, depths def load_llff_data(basedir, factor=8, width=None, height=None, recenter=True, rerotate=True, bd_factor=.75, spherify=False, path_zflat=False, load_depths=False, movie_render_kwargs={}, args=None): poses, bds, imgs = _load_data(basedir, factor=factor, width=width, height=height, load_depths=load_depths, args=args) # factor=8 downsamples original imgs by 8x print('Loaded', basedir, bds.min(), bds.max()) if load_depths: depths = depths[0] else: depths = 0 # Correct rotation matrix ordering and move variable dim to axis 0 poses = np.concatenate([poses[:, 1:2, :], -poses[:, 0:1, :], poses[:, 2:, :]], 1) poses = np.moveaxis(poses, -1, 0).astype(np.float32) imgs = np.moveaxis(imgs, -1, 0).astype(np.float32) images = imgs bds = np.moveaxis(bds, -1, 0).astype(np.float32) # Rescale if bd_factor is provided if bds.min() < 0 and bd_factor is not None: print('Found negative z values from SfM sparse points!?') print('Please try bd_factor=None') import sys; sys.exit() sc = 1. if bd_factor is None else 1./(bds.min() * bd_factor) poses[:,:3,3] *= sc # [B, 3, 5] bds *= sc depths *= sc if recenter: poses = recenter_poses(poses) # [B, 3, 5] if spherify: poses, radius, bds, depths = spherify_poses(poses, bds, depths) # [B, 3, 5] if rerotate: poses = rerotate_poses(poses) # [B, 3, 5] ### generate spiral poses for rendering fly-through movie centroid = poses[:,:3,3].mean(0) radcircle = movie_render_kwargs.get('scale_r', 1) * np.linalg.norm(poses[:,:3,3] - centroid, axis=-1).mean() centroid[0] += movie_render_kwargs.get('shift_x', 0) centroid[1] += movie_render_kwargs.get('shift_y', 0) centroid[2] += movie_render_kwargs.get('shift_z', 0) new_up_rad = movie_render_kwargs.get('pitch_deg', 0) * np.pi / 180 target_y = radcircle * np.tan(new_up_rad) render_poses = [] for th in np.linspace(0., 2.*np.pi, 200): camorigin = np.array([radcircle * np.cos(th), 0, radcircle * np.sin(th)]) if movie_render_kwargs.get('flip_up', False): up = np.array([0,1.,0]) else: up = np.array([0,-1.,0]) vec2 = normalize(camorigin) vec0 = normalize(np.cross(vec2, up)) vec1 = normalize(np.cross(vec2, vec0)) pos = camorigin + centroid # rotate to align with new pitch rotation lookat = -vec2 lookat[1] = target_y lookat = normalize(lookat) vec2 = -lookat vec1 = normalize(np.cross(vec2, vec0)) p = np.stack([vec0, vec1, vec2, pos], 1) render_poses.append(p) render_poses = np.stack(render_poses, 0) render_poses = np.concatenate([render_poses, np.broadcast_to(poses[0,:3,-1:], render_poses[:,:3,-1:].shape)], -1) else: c2w = poses_avg(poses) print('recentered', c2w.shape) print(c2w[:3,:4]) ## Get spiral # Get average pose up = normalize(poses[:, :3, 1].sum(0)) # Find a reasonable "focus depth" for this dataset close_depth, inf_depth = bds.min()*.9, bds.max()*5. dt = .75 mean_dz = 1./(((1.-dt)/close_depth + dt/inf_depth)) focal = mean_dz * movie_render_kwargs.get('scale_f', 1) # Get radii for spiral path zdelta = movie_render_kwargs.get('zdelta', 0.5) zrate = movie_render_kwargs.get('zrate', 1.0) tt = poses[:,:3,3] # ptstocam(poses[:3,3,:].T, c2w).T rads = np.percentile(np.abs(tt), 90, 0) * movie_render_kwargs.get('scale_r', 1) c2w_path = c2w N_views = 120 N_rots = movie_render_kwargs.get('N_rots', 1) if path_zflat: # zloc = np.percentile(tt, 10, 0)[2] zloc = -close_depth * .1 c2w_path[:3,3] = c2w_path[:3,3] + zloc * c2w_path[:3,2] rads[2] = 0. N_rots = 1 N_views/=2 # Generate poses for spiral path render_poses = render_path_spiral(c2w_path, up, rads, focal, zdelta, zrate=zrate, rots=N_rots, N=N_views) render_poses = torch.Tensor(render_poses) c2w = poses_avg(poses) print('Data:') print(poses.shape, images.shape, bds.shape) dists = np.sum(np.square(c2w[:3,3] - poses[:,:3,3]), -1) i_test = np.argmin(dists) print('HOLDOUT view is', i_test) images = images.astype(np.float32) poses = poses.astype(np.float32) return images, depths, poses, bds, render_poses, i_test