File size: 2,081 Bytes
4e7b4da |
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 |
import os
import random
import json
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.transforms.functional as F
import numpy as np
from einops import rearrange, repeat
class Camera(object):
def __init__(self, entry):
fx, fy, cx, cy = entry[1:5]
self.fx = fx
self.fy = fy
self.cx = cx
self.cy = cy
w2c_mat = np.array(entry[7:]).reshape(3, 4)
w2c_mat_4x4 = np.eye(4)
w2c_mat_4x4[:3, :] = w2c_mat
self.w2c_mat = w2c_mat_4x4
self.c2w_mat = np.linalg.inv(w2c_mat_4x4)
def load_cameras(path):
with open(path, 'r') as f:
poses = f.readlines()
poses = [pose.strip().split(' ') for pose in poses[1:]]
cam_params = [[float(x) for x in pose] for pose in poses]
cam_params = [Camera(cam_param) for cam_param in cam_params]
return cam_params
def get_relative_pose(cam_params):
abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
source_cam_c2w = abs_c2ws[0]
cam_to_origin = 0
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, -cam_to_origin],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
abs2rel = target_cam_c2w @ abs_w2cs[0]
ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
ret_poses = np.array(ret_poses, dtype=np.float32)
return ret_poses
def get_K(intrinsics, size, do_normalize=False):
def normalize_intrinsic(x, size):
h, w = size
x[:,:,0:1] = x[:,:,0:1] / w
x[:,:,1:2] = x[:,:,1:2] / h
return x
b, _, t, _ = intrinsics.shape
K = torch.zeros((b, t, 9), dtype=intrinsics.dtype, device=intrinsics.device)
fx, fy, cx, cy = intrinsics.squeeze(1).chunk(4, dim=-1)
K[:,:,0:1] = fx
K[:,:,2:3] = cx
K[:,:,4:5] = fy
K[:,:,5:6] = cy
K[:,:,8:9] = 1.0
K = rearrange(K, "b t (h w) -> b t h w", h=3, w=3)
if do_normalize:
K = normalize_intrinsic(K, size)
return K |