Spaces:
Paused
Paused
| import os | |
| import glob | |
| 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) | |
| def load_tankstemple_data(basedir, movie_render_kwargs={}, args=None): | |
| pose_paths = sorted(glob.glob(os.path.join(basedir, 'pose', '*txt'))) | |
| rgb_paths = sorted(glob.glob(os.path.join(basedir, 'rgb', '*jpg'))) | |
| all_poses = [] | |
| all_imgs = [] | |
| i_split = [[], []] | |
| for i, (pose_path, rgb_path) in enumerate(zip(pose_paths, rgb_paths)): | |
| i_set = int(os.path.split(rgb_path)[-1][0]) | |
| all_poses.append(np.loadtxt(pose_path).astype(np.float32)) | |
| all_imgs.append((imageio.imread(rgb_path) / 255.).astype(np.float32)) | |
| i_split[i_set].append(i) | |
| imgs = np.stack(all_imgs, 0) | |
| poses = np.stack(all_poses, 0) | |
| i_split.append(i_split[-1]) | |
| height, width = imgs.shape[1:3] | |
| if args is not None and args.distill_active: | |
| fts_dict = load_features(file=os.path.join(basedir, "features.pt"), imhw=(height, width)) | |
| fts = torch.stack(list(fts_dict.values())).permute(2,3,1,0) | |
| else: | |
| fts = torch.zeros([height, width, 0, imgs.shape[-1]]) | |
| path_intrinsics = os.path.join(basedir, 'intrinsics.txt') | |
| H, W = imgs[0].shape[:2] | |
| K = np.loadtxt(path_intrinsics) | |
| focal = float(K[0,0]) | |
| ### generate spiral poses for rendering fly-through movie | |
| centroid = poses[:,:3,3].mean(0) | |
| radcircle = movie_render_kwargs.get('scale_r', 1.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) | |
| 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_vec', 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) | |
| lookat *= -1 | |
| 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) | |
| return imgs, poses, render_poses, [H, W, focal], K, i_split, fts | |