Nerfstudio Gaussian Splatting 相机姿态可微/求导的具体实现方式

源码

代码

class _ProjectGaussians(Function):
"""Project 3D gaussians to 2D."""
@staticmethod
def forward(
ctx,
means3d: Float[Tensor, "*batch 3"],
scales: Float[Tensor, "*batch 3"],
glob_scale: float,
quats: Float[Tensor, "*batch 4"],
viewmat: Float[Tensor, "4 4"],
fx: float,
fy: float,
cx: float,
cy: float,
img_height: int,
img_width: int,
block_width: int,
clip_thresh: float = 0.01,
):
num_points = means3d.shape[-2]
if num_points < 1 or means3d.shape[-1] != 3:
raise ValueError(f"Invalid shape for means3d: {means3d.shape}")
(
cov3d,
xys,
depths,
radii,
conics,
compensation,
num_tiles_hit,
) = _C.project_gaussians_forward(
num_points,
means3d,
scales,
glob_scale,
quats,
viewmat,
fx,
fy,
cx,
cy,
img_height,
img_width,
block_width,
clip_thresh,
)
# Save non-tensors.
ctx.img_height = img_height
ctx.img_width = img_width
ctx.num_points = num_points
ctx.glob_scale = glob_scale
ctx.fx = fx
ctx.fy = fy
ctx.cx = cx
ctx.cy = cy
# Save tensors.
ctx.save_for_backward(
means3d,
scales,
quats,
viewmat,
cov3d,
radii,
conics,
compensation,
)
return (xys, depths, radii, conics, compensation, num_tiles_hit, cov3d)
@staticmethod
def backward(
ctx,
v_xys,
v_depths,
v_radii,
v_conics,
v_compensation,
v_num_tiles_hit,
v_cov3d,
):
(
means3d,
scales,
quats,
viewmat,
cov3d,
radii,
conics,
compensation,
) = ctx.saved_tensors
(v_cov2d, v_cov3d, v_mean3d, v_scale, v_quat) = _C.project_gaussians_backward(
ctx.num_points,
means3d,
scales,
ctx.glob_scale,
quats,
viewmat,
ctx.fx,
ctx.fy,
ctx.cx,
ctx.cy,
ctx.img_height,
ctx.img_width,
cov3d,
radii,
conics,
compensation,
v_xys,
v_depths,
v_conics,
v_compensation,
)
if viewmat.requires_grad:
v_viewmat = torch.zeros_like(viewmat)
R = viewmat[..., :3, :3]
# Denote ProjectGaussians for a single Gaussian (mean3d, q, s)
# viemwat = [R, t] as:
#
#   f(mean3d, q, s, R, t, intrinsics)
#       = g(R @ mean3d + t,
#           R @ cov3d_world(q, s) @ R^T ))
#
# Then, the Jacobian w.r.t., t is:
#
#   d f / d t = df / d mean3d @ R^T
#
# and, in the context of fine tuning camera poses, it is reasonable
# to assume that
#
#   d f / d R_ij =~ \sum_l d f / d t_l * d (R @ mean3d)_l / d R_ij
#                = d f / d_t_i * mean3d[j]
#
# Gradients for R and t can then be obtained by summing over
# all the Gaussians.
v_mean3d_cam = torch.matmul(v_mean3d, R.transpose(-1, -2))
# gradient w.r.t. view matrix translation
v_viewmat[..., :3, 3] = v_mean3d_cam.sum(-2)
# gradent w.r.t. view matrix rotation
for j in range(3):
for l in range(3):
v_viewmat[..., j, l] = torch.dot(
v_mean3d_cam[..., j], means3d[..., l]
)
else:
v_viewmat = None
# Return a gradient for each input.
return (
# means3d: Float[Tensor, "*batch 3"],
v_mean3d,
# scales: Float[Tensor, "*batch 3"],
v_scale,
# glob_scale: float,
None,
# quats: Float[Tensor, "*batch 4"],
v_quat,
# viewmat: Float[Tensor, "4 4"],
v_viewmat,
# fx: float,
None,
# fy: float,
None,
# cx: float,
None,
# cy: float,
None,
# img_height: int,
None,
# img_width: int,
None,
# block_width: int,
None,
# clip_thresh,
None,
)

解释

这段代码是用于计算视图矩阵（view matrix）的梯度，视图矩阵通常用于3D图形和计算机视觉中，用于将世界坐标转换为相机坐标。在深度学习中，视图矩阵通常用于表示相机的位置和朝向。

代码的目的是计算一个由多个高斯分布（每个分布有均值mean3d、协方差cov3d_world、以及其他参数q和s）组成的投影高斯分布的梯度。这些高斯分布通过视图矩阵viewmat（包含旋转矩阵R和平移向量t）进行投影，并考虑了相机内参intrinsics。

代码的逻辑如下：

1. 检查视图矩阵是否需要梯度：if viewmat.requires_grad: 检查视图矩阵是否需要计算梯度。如果需要，则继续执行后续代码。

2. 初始化梯度张量：v_viewmat = torch.zeros_like(viewmat) 创建一个与视图矩阵形状相同的全零张量，用于存储梯度。

3. 提取旋转矩阵：R = viewmat[..., :3, :3] 从视图矩阵中提取旋转矩阵部分。

4. 计算均值在相机坐标系下的表示：v_mean3d_cam = torch.matmul(v_mean3d, R.transpose(-1, -2)) 将世界坐标系下的均值通过旋转矩阵转换到相机坐标系。

5. 计算相对于平移向量的梯度：v_viewmat[..., :3, 3] = v_mean3d_cam.sum(-2) 计算所有高斯分布的梯度相对于平移向量的和。

6. 计算相对于旋转矩阵的梯度：通过两层循环，计算所有高斯分布的梯度相对于旋转矩阵的和。外层循环遍历旋转矩阵的行，内层循环遍历旋转矩阵的列。对于每个元素(j, l)，计算v_mean3d_cam[..., j]和means3d[..., l]的点积，并将结果赋给v_viewmat[..., j, l]。

7. 返回梯度：最后，函数返回计算得到的梯度v_viewmat。

总的来说，这段代码是用于计算视图矩阵的梯度，这些梯度可以用于优化相机的位置和朝向，以最小化投影高斯分布的某种损失函数。在实际应用中，这通常在训练神经网络时用于端到端的优化，以提高模型的性能。