tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code
2ca0695fef
nerfstudio/nerfactory/cameras/utils.py
Matthew Tancik 6ee6934581
Update dataset loaders (#319) 
* Update dataset loaders

* Set ngp default to blender
2022-08-31 17:30:56 -07:00

462 lines
14 KiB
Python
Raw Blame History

# 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 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
View Git Blame Copy Permalink