diff --git a/splax/_project/_kernels.py b/splax/_project/_kernels.py index 6c96b7d..f82b5ac 100644 --- a/splax/_project/_kernels.py +++ b/splax/_project/_kernels.py @@ -208,17 +208,9 @@ def _project_kernel( return V3 = M * wp.transpose(M) - # EWA projection of the covariance - tan_fovx = 0.5 * wp.float32(img_w) / fx - tan_fovy = 0.5 * wp.float32(img_h) / fy - lim_x = 1.3 * tan_fovx - lim_y = 1.3 * tan_fovy - tx = p_view[0] - ty = p_view[1] - tz = p_view[2] - tx = tz * wp.min(lim_x, wp.max(-lim_x, tx / tz)) - ty = tz * wp.min(lim_y, wp.max(-lim_y, ty / tz)) - J = _ewa_jacobian(tx, ty, tz, fx, fy, k1, k2, p1, p2, k3) + # EWA projection of the covariance, with the Jacobian at the position clamped to the FOV limits + p_clamped = _clamp_to_fov(p_view, _fov_limits(img_w, img_h, fx, fy)) + J = _ewa_jacobian(p_clamped[0], p_clamped[1], p_clamped[2], fx, fy, k1, k2, p1, p2, k3) T = J * W cov = T * V3 * wp.transpose(T) # 0.3 px screen-space dilation, the standard 3DGS low-pass guard @@ -272,7 +264,7 @@ def _project_kernel( # region backward kernels # cov3d is recomputed in the backward instead of cached, saving memory on the order of 6N. The EWA -# Jacobian is rebuilt from the unclamped camera-space position, the standard gsplat approximation. +# Jacobian is rebuilt at the position clamped to the FOV limits. def _project_bwd_warp( @@ -290,6 +282,8 @@ def _project_bwd_warp( n_gaussians: int, n_transforms: int, has_transforms: bool, + img_h: int, + img_w: int, fx: float, fy: float, k1: float, @@ -357,6 +351,8 @@ def _project_bwd_warp( sel_vdepth, sel_vconic, sel_transforms, + img_h, + img_w, fx, fy, k1, @@ -411,6 +407,8 @@ def _project_bwd_kernel( sel_vdepth: wp.bool, sel_vconic: wp.bool, sel_transforms: wp.bool, + img_h: wp.int32, + img_w: wp.int32, fx: wp.float32, fy: wp.float32, k1: wp.float32, @@ -460,8 +458,10 @@ def _project_bwd_kernel( V = M * wp.transpose(M) R = wp.quat_to_matrix(quat) s = glob_scale * scales[s_idx] - # Note that J is rebuilt from the unclamped p (the gsplat approximation) - J = _ewa_jacobian(p[0], p[1], p[2], fx, fy, k1, k2, p1, p2, k3) + # J at the position clamped to the FOV limits + lim = _fov_limits(img_w, img_h, fx, fy) + p_clamped = _clamp_to_fov(p, lim) + J = _ewa_jacobian(p_clamped[0], p_clamped[1], p_clamped[2], fx, fy, k1, k2, p1, p2, k3) T = J * W # Gradient computation, starting with the projection conic = conics[c_idx] @@ -471,7 +471,7 @@ def _project_bwd_kernel( S = X * G * X vcov2d = wp.vec3(-S[0, 0], -(S[0, 1] + S[1, 0]), -S[1, 1]) v_p, v_T, v_V = _proj_vjp( - p, T, W, V, fx, fy, k1, k2, p1, p2, k3, v_xy_in[vx_idx], v_depth_in[vd_idx], vcov2d + p, T, W, V, fx, fy, k1, k2, p1, p2, k3, lim, v_xy_in[vx_idx], v_depth_in[vd_idx], vcov2d ) v_mean_world = wp.transpose(W) * v_p # world-space mean cotangent v_M_world = (v_V + wp.transpose(v_V)) * M # covariance-factor cotangent @@ -532,6 +532,7 @@ def _proj_vjp( p1: wp.float32, p2: wp.float32, k3: wp.float32, + lim: wp.vec2, v_xy: wp.vec2, v_depth: wp.float32, vcov2d: wp.vec3, @@ -549,16 +550,10 @@ def _proj_vjp( tz = p[2] rz = 1.0 / tz rz2 = rz * rz - rz3 = rz2 * rz rw = 1.0 / (tz + 1e-6) nx = tx * rz ny = ty * rz - r2 = nx * nx + ny * ny d = _lens_jacobian(nx, ny, k1, k2, p1, p2, k3) - a11 = fx * d[0] - a12 = fx * d[1] - a21 = fy * d[1] - a22 = fy * d[2] # center path, the screen cotangent crosses the lens before the divide vpx = fx * v_xy[0] @@ -578,32 +573,49 @@ def _proj_vjp( v_T = 2.0 * v_cov * T * V # v_cov is symmetric, so v_cov T V + v_cov^T T V = 2 v_cov T V v_J = v_T * wp.transpose(W) - # covariance path, the perspective terms with the lens jacobian held fixed - v_t_x = -rz2 * (a11 * v_J[0, 2] + a21 * v_J[1, 2]) - v_t_y = -rz2 * (a12 * v_J[0, 2] + a22 * v_J[1, 2]) + # covariance path, at the clamped normalized coordinates + u = wp.min(lim[0], wp.max(-lim[0], nx)) + v = wp.min(lim[1], wp.max(-lim[1], ny)) + dc = _lens_jacobian(u, v, k1, k2, p1, p2, k3) + a11 = fx * dc[0] + a12 = fx * dc[1] + a21 = fy * dc[1] + a22 = fy * dc[2] + # the depth with u and v held fixed v_t_z = -rz2 * (a11 * v_J[0, 0] + a12 * v_J[0, 1] + a21 * v_J[1, 0] + a22 * v_J[1, 1]) - v_t_z += 2.0 * rz3 * ((a11 * tx + a12 * ty) * v_J[0, 2] + (a21 * tx + a22 * ty) * v_J[1, 2]) - - # covariance path, the lens jacobian itself moving with the normalized coordinates - va11 = v_J[0, 0] * rz - v_J[0, 2] * tx * rz2 - va12 = v_J[0, 1] * rz - v_J[0, 2] * ty * rz2 - va21 = v_J[1, 0] * rz - v_J[1, 2] * tx * rz2 - va22 = v_J[1, 1] * rz - v_J[1, 2] * ty * rz2 + v_t_z += rz2 * ((a11 * u + a12 * v) * v_J[0, 2] + (a21 * u + a22 * v) * v_J[1, 2]) + # u and v, the perspective terms with the lens jacobian held fixed + v_u = -rz * (a11 * v_J[0, 2] + a21 * v_J[1, 2]) + v_v = -rz * (a12 * v_J[0, 2] + a22 * v_J[1, 2]) + + # u and v, the lens jacobian itself moving with the normalized coordinates + va11 = (v_J[0, 0] - v_J[0, 2] * u) * rz + va12 = (v_J[0, 1] - v_J[0, 2] * v) * rz + va21 = (v_J[1, 0] - v_J[1, 2] * u) * rz + va22 = (v_J[1, 1] - v_J[1, 2] * v) * rz vd11 = fx * va11 vd12 = fx * va12 + fy * va21 vd22 = fy * va22 + r2 = u * u + v * v g = k1 + r2 * (2.0 * k2 + 3.0 * k3 * r2) h = 2.0 * k2 + 6.0 * k3 * r2 # d(g)/d(r2) # d11 and d22 share their cross derivative with d12, so three terms cover all six partials - dxx = 6.0 * nx * g + 4.0 * nx * nx * nx * h + 6.0 * p2 - dxy = 2.0 * ny * g + 4.0 * nx * nx * ny * h + 2.0 * p1 - dyx = 2.0 * nx * g + 4.0 * nx * ny * ny * h + 2.0 * p2 - dyy = 6.0 * ny * g + 4.0 * ny * ny * ny * h + 6.0 * p1 - v_nx = vd11 * dxx + vd12 * dxy + vd22 * dyx - v_ny = vd11 * dxy + vd12 * dyx + vd22 * dyy - v_t_x += v_nx * rz - v_t_y += v_ny * rz - v_t_z -= (v_nx * tx + v_ny * ty) * rz2 + dxx = 6.0 * u * g + 4.0 * u * u * u * h + 6.0 * p2 + dxy = 2.0 * v * g + 4.0 * u * u * v * h + 2.0 * p1 + dyx = 2.0 * u * g + 4.0 * u * v * v * h + 2.0 * p2 + dyy = 6.0 * v * g + 4.0 * v * v * v * h + 6.0 * p1 + v_u += vd11 * dxx + vd12 * dxy + vd22 * dyx + v_v += vd11 * dxy + vd12 * dyx + vd22 * dyy + + # through the clamp, u = x / z and v = y / z where they are not clamped + v_t_x = wp.float32(0.0) + v_t_y = wp.float32(0.0) + if nx > -lim[0] and nx < lim[0]: + v_t_x = v_u * rz + v_t_z -= v_u * tx * rz2 + if ny > -lim[1] and ny < lim[1]: + v_t_y = v_v * rz + v_t_z -= v_v * ty * rz2 v_p = wp.vec3(vvx + v_t_x, vvy + v_t_y, vvz + v_t_z) return v_p, v_T, v_V @@ -772,6 +784,25 @@ def _lens_jacobian( ) +@wp.func +def _fov_limits(img_w: wp.int32, img_h: wp.int32, fx: wp.float32, fy: wp.float32) -> wp.vec2: + """Bounds of |x / z| and |y / z| at which the projection evaluates the EWA Jacobian. + + The affine EWA approximation inflates gaussians far outside the image, so the Jacobian is taken + at the camera-space position clamped to 1.3 times the half field of view. + """ + return wp.vec2(1.3 * (0.5 * wp.float32(img_w) / fx), 1.3 * (0.5 * wp.float32(img_h) / fy)) + + +@wp.func +def _clamp_to_fov(p: wp.vec3, lim: wp.vec2) -> wp.vec3: + """Clamp a camera-space position to the FOV limits at its depth.""" + tz = p[2] + tx = tz * wp.min(lim[0], wp.max(-lim[0], p[0] / tz)) + ty = tz * wp.min(lim[1], wp.max(-lim[1], p[1] / tz)) + return wp.vec3(tx, ty, tz) + + @wp.func def _ewa_jacobian( tx: wp.float32, diff --git a/splax/_project/_project.py b/splax/_project/_project.py index e51b5d2..1193243 100644 --- a/splax/_project/_project.py +++ b/splax/_project/_project.py @@ -306,6 +306,8 @@ def _project_bwd( n, k, has_transforms, + img_shape[0], + img_shape[1], f[0], f[1], dist[0], diff --git a/tests/unit/test_project.py b/tests/unit/test_project.py index 84fa5d3..2e36846 100644 --- a/tests/unit/test_project.py +++ b/tests/unit/test_project.py @@ -289,6 +289,31 @@ def test_project_grad_distorted(): assert_finite_difference(loss, args, grads, eps=1e-3, rtol=5e-4) +@pytest.mark.parametrize( + "lens", [{}, {"dist": (-0.3, 0.08, 2e-3, -1e-3, 0.0)}], ids=["pinhole", "distorted"] +) +def test_project_grad_clamped(lens: dict): + """Match the gradient of gaussians beyond the FOV limits against finite differences.""" + H = W = 96 + kw = camera(H, W) | lens + lim = 1.3 * 0.5 * W / kw["f"][0] # FOV limits in a square image + axis = jnp.linspace(-1.75, 1.75, 8) * lim # Place Gaussians inside and outside the FOV limit + uv = jnp.stack(jnp.meshgrid(axis, axis), -1).reshape(-1, 2) + z = 2.0 + means = jnp.concatenate([uv, jnp.ones((64, 1))], 1) * z - VIEWMAT[:3, 3] + keys = jax.random.split(jax.random.key(0)) + scales = jax.random.uniform(keys[0], (64, 3), minval=0.3, maxval=0.8) + quats = jax.random.normal(keys[1], (64, 4)) + args = (means, scales, quats, VIEWMAT) + project = jax.jit(partial(splax.project, opacities=jnp.ones(64), **kw)) + + def loss(*args: jax.Array) -> jax.Array: + return jnp.sum(project(*args)[3] ** 2) + + grads = jax.jit(jax.grad(loss, argnums=(0, 1, 2, 3)))(*args) + assert_finite_difference(loss, args, grads) + + def test_project_grad_distorted_differs_from_pinhole(): """The lens must move the gradient, not only the forward projection.""" n, H, W = 400, 96, 96