Examples¶
Render a scene¶
Fetch a splat and write a single image. The camera is framed from the splat's own extent, so the script works for any scene in the dataset without per-scene metadata.
"""Render a splat from the splax test-data repository to an image.
The scene is downloaded into the local splax cache on first use and reused afterwards. The camera is
framed from the splat itself, so no per-scene metadata is needed.
Usage:
python examples/render_scene.py
python examples/render_scene.py --scene lego --out lego.png --res 800
"""
from __future__ import annotations
import argparse
import logging
from functools import partial
from pathlib import Path
import imageio.v3 as iio
import jax
import jax.numpy as jnp
import numpy as np
import splax
logger = logging.getLogger(__name__)
EXAMPLES = Path(__file__).parent
BASE = "https://huggingface.co/datasets/amacati/splax-test-data/resolve/main"
def frame_camera(means: jax.Array, direction: tuple, up: tuple) -> np.ndarray:
"""Look at the splat centre from ``direction``, backed off far enough to frame all of it.
Args:
means: Gaussian centers, shape ``(N, 3)``.
direction: World-space direction from the centre towards the camera.
up: World up direction.
Returns:
A ``(4, 4)`` world-to-camera matrix.
"""
centre = np.asarray(means.mean(axis=0))
radius = float(jnp.linalg.norm(means - centre, axis=-1).max())
offset = np.asarray(direction, float)
return splax.utils.look_at(centre + 2.5 * radius * offset / np.linalg.norm(offset), centre, up)
def main(scene: str, out: Path, res: int, fov: float, direction: tuple, up: tuple):
splats = splax.io.load_ply(splax.io.fetch(f"{BASE}/scenes/{scene}.ply"))
logger.info(f"loaded {splats[0].shape[0]} gaussians from {scene}.ply")
focal = 0.5 * res / np.tan(0.5 * np.deg2rad(fov))
render = jax.jit(
partial(splax.render, img_shape=(res, res), f=(focal, focal), background=jnp.ones(3))
)
img, _ = render(*splats, viewmat=jnp.asarray(frame_camera(splats[0], direction, up)))
iio.imwrite(out, np.asarray(jnp.clip(img, 0.0, 1.0) * 255, np.uint8))
logger.info(f"wrote {out}")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--scene", default="lego", help="scene name under scenes/ in the dataset")
parser.add_argument("--out", type=Path, default=EXAMPLES / "render.png")
parser.add_argument("--res", type=int, default=800)
parser.add_argument("--fov", type=float, default=40.0, help="horizontal field of view (deg)")
parser.add_argument("--direction", type=float, nargs=3, default=(1.0, -1.0, 0.6))
parser.add_argument("--up", type=float, nargs=3, default=(0.0, 0.0, 1.0))
args = parser.parse_args()
main(args.scene, args.out, args.res, args.fov, tuple(args.direction), tuple(args.up))
Join two splats and move one¶
Concatenate two copies of the lego splat into one set of arrays and drive the second with a rigid
transform, writing a GIF of the orbit. Both copies share every kernel launch, and only the
(1, 4, 4) transform changes per frame. See
dynamic scene composition for the mechanism.
"""Join two splats into one scene and drive one of them with a rigid transform.
Two copies of the lego splat are concatenated into a single set of arrays. The second copy is
declared as a movable slice, so it follows a pose while the first stays put. The splat is never
copied per frame, only the ``(1, 4, 4)`` transform changes.
Usage:
python examples/compose_splats.py
python examples/compose_splats.py --out orbit.gif --frames 60
"""
from __future__ import annotations
import argparse
import logging
import os
from functools import partial
from pathlib import Path
import imageio.v3 as iio
import jax
import jax.numpy as jnp
import numpy as np
os.environ.setdefault("SCIPY_ARRAY_API", "1")
from scipy.spatial.transform import RigidTransform, Rotation
import splax
logger = logging.getLogger(__name__)
EXAMPLES = Path(__file__).parent
BASE = "https://huggingface.co/datasets/amacati/splax-test-data/resolve/main"
def orbit_pose(centre: np.ndarray, offset: np.ndarray, angle: float) -> np.ndarray:
"""Yaw a splat about its own centre and carry it around by ``offset``.
Args:
centre: World position the splat rotates about, shape ``(3,)``.
offset: World translation applied after the rotation, shape ``(3,)``.
angle: Yaw about the world up axis, in radians.
Returns:
A ``(4, 4)`` world-space rigid transform.
"""
to_origin = RigidTransform.from_translation(-centre)
spin = RigidTransform.from_rotation(Rotation.from_euler("z", angle))
return (RigidTransform.from_translation(centre + offset) * spin * to_origin).as_matrix()
def main(out: Path, res: int, fov: float, frames: int, distance: float, up: tuple):
splat = splax.io.load_ply(splax.io.fetch(f"{BASE}/scenes/lego.ply"))
n = splat[0].shape[0]
centre = np.asarray(splat[0].mean(axis=0))
radius = float(jnp.linalg.norm(splat[0] - centre, axis=-1).max())
logger.info(f"loaded {n} gaussians, composing a scene of {2 * n}")
# One splat holding both copies back to back. Gaussians [n, 2n) follow transform 0, the rest
# stay static, so the two halves share every kernel launch.
scene = tuple(jnp.concatenate([array, array]) for array in splat)
slices = ((n, 2 * n),)
reach = radius * (1.0 + distance)
eye = centre + 2.5 * reach * np.array([1.0, -1.0, 0.5]) / np.linalg.norm([1.0, -1.0, 0.5])
focal = 0.5 * res / np.tan(0.5 * np.deg2rad(fov))
render = jax.jit(
partial(
splax.render,
viewmat=jnp.asarray(splax.utils.look_at(eye, centre, up)),
background=jnp.ones(3),
img_shape=(res, res),
f=(focal, focal),
gaussian_slices=slices,
)
)
images = []
for angle in np.linspace(0.0, 2 * np.pi, frames, endpoint=False):
offset = radius * distance * np.array([np.cos(angle), np.sin(angle), 0.0])
pose = orbit_pose(centre, offset, angle)
img, _ = render(*scene, gaussian_transforms=jnp.asarray(pose)[None])
images.append(np.asarray(jnp.clip(img, 0.0, 1.0) * 255, np.uint8))
iio.imwrite(out, images, duration=frames // 30, loop=0)
logger.info(f"wrote {out} with {frames} frames")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--out", type=Path, default=EXAMPLES / "compose.gif")
parser.add_argument("--res", type=int, default=512)
parser.add_argument("--fov", type=float, default=40.0, help="horizontal field of view (deg)")
parser.add_argument("--frames", type=int, default=60)
parser.add_argument("--distance", type=float, default=1.5, help="orbit radius, in splat radii")
parser.add_argument("--up", type=float, nargs=3, default=(0.0, 0.0, 1.0))
args = parser.parse_args()
main(args.out, args.res, args.fov, args.frames, args.distance, tuple(args.up))
Serve a moving object to a browser¶
Upload the hall splat once and fly a drone splat around a circle every frame, moving it without re-uploading its gaussians. See Viewer.
"""Serve a splat scene with a flying drone through ``splax.viewer``.
The hall splat is uploaded once and stays static. The drone splat is flown around a circle every
frame with ``Viewer.update_pose``, which moves it without re-uploading its gaussians.
Usage:
python examples/viewer_demo.py
python examples/viewer_demo.py --radius 1.0 --height 1.5 --port 8080
Open http://localhost:8080 in a browser, then stop with Ctrl+C.
"""
from __future__ import annotations
import argparse
import logging
import os
import time
from typing import TYPE_CHECKING
import numpy as np
os.environ.setdefault("JAX_PLATFORMS", "cpu")
os.environ.setdefault("SCIPY_ARRAY_API", "1")
from scipy.spatial.transform import Rotation
import splax
from splax.viewer import Viewer
if TYPE_CHECKING:
import viser
logger = logging.getLogger(__name__)
BASE = "https://huggingface.co/datasets/amacati/splats/resolve/main"
# The hall spans roughly 39 x 12 x 7 m with its floor near z = 0, so the drone flies over
# the middle of it.
CENTRE = np.array([1.0, 0.5, 0.0])
def main(hall: str, drone: str, port: int, radius: float, height: float, freq: float):
viewer = Viewer(port=port)
logger.info(f"loading {hall}")
viewer.add_splats("hall", *splax.io.load_ply(splax.io.fetch(f"{BASE}/{hall}.ply")))
logger.info(f"loading {drone}")
viewer.add_splats("drone", *splax.io.load_ply(splax.io.fetch(f"{BASE}/{drone}.ply")))
# Focus on the drone
focus = CENTRE + np.array([0.0, 0.0, height])
stand = (radius + 0.6) / np.sqrt(2.0)
@viewer.server.on_client_connect
def _(client: viser.ClientHandle):
client.camera.position = focus + np.array([stand, -stand, 0.25])
client.camera.look_at = focus
logger.info(f"viewer running at http://localhost:{port} -- Ctrl+C to stop")
start = time.time()
try:
while time.time() - start < args.duration:
angle = 2 * np.pi * freq * (time.time() - start)
position = focus + radius * np.array([np.cos(angle), np.sin(angle), 0.0])
# Yaw along the direction of travel, a quarter turn ahead of the orbit angle.
wxyz = Rotation.from_euler("z", angle + np.pi / 2).as_quat(scalar_first=True)
viewer.update_pose("drone", position, wxyz)
time.sleep(1 / 30)
except KeyboardInterrupt:
viewer.close()
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--hall", default="robot_hall", help="static scene in the splat repository")
parser.add_argument("--drone", default="cf21B_500", help="flying object in the repository")
parser.add_argument("--port", type=int, default=8080)
parser.add_argument("--radius", type=float, default=0.8, help="flight circle radius (m)")
parser.add_argument("--height", type=float, default=1.2, help="flight height (m)")
parser.add_argument("--freq", type=float, default=0.1, help="circle frequency (Hz)")
parser.add_argument("--duration", type=float, default=30.0, help="demo duration (s)")
args = parser.parse_args()
main(args.hall, args.drone, args.port, args.radius, args.height, args.freq)