Rendering

splax.render is the rendering entry point. The call returns an (image, alpha) pair, where image is the (H, W, 3) render and alpha the (H, W) accumulated coverage. Gradients are covered under Training.

We first load a splat and prepare a view matrix:

from functools import partial

import jax
import jax.numpy as jnp
import splax

SCENE = "https://huggingface.co/datasets/amacati/splax-test-data/resolve/main/scenes/lego.ply"
splats = splax.io.load_ply(splax.io.fetch(SCENE))
means, log_scales, quats, sh_colors, logit_opacities = splats
H, W, fx, fy = 400, 400, 400.0, 400.0
viewmat = splax.utils.look_at(jnp.array((0.0, -3.0, 1.0)), jnp.zeros(3), up=(0.0, 0.0, 1.0))
eyes = jnp.array([(0.0, -3.0, 1.0), (3.0, 0.0, 1.0)])
viewmats = splax.utils.look_at(eyes, jnp.zeros(3), up=(0.0, 0.0, 1.0))
slices = ((100, 1000), (1000, 1500))
poses = jnp.broadcast_to(jnp.eye(4), (len(slices), 4, 4))
pose_batch = jnp.broadcast_to(poses, (len(viewmats), len(slices), 4, 4))

img, _ = splax.render(
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    viewmat=viewmat,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
)  # (H, W, 3)

Inputs

render takes unconstrained parameters, so an optimizer can update them directly.

Argument	Shape	Meaning
means3d	(N, 3)	World positions
log_scales	(N, 3)	Log of the per-axis scales
quats	(N, 4)	wxyz quaternions
sh_colors	(N, 3)	Degree-0 SH color coefficients, 0 is mid grey
logit_opacities	(N,)	Opacity logits

splax.io.apply_activations and splax.io.invert_activations convert to and from linear scales, RGB colors, and [0, 1] opacities, see IO.

Camera conventions

viewmat is a (4, 4) world-to-camera matrix in the OpenCV convention (+z forward, +y down, +x right), consistent with COLMAP's output files. NeRF and OpenGL poses (-z forward) must be converted first with splax.utils.nerf_camera.

f is the focal length (fx, fy) in pixels and c is the principal point (cx, cy) in pixels, where the optical axis meets the image plane. It defaults to the image center (W / 2, H / 2). Calibrated real cameras, e.g. with COLMAP intrinsics, provide their own off-center values. img_shape is (H, W). glob_scale multiplies every gaussian scale, and clip_thresh is the near-plane depth cutoff.

dist holds the Brown-Conrady coefficients (k1, k2, p1, p2, k3) that COLMAP writes for its SIMPLE_RADIAL, RADIAL, OPENCV and FULL_OPENCV cameras, and defaults to zero for an ideal lens, see Lens distortion. Fisheye and other non-polynomial models are not covered.

img_shape, f, c, and dist size the kernel launch, so they are static under jax.jit, see Jitting.

Backgrounds

background is a 3-dimensional RGB color composited behind the splat where transmittance remains. It is a constant and is not differentiated.

Lens distortion

Real cameras are not ideal, and a reconstruction fitted to their photographs is only correct when the renderer reproduces the same lens. dist renders the scene through a Brown-Conrady lens.

render_dist = jax.jit(splax.render, static_argnames=("img_shape", "f", "c", "dist"))
distorted, _ = render_dist(
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    viewmat=viewmat,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
    dist=(-0.3, 0.08, 0.0, 0.0, 0.0),
)

The coefficients are static and carry no gradient, so a scene can be fitted through its lens, but the lens itself cannot be calibrated this way.

Antialiased mode

antialiased=True applies the Mip-Splatting opacity compensation, cancelling the area inflation that thin gaussians gain from the projection. Use the same setting at inference that a model was trained with.

Dynamic scene composition

Composed scenes can move whole sections of gaussians with rigid transforms to immitate moving objects without copying the splats. gaussian_transforms is a (K, 4, 4) stack of world-space transforms and gaussian_slices the K matching, non-overlapping (start, stop) index ranges. The gaussians in slice k move by gaussian_transforms[k]. Everything outside the slices stays static.

img, _ = splax.render(
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    viewmat=viewmat,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
    gaussian_transforms=poses,  # (K, 4, 4)
    gaussian_slices=((100, 1000), (1000, 1500)),
)

Batched dynamics work through jax.vmap over the transform stack. Every batch element renders the same shared splat with its objects at different poses.

move = partial(
    splax.render,
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    viewmat=viewmat,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
    gaussian_slices=slices,
)
imgs, _ = jax.vmap(move)(gaussian_transforms=pose_batch)  # (B, K, 4, 4) -> (B, H, W, 3)

Examples has a runnable version that joins two splats and orbits one of them.

Omitting both arguments renders the splat as one static scene. The transforms are differentiable, see object pose gradients.

Jitting

img_shape, f, c, and dist size the kernel launch and gaussian_slices indexes it, so all five are static. Under jax.jit either declare them or close over them, otherwise the call raises.

render_jit = jax.jit(splax.render, static_argnames=("img_shape", "f", "c"))
img, _ = render_jit(
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    viewmat=viewmat,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
)

Closing over them with functools.partial leaves the batched argument as the only input, which is what jax.vmap maps over. Keyword arguments map along their leading axis.

render_at = partial(
    splax.render,
    means,
    log_scales,
    quats,
    sh_colors,
    logit_opacities,
    background=jnp.ones(3),
    img_shape=(H, W),
    f=(fx, fy),
)
imgs, _ = jax.jit(jax.vmap(render_at))(viewmat=viewmats)  # (B, H, W, 3)

A static value is baked into the compiled kernel, so a resolution sweep or a change of the slice layout compiles once per distinct value.

Low-level primitives

splax.render composes two primitives that are public in their own right.

..Warning:: All low-level primitives consume activated arrays, not unconstrained parameters. Use splax.io.apply_activations/ splax.io.invert_activations to convert the parameters.

• splax.project maps gaussians to the 2D screen-space.
• splax.rasterize blends the projected gaussians into a (H, W, 3) image and its (H, W) alpha.
• splax.rasterize_depth blends into a (H, W, 4) image whose fourth channel is the expected depth, plus the same (H, W) alpha.

Both rasterization primitives must be passed the same opacities splax.project ran on. Failing to do so will result in crashes or incorrect renderings. Rasterization takes an antialiased keyword, so the compensation needs no separate call.

splax maintains a scratch memory pool for intermediate arrays used by the backend. splax.clear_cache releases this memory, for example before switching to a different workload size.