mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 19:33:46 +08:00
537 lines
17 KiB
Python
537 lines
17 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# pylint: skip-file
|
|
"""
|
|
Camera transformation helper code.
|
|
"""
|
|
|
|
import math
|
|
from typing import Literal, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torchtyping import TensorType
|
|
|
|
_EPS = np.finfo(float).eps * 4.0
|
|
|
|
|
|
def unit_vector(data, axis=None, out=None):
|
|
"""Return ndarray normalized by length, i.e. Euclidean norm, along axis.
|
|
>>> v0 = np.random.random(3)
|
|
>>> v1 = unit_vector(v0)
|
|
>>> np.allclose(v1, v0 / np.linalg.norm(v0))
|
|
True
|
|
>>> v0 = np.random.rand(5, 4, 3)
|
|
>>> v1 = unit_vector(v0, axis=-1)
|
|
>>> v2 = v0 / np.expand_dims(np.sqrt(np.sum(v0*v0, axis=2)), 2)
|
|
>>> np.allclose(v1, v2)
|
|
True
|
|
>>> v1 = unit_vector(v0, axis=1)
|
|
>>> v2 = v0 / np.expand_dims(np.sqrt(np.sum(v0*v0, axis=1)), 1)
|
|
>>> np.allclose(v1, v2)
|
|
True
|
|
>>> v1 = np.empty((5, 4, 3))
|
|
>>> unit_vector(v0, axis=1, out=v1)
|
|
>>> np.allclose(v1, v2)
|
|
True
|
|
>>> list(unit_vector([]))
|
|
[]
|
|
>>> list(unit_vector([1]))
|
|
[1.0]
|
|
"""
|
|
if out is None:
|
|
data = np.array(data, dtype=np.float64, copy=True)
|
|
if data.ndim == 1:
|
|
data /= math.sqrt(np.dot(data, data))
|
|
return data
|
|
else:
|
|
if out is not data:
|
|
out[:] = np.array(data, copy=False)
|
|
data = out
|
|
length = np.atleast_1d(np.sum(data * data, axis))
|
|
np.sqrt(length, length)
|
|
if axis is not None:
|
|
length = np.expand_dims(length, axis)
|
|
data /= length
|
|
if out is None:
|
|
return data
|
|
|
|
|
|
def quaternion_from_matrix(matrix, isprecise=False):
|
|
"""Return quaternion from rotation matrix.
|
|
If isprecise is True, the input matrix is assumed to be a precise rotation
|
|
matrix and a faster algorithm is used.
|
|
>>> q = quaternion_from_matrix(np.identity(4), True)
|
|
>>> np.allclose(q, [1, 0, 0, 0])
|
|
True
|
|
>>> q = quaternion_from_matrix(np.diag([1, -1, -1, 1]))
|
|
>>> np.allclose(q, [0, 1, 0, 0]) or np.allclose(q, [0, -1, 0, 0])
|
|
True
|
|
>>> R = rotation_matrix(0.123, (1, 2, 3))
|
|
>>> q = quaternion_from_matrix(R, True)
|
|
>>> np.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786])
|
|
True
|
|
>>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0],
|
|
... [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]]
|
|
>>> q = quaternion_from_matrix(R)
|
|
>>> np.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611])
|
|
True
|
|
>>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0],
|
|
... [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]]
|
|
>>> q = quaternion_from_matrix(R)
|
|
>>> np.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603])
|
|
True
|
|
>>> R = random_rotation_matrix()
|
|
>>> q = quaternion_from_matrix(R)
|
|
>>> is_same_transform(R, quaternion_matrix(q))
|
|
True
|
|
>>> R = euler_matrix(0.0, 0.0, np.pi/2.0)
|
|
>>> np.allclose(quaternion_from_matrix(R, isprecise=False),
|
|
... quaternion_from_matrix(R, isprecise=True))
|
|
True
|
|
"""
|
|
M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]
|
|
if isprecise:
|
|
q = np.empty((4,))
|
|
t = np.trace(M)
|
|
if t > M[3, 3]:
|
|
q[0] = t
|
|
q[3] = M[1, 0] - M[0, 1]
|
|
q[2] = M[0, 2] - M[2, 0]
|
|
q[1] = M[2, 1] - M[1, 2]
|
|
else:
|
|
i, j, k = 1, 2, 3
|
|
if M[1, 1] > M[0, 0]:
|
|
i, j, k = 2, 3, 1
|
|
if M[2, 2] > M[i, i]:
|
|
i, j, k = 3, 1, 2
|
|
t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]
|
|
q[i] = t
|
|
q[j] = M[i, j] + M[j, i]
|
|
q[k] = M[k, i] + M[i, k]
|
|
q[3] = M[k, j] - M[j, k]
|
|
q *= 0.5 / math.sqrt(t * M[3, 3])
|
|
else:
|
|
m00 = M[0, 0]
|
|
m01 = M[0, 1]
|
|
m02 = M[0, 2]
|
|
m10 = M[1, 0]
|
|
m11 = M[1, 1]
|
|
m12 = M[1, 2]
|
|
m20 = M[2, 0]
|
|
m21 = M[2, 1]
|
|
m22 = M[2, 2]
|
|
# symmetric matrix K
|
|
K = np.array(
|
|
[
|
|
[m00 - m11 - m22, 0.0, 0.0, 0.0],
|
|
[m01 + m10, m11 - m00 - m22, 0.0, 0.0],
|
|
[m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
|
|
[m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],
|
|
]
|
|
)
|
|
K /= 3.0
|
|
# quaternion is eigenvector of K that corresponds to largest eigenvalue
|
|
w, V = np.linalg.eigh(K)
|
|
q = V[np.array([3, 0, 1, 2]), np.argmax(w)]
|
|
if q[0] < 0.0:
|
|
np.negative(q, q)
|
|
return q
|
|
|
|
|
|
def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True):
|
|
"""Return spherical linear interpolation between two quaternions.
|
|
>>> q0 = random_quaternion()
|
|
>>> q1 = random_quaternion()
|
|
>>> q = quaternion_slerp(q0, q1, 0)
|
|
>>> np.allclose(q, q0)
|
|
True
|
|
>>> q = quaternion_slerp(q0, q1, 1, 1)
|
|
>>> np.allclose(q, q1)
|
|
True
|
|
>>> q = quaternion_slerp(q0, q1, 0.5)
|
|
>>> angle = math.acos(np.dot(q0, q))
|
|
>>> (np.allclose(2, math.acos(np.dot(q0, q1)) / angle) or
|
|
np.allclose(2, math.acos(-np.dot(q0, q1)) / angle))
|
|
True
|
|
"""
|
|
q0 = unit_vector(quat0[:4])
|
|
q1 = unit_vector(quat1[:4])
|
|
if q0 is None or q1 is None:
|
|
raise ValueError("Input quaternions invalid.")
|
|
if fraction == 0.0:
|
|
return q0
|
|
elif fraction == 1.0:
|
|
return q1
|
|
d = np.dot(q0, q1)
|
|
if abs(abs(d) - 1.0) < _EPS:
|
|
return q0
|
|
if shortestpath and d < 0.0:
|
|
# invert rotation
|
|
d = -d
|
|
np.negative(q1, q1)
|
|
angle = math.acos(d) + spin * math.pi
|
|
if abs(angle) < _EPS:
|
|
return q0
|
|
isin = 1.0 / math.sin(angle)
|
|
q0 *= math.sin((1.0 - fraction) * angle) * isin
|
|
q1 *= math.sin(fraction * angle) * isin
|
|
q0 += q1
|
|
return q0
|
|
|
|
|
|
def quaternion_matrix(quaternion):
|
|
"""Return homogeneous rotation matrix from quaternion.
|
|
>>> M = quaternion_matrix([0.99810947, 0.06146124, 0, 0])
|
|
>>> np.allclose(M, rotation_matrix(0.123, [1, 0, 0]))
|
|
True
|
|
>>> M = quaternion_matrix([1, 0, 0, 0])
|
|
>>> np.allclose(M, np.identity(4))
|
|
True
|
|
>>> M = quaternion_matrix([0, 1, 0, 0])
|
|
>>> np.allclose(M, np.diag([1, -1, -1, 1]))
|
|
True
|
|
"""
|
|
q = np.array(quaternion, dtype=np.float64, copy=True)
|
|
n = np.dot(q, q)
|
|
if n < _EPS:
|
|
return np.identity(4)
|
|
q *= math.sqrt(2.0 / n)
|
|
q = np.outer(q, q)
|
|
return np.array(
|
|
[
|
|
[1.0 - q[2, 2] - q[3, 3], q[1, 2] - q[3, 0], q[1, 3] + q[2, 0], 0.0],
|
|
[q[1, 2] + q[3, 0], 1.0 - q[1, 1] - q[3, 3], q[2, 3] - q[1, 0], 0.0],
|
|
[q[1, 3] - q[2, 0], q[2, 3] + q[1, 0], 1.0 - q[1, 1] - q[2, 2], 0.0],
|
|
[0.0, 0.0, 0.0, 1.0],
|
|
]
|
|
)
|
|
|
|
|
|
def get_interpolated_poses(poseA, poseB, steps=10):
|
|
"""Return interpolation of poses with specified number of steps.
|
|
Args:
|
|
poseA
|
|
poseB
|
|
Return:
|
|
list of poses
|
|
"""
|
|
|
|
quatA = quaternion_from_matrix(poseA[:3, :3])
|
|
quatB = quaternion_from_matrix(poseB[:3, :3])
|
|
|
|
ts = np.linspace(0, 1, steps)
|
|
quats = [quaternion_slerp(quatA, quatB, t) for t in ts]
|
|
trans = [(1 - t) * poseA[:3, 3] + t * poseB[:3, 3] for t in ts]
|
|
|
|
posesAB = []
|
|
for quat, tran in zip(quats, trans):
|
|
pose = np.identity(4)
|
|
pose[:3, :3] = quaternion_matrix(quat)[:3, :3]
|
|
pose[:3, 3] = tran
|
|
posesAB.append(pose)
|
|
return posesAB
|
|
|
|
|
|
def get_interpolated_K(KA, KB, steps=10):
|
|
"""
|
|
Interpolate between two camera poses.
|
|
Args:
|
|
Returns:
|
|
"""
|
|
Ks = []
|
|
ts = np.linspace(0, 1, steps)
|
|
for t in ts:
|
|
newK = KA * (1.0 - t) + KB * t
|
|
Ks.append(newK)
|
|
return Ks
|
|
|
|
|
|
def get_interpolated_poses_many(
|
|
poses: TensorType["num_poses", 3, 4],
|
|
Ks: TensorType["num_poses", 3, 3],
|
|
steps_per_transition=10,
|
|
) -> Tuple[TensorType["num_poses", 3, 4], TensorType["num_poses", 3, 3]]:
|
|
"""Return interpolated poses for many camera poses.
|
|
|
|
Args:
|
|
poses: list of camera poses
|
|
Ks: list of camera intrinsics
|
|
steps_per_transition: number of steps per transition
|
|
|
|
Returns:
|
|
tuple of new poses and intrinsics
|
|
"""
|
|
traj = []
|
|
Ks = []
|
|
for idx in range(poses.shape[0] - 1):
|
|
poseA = poses[idx]
|
|
poseB = poses[idx + 1]
|
|
posesAB = get_interpolated_poses(poseA, poseB, steps=steps_per_transition)
|
|
traj += posesAB
|
|
Ks += get_interpolated_K(Ks[idx], Ks[idx + 1], steps_per_transition)
|
|
return torch.stack(traj, dim=0), torch.stack(Ks, dim=0)
|
|
|
|
|
|
def get_swirl_poses(pose):
|
|
"""
|
|
Returns a swirl around the camera poses.
|
|
# TODO(ethan): add parameters and make this work
|
|
"""
|
|
N_frames = 30 * 4
|
|
dx = np.linspace(0, 0.03, N_frames)
|
|
dy = np.linspace(0, -0.1, N_frames)
|
|
dz = np.linspace(0, 0.5, N_frames)
|
|
poses = np.tile(pose, (N_frames, 1, 1))
|
|
for i in range(N_frames):
|
|
poses[i, 0, 3] += dx[i]
|
|
poses[i, 1, 3] += dy[i]
|
|
poses[i, 2, 3] += dz[i]
|
|
return poses
|
|
|
|
|
|
def normalize(x):
|
|
"""Returns a normalized vector."""
|
|
return x / torch.linalg.norm(x)
|
|
|
|
|
|
def viewmatrix(lookat, up, pos):
|
|
"""Returns a camera transformation matrix.
|
|
|
|
Args:
|
|
lookat: The direction the camera is looking.
|
|
up: The upward direction of the camera.
|
|
pos: The position of the camera.
|
|
|
|
Returns:
|
|
torch.Tensor: A camera transformation matrix.
|
|
"""
|
|
vec2 = normalize(lookat)
|
|
vec1_avg = normalize(up)
|
|
vec0 = normalize(torch.cross(vec1_avg, vec2))
|
|
vec1 = normalize(torch.cross(vec2, vec0))
|
|
m = torch.stack([vec0, vec1, vec2, pos], 1)
|
|
return m
|
|
|
|
|
|
def get_distortion_params(
|
|
k1: float = 0.0,
|
|
k2: float = 0.0,
|
|
k3: float = 0.0,
|
|
k4: float = 0.0,
|
|
p1: float = 0.0,
|
|
p2: float = 0.0,
|
|
):
|
|
"""Returns a distortion parameters matrix.
|
|
|
|
Args:
|
|
k1: The first radial distortion parameter.
|
|
k2: The second radial distortion parameter.
|
|
k3: The third radial distortion parameter.
|
|
k4: The fourth radial distortion parameter.
|
|
p1: The first tangential distortion parameter.
|
|
p2: The second tangential distortion parameter.
|
|
Returns:
|
|
torch.Tensor: A distortion parameters matrix.
|
|
"""
|
|
return torch.Tensor([k1, k2, k3, k4, p1, p2])
|
|
|
|
|
|
# Addapted from MultiNeRF
|
|
# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509
|
|
def radial_and_tangential_undistort(
|
|
coords: TensorType["num_coords":..., 2],
|
|
distortion_params: TensorType["num_coords":..., 6],
|
|
eps: float = 1e-9,
|
|
max_iterations=10,
|
|
) -> TensorType["num_coords":..., 2]:
|
|
"""Computes undistorted coords given opencv distortion parameters.
|
|
|
|
Args:
|
|
coords: The distorted coordinates.
|
|
distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
|
|
eps: The epsilon for the convergence.
|
|
max_iterations: The maximum number of iterations to perform.
|
|
|
|
Returns:
|
|
The undistorted coordinates.
|
|
"""
|
|
|
|
# Initialize from the distorted point.
|
|
x = torch.clone(coords[..., 0])
|
|
y = torch.clone(coords[..., 1])
|
|
|
|
for _ in range(max_iterations):
|
|
fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(
|
|
x=x, y=y, xd=coords[..., 0], yd=coords[..., 1], distortion_params=distortion_params
|
|
)
|
|
denominator = fy_x * fx_y - fx_x * fy_y
|
|
x_numerator = fx * fy_y - fy * fx_y
|
|
y_numerator = fy * fx_x - fx * fy_x
|
|
step_x = torch.where(torch.abs(denominator) > eps, x_numerator / denominator, torch.zeros_like(denominator))
|
|
step_y = torch.where(torch.abs(denominator) > eps, y_numerator / denominator, torch.zeros_like(denominator))
|
|
|
|
x = x + step_x
|
|
y = y + step_y
|
|
|
|
return torch.stack([x, y], dim=-1)
|
|
|
|
|
|
# Adapted from MultiNeRF:
|
|
# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474
|
|
def _compute_residual_and_jacobian(
|
|
x: TensorType["num_coords":...],
|
|
y: TensorType["num_coords":...],
|
|
xd: TensorType["num_coords":...],
|
|
yd: TensorType["num_coords":...],
|
|
distortion_params: TensorType["num_coords":..., 6],
|
|
) -> Tuple[
|
|
TensorType["num_coords":...],
|
|
TensorType["num_coords":...],
|
|
TensorType["num_coords":...],
|
|
TensorType["num_coords":...],
|
|
TensorType["num_coords":...],
|
|
TensorType["num_coords":...],
|
|
]:
|
|
"""Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians.
|
|
|
|
Args:
|
|
x: The updated x coordinates.
|
|
y: The updated y coordinates.
|
|
xd: The distorted x coordinates.
|
|
yd: The distorted y coordinates.
|
|
distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
|
|
|
|
Returns:
|
|
The residuals (fx, fy) and jacobians (fx_x, fx_y, fy_x, fy_y).
|
|
"""
|
|
|
|
k1 = distortion_params[..., 0]
|
|
k2 = distortion_params[..., 1]
|
|
k3 = distortion_params[..., 2]
|
|
k4 = distortion_params[..., 3]
|
|
p1 = distortion_params[..., 4]
|
|
p2 = distortion_params[..., 5]
|
|
|
|
# let r(x, y) = x^2 + y^2;
|
|
# d(x, y) = 1 + k1 * r(x, y) + k2 * r(x, y) ^2 + k3 * r(x, y)^3 +
|
|
# k4 * r(x, y)^4;
|
|
r = x * x + y * y
|
|
d = 1.0 + r * (k1 + r * (k2 + r * (k3 + r * k4)))
|
|
|
|
# The perfect projection is:
|
|
# xd = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2);
|
|
# yd = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2);
|
|
#
|
|
# Let's define
|
|
#
|
|
# fx(x, y) = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2) - xd;
|
|
# fy(x, y) = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2) - yd;
|
|
#
|
|
# We are looking for a solution that satisfies
|
|
# fx(x, y) = fy(x, y) = 0;
|
|
fx = d * x + 2 * p1 * x * y + p2 * (r + 2 * x * x) - xd
|
|
fy = d * y + 2 * p2 * x * y + p1 * (r + 2 * y * y) - yd
|
|
|
|
# Compute derivative of d over [x, y]
|
|
d_r = k1 + r * (2.0 * k2 + r * (3.0 * k3 + r * 4.0 * k4))
|
|
d_x = 2.0 * x * d_r
|
|
d_y = 2.0 * y * d_r
|
|
|
|
# Compute derivative of fx over x and y.
|
|
fx_x = d + d_x * x + 2.0 * p1 * y + 6.0 * p2 * x
|
|
fx_y = d_y * x + 2.0 * p1 * x + 2.0 * p2 * y
|
|
|
|
# Compute derivative of fy over x and y.
|
|
fy_x = d_x * y + 2.0 * p2 * y + 2.0 * p1 * x
|
|
fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y
|
|
|
|
return fx, fy, fx_x, fx_y, fy_x, fy_y
|
|
|
|
|
|
def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
|
|
"""Compute the rotation matrix that rotates vector a to vector b.
|
|
|
|
Args:
|
|
a: The vector to rotate.
|
|
b: The vector to rotate to.
|
|
Returns:
|
|
The rotation matrix.
|
|
"""
|
|
a = a / torch.linalg.norm(a)
|
|
b = b / torch.linalg.norm(b)
|
|
v = torch.cross(a, b)
|
|
c = torch.dot(a, b)
|
|
# If vectors are exactly opposite, we add a little noise to one of them
|
|
if c < -1 + 1e-8:
|
|
eps = (torch.rand(3) - 0.5) * 0.01
|
|
return rotmat(a + eps, b)
|
|
s = torch.linalg.norm(v)
|
|
skew_sym_mat = torch.Tensor(
|
|
[
|
|
[0, -v[2], v[1]],
|
|
[v[2], 0, -v[0]],
|
|
[-v[1], v[0], 0],
|
|
]
|
|
)
|
|
return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
|
|
|
|
|
|
def auto_orient_poses(
|
|
poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up"] = "up"
|
|
) -> TensorType["num_poses":..., 3, 4]:
|
|
"""Orients and centers the poses. We provide two methods for orientation: pca and up.
|
|
|
|
pca: Orient the poses so that the principal component of the points is aligned with the axes.
|
|
This method works well when all of the cameras are in the same plane.
|
|
|
|
up: Orient the poses so that the average up vector is aligned with the z axis.
|
|
This method works well when images are not at arbitrary angles.
|
|
|
|
Args:
|
|
poses: The poses to orient.
|
|
method: The method to use for orientation. Either "pca" or "up".
|
|
|
|
Returns:
|
|
The oriented poses.
|
|
"""
|
|
|
|
translation = poses[..., :3, 3]
|
|
|
|
mean_translation = torch.mean(translation, dim=0)
|
|
translation = translation - mean_translation
|
|
|
|
if method == "pca":
|
|
_, eigvec = torch.linalg.eigh(translation.T @ translation)
|
|
eigvec = torch.flip(eigvec, dims=(-1,))
|
|
|
|
if torch.linalg.det(eigvec) < 0:
|
|
eigvec[:, 2] = -eigvec[:, 2]
|
|
|
|
transform = torch.cat([eigvec, eigvec @ -mean_translation[..., None]], dim=-1)
|
|
oriented_poses = transform @ poses
|
|
|
|
if oriented_poses.mean(axis=0)[2, 1] < 0:
|
|
oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
|
|
elif method == "up":
|
|
up = torch.mean(poses[:, :3, 1], dim=0)
|
|
up = up / torch.linalg.norm(up)
|
|
|
|
rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
|
|
transform = torch.cat([rotation, rotation @ -mean_translation[..., None]], dim=-1)
|
|
oriented_poses = transform @ poses
|
|
|
|
return oriented_poses
|