Spaces:
Paused
Paused
File size: 3,720 Bytes
b177539 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import os
import torch
import numpy as np
import imageio
import json
import torch.nn.functional as F
import cv2
def normalize(x):
return x / np.linalg.norm(x)
trans_t = lambda t : np.array([
[1,0,0,0],
[0,1,0,0],
[0,0,1,t],
[0,0,0,1]]).astype(np.float32)
trans_center = lambda centroid : np.array([
[1,0,0,centroid[0]],
[0,1,0,centroid[1]],
[0,0,1,centroid[2]],
[0,0,0,1]]).astype(np.float32)
rot_phi = lambda phi : np.array([ # rot dir: +y -> +z
[1,0,0,0],
[0,np.cos(phi),-np.sin(phi),0],
[0,np.sin(phi), np.cos(phi),0],
[0,0,0,1]]).astype(np.float32)
rot_theta = lambda th : np.array([ # rot dir: +x -> +z
[np.cos(th),0,-np.sin(th),0],
[0,1,0,0],
[np.sin(th),0, np.cos(th),0],
[0,0,0,1]]).astype(np.float32)
rot_gamma = lambda ga : np.array([ # rot dir: +x -> +y
[np.cos(ga),-np.sin(ga),0,0],
[np.sin(ga), np.cos(ga),0,0],
[0,0,1,0],
[0,0,0,1]]).astype(np.float32)
def pose_spherical(gamma, phi, t):
c2w = np.array([
[1,0,0,0],
[0,1,0,0],
[0,0,1,0],
[0,0,0,1]]).astype(np.float32)
c2w = rot_phi(phi/180.*np.pi) @ c2w
c2w = rot_gamma(gamma/180.*np.pi) @ c2w
c2w[:3, 3] = t
return c2w
def load_lerf_data(basedir, factor=2, args=None, movie_render_kwargs={}):
with open(os.path.join(basedir, 'transforms.json'), 'r') as fp:
metas = json.load(fp)
imgs = []
poses = []
intrinsics = []
fts = []
skip = 1
for frame in metas['frames'][::skip]:
fname = os.path.join(basedir, frame['file_path'])
just_fname = fname.split('/')[-1]
if factor >= 2:
fname = os.path.join(basedir, 'images_{}'.format(factor), just_fname)
else:
fname = os.path.join(basedir, 'images', just_fname)
imgs.append(imageio.imread(fname))
poses.append(np.array(frame['transform_matrix']))
K = np.array([
[frame['fl_x']/factor, 0, frame['cx']/factor],
[0, frame['fl_y']/factor, frame['cy']/factor],
[0, 0, 1]
]).astype(np.float32)
intrinsics.append(K)
imgs = (np.array(imgs) / 255.).astype(np.float32) # keep all 4 channels (RGBA)
poses = np.array(poses).astype(np.float32)
intrinsics = np.array(intrinsics).astype(np.float32)
f_avg = (intrinsics[:, 0, 0] + intrinsics[:, 1, 1]).mean() / 2.
i_test = np.arange(0, int(poses.shape[0]), 8)
i_val = i_test
i_train = np.array([i for i in np.arange(int(poses.shape[0])) if
(i not in i_test and i not in i_val)])
i_split = [i_train, i_val, i_test]
H, W = imgs[0].shape[:2]
poses_ = poses.copy()
centroid = poses_[:,:3,3].mean(0)
radcircle = movie_render_kwargs.get('scale_r', 0) * 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)
up_rad = movie_render_kwargs.get('pitch_deg', 0)
# render_poses = torch.stack([pose_spherical(angle, up_rad, centroid) for angle in np.linspace(-180,180,80+1)[:-1]], 0)
render_poses = []
camera_o = np.zeros_like(centroid)
num_render = 90
for th in np.linspace(0., 360., num_render):
camera_o[0] = centroid[0] + radcircle * np.cos(th/180.*np.pi)
camera_o[1] = centroid[1] + radcircle * np.sin(th/180.*np.pi)
camera_o[2] = centroid[2]
render_poses.append(pose_spherical(th+90.0, up_rad, camera_o))
render_poses = np.stack(render_poses, axis=0)
return imgs, poses, render_poses, [H, W, f_avg], intrinsics, i_split
|