Bundle adjustment

In this notebook, we solve a bundle adjustment problem: jointly refining camera poses and 3D point positions to minimize reprojection error.

Inputs: Initial camera poses, 3D points, and 2D observations from BAL dataset
Outputs: Refined camera poses and 3D point cloud

Features used:

• SE3Var for camera poses (batched)

• Custom Point3Var for 3D landmark positions

• Batched reprojection costs with Huber loss

• Trust region solver for large-scale optimization

• Automatic Schur-complement elimination of the landmark block

Show code cell source

Hide code cell source

import sys
from loguru import logger

logger.remove()
logger.add(sys.stdout, format="<level>{level: <8}</level> | {message}");

import jax
import jax.numpy as jnp
import jaxlie
import jaxls

Load BAL dataset

Download and parse a problem from the Dubrovnik dataset. The BAL format stores:

• Camera parameters: Rodrigues rotation (3), translation (3), focal length (1), distortion k1, k2 (2)

• 3D point positions

• 2D observations with camera and point indices

Show code cell source

Hide code cell source

import bz2
import urllib.request
from pathlib import Path


def download_bal_dataset(url: str, cache_dir: Path = Path("/tmp/bal_data")) -> Path:
    """Download and cache a BAL dataset.

    Args:
        url: URL to download the dataset from
        cache_dir: Directory to cache downloaded files

    Returns:
        Path to the downloaded (cached) file
    """
    cache_dir.mkdir(parents=True, exist_ok=True)
    filename = url.split("/")[-1]
    local_path = cache_dir / filename

    if not local_path.exists():
        logger.info(f"Downloading {filename}...")
        urllib.request.urlretrieve(url, local_path)
    return local_path


def parse_bal_file(
    path: Path,
) -> tuple[
    jax.Array,  # camera_params: (n_cameras, 9)
    jax.Array,  # points_3d: (n_points, 3)
    jax.Array,  # observations: (n_obs, 2) pixel coords
    jax.Array,  # camera_indices: (n_obs,)
    jax.Array,  # point_indices: (n_obs,)
]:
    """Parse a BAL dataset file.

    Args:
        path: Path to the BAL dataset file (bz2 compressed)

    Returns:
        Tuple of (camera_params, points_3d, observations, camera_indices, point_indices)
    """
    with bz2.open(path, "rt") as f:
        n_cameras, n_points, n_obs = map(int, f.readline().split())

        # Read observations.
        camera_indices = []
        point_indices = []
        observations = []
        for _ in range(n_obs):
            parts = f.readline().split()
            camera_indices.append(int(parts[0]))
            point_indices.append(int(parts[1]))
            observations.append([float(parts[2]), float(parts[3])])

        # Read camera parameters (9 values each)
        camera_params = []
        for _ in range(n_cameras):
            params = [float(f.readline()) for _ in range(9)]
            camera_params.append(params)

        # Read 3D points (3 values each)
        points_3d = []
        for _ in range(n_points):
            point = [float(f.readline()) for _ in range(3)]
            points_3d.append(point)

    return (
        jnp.array(camera_params),
        jnp.array(points_3d),
        jnp.array(observations),
        jnp.array(camera_indices),
        jnp.array(point_indices),
    )

# Download Dubrovnik dataset (16 cameras, 22,106 points)
bal_url = "https://grail.cs.washington.edu/projects/bal/data/dubrovnik/problem-16-22106-pre.txt.bz2"
bal_path = download_bal_dataset(bal_url)

camera_params, points_3d, observations, camera_indices, point_indices = parse_bal_file(
    bal_path
)
n_cameras, n_points, n_obs = (
    camera_params.shape[0],
    points_3d.shape[0],
    observations.shape[0],
)

print(f"camera_params:  {camera_params.shape}")
print(f"points_3d:      {points_3d.shape}")
print(f"observations:   {observations.shape}")
print(f"camera_indices: {camera_indices.shape}")
print(f"point_indices:  {point_indices.shape}")

camera_params:  (16, 9)
points_3d:      (22106, 3)
observations:   (83718, 2)
camera_indices: (83718,)
point_indices:  (83718,)

Camera model

The BAL camera model uses:

• World-to-camera: \(P = R \cdot X + t\)

• Perspective projection: \(p = -P_{xy} / P_z\)

• Radial distortion: \(r(p) = 1 + k_1 ||p||^2 + k_2 ||p||^4\)

• Final projection: \(p' = f \cdot r(p) \cdot p\)

@jax.jit
def project_point(
    point_world: jax.Array,
    T_camera_world: jaxlie.SE3,
    focal: float,
    k1: float,
    k2: float,
) -> jax.Array:
    """Project 3D point to 2D using BAL camera model.

    Args:
        point_world: 3D point in world frame (3,)
        T_camera_world: Camera pose (world to camera transform)
        focal: Focal length
        k1: First radial distortion coefficient
        k2: Second radial distortion coefficient

    Returns:
        2D projected point (2,)
    """
    # Transform point to camera frame.
    point_cam = T_camera_world @ point_world

    # Perspective projection (BAL convention: -P_xy / P_z)
    p = -point_cam[:2] / point_cam[2]

    # Radial distortion.
    r_sq = jnp.sum(p**2)
    distortion = 1.0 + k1 * r_sq + k2 * r_sq**2

    return focal * distortion * p

BAL to SE3 conversion

BAL stores camera extrinsics as Rodrigues rotation (axis-angle) + translation. We use jaxlie.SO3.exp() to convert the Rodrigues vector directly to SO3:

@jax.jit
def bal_params_to_se3(params: jax.Array) -> jaxlie.SE3:
    """Convert BAL camera parameters to SE3 pose.

    BAL stores rotation as Rodrigues vector (axis-angle).

    Args:
        params: Camera parameters [rodrigues(3), translation(3)]

    Returns:
        SE3 camera pose
    """
    return jaxlie.SE3.from_rotation_and_translation(
        rotation=jaxlie.SO3.exp(params[:3]),
        translation=params[3:6],
    )


# Convert all cameras.
initial_poses = jax.vmap(bal_params_to_se3)(camera_params[:, :6])
focal_lengths = camera_params[:, 6]
distortion_k1 = camera_params[:, 7]
distortion_k2 = camera_params[:, 8]

# Add noise to 3D points to make optimization more interesting.
key = jax.random.PRNGKey(0)
point_noise = jax.random.normal(key, points_3d.shape) * 0.5
points_3d_noisy = points_3d + point_noise

print(f"Initial poses shape: {initial_poses.wxyz_xyz.shape}")
print(
    f"Focal lengths range: [{float(focal_lengths.min()):.1f}, {float(focal_lengths.max()):.1f}]"
)
print("Added noise to 3D points: std=0.5 units")

Initial poses shape: (16, 7)
Focal lengths range: [809.5, 1918.2]
Added noise to 3D points: std=0.5 units

Variables and costs

class Point3Var(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(3)):
    """3D landmark position."""


# Create batched variables.
camera_vars = jaxls.SE3Var(id=jnp.arange(n_cameras))
point_vars = Point3Var(id=jnp.arange(n_points))

print(f"Camera variables: {n_cameras}")
print(f"Point variables: {n_points}")

Camera variables: 16
Point variables: 22106

@jaxls.Cost.factory
def reprojection_cost(
    vals: jaxls.VarValues,
    camera_var: jaxls.SE3Var,
    point_var: Point3Var,
    observed_px: jax.Array,
    focal: float,
    k1: float,
    k2: float,
) -> jax.Array:
    """Reprojection error with Huber loss for robustness."""
    pose = vals[camera_var]
    point = vals[point_var]
    projected = project_point(point, pose, focal, k1, k2)
    residual = projected - observed_px

    # IRLS-style Huber weighting for robustness to outliers.
    # For |r| <= delta: weight = 1 (quadratic region)
    # For |r| > delta: weight = delta / |r| (linear region)
    # stop_gradient prevents instabilities from differentiating through weights.
    delta = 2.0  # pixels
    abs_r = jnp.abs(residual) + 1e-8
    weight = jax.lax.stop_gradient(jnp.where(abs_r > delta, delta / abs_r, 1.0))
    return residual * jnp.sqrt(weight)

Problem construction

Create batched reprojection costs for all observations:

# Create a single batched cost for all observations.
costs = [
    reprojection_cost(
        jaxls.SE3Var(id=camera_indices),
        Point3Var(id=point_indices),
        observations,
        focal_lengths[camera_indices],
        distortion_k1[camera_indices],
        distortion_k2[camera_indices],
    )
]

print(f"Created 1 batched cost representing {n_obs} observations")

Created 1 batched cost representing 83718 observations

# Initial values (with added noise)
initial_vals = jaxls.VarValues.make(
    [
        camera_vars.with_value(initial_poses),
        point_vars.with_value(points_3d_noisy),
    ]
)

# Create the problem.
problem = jaxls.LeastSquaresProblem(costs, [camera_vars, point_vars])

# Visualize the problem structure structure.
problem.show()

# Analyze the problem.
problem = problem.analyze()

INFO     | Building optimization problem with 83718 terms and 22122 variables: 83718 costs, 0 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing group with 83718 costs, 2 variables each: reprojection_cost
INFO     | Variable elimination: eliminating Point3Var (66318 of 66414 tangent dims); reduced system is 96-dimensional

Solving

This is a large, sparse problem: thousands of points, each seen by a handful of cameras, and no cost couples two points. jaxls detects that the point block of the Hessian is block-diagonal and eliminates it automatically via the Schur complement (see Schur complement elimination). The linear solver then runs on the much smaller, better-conditioned camera-only system, and the points are recovered by back-substitution. No change to the solve call is needed; the choice is reported in the analyze() log above (“Variable elimination: eliminating Point3Var …”).

After elimination the reduced system is tiny (one block per camera), so we solve it with dense_cholesky. conjugate_gradient (the default) and cholmod use the same elimination; cholmod factors the reduced system sparse-directly, the Ceres/g2o-style “Schur + sparse-direct” combination, which is the better choice when the reduced camera system is itself large and sparse.

solution = problem.solve(initial_vals, linear_solver="dense_cholesky")

INFO     |  step #0: cost=6650404.5000 lambd=0.0005
INFO     |      - reprojection_cost(83718): 6650404.50000 (avg 39.71908)
INFO     |      accepted=True ATb_norm=1.16e+07 cost_prev=6650404.5000 cost_new=570236.6250
INFO     |  step #1: cost=570236.6250 lambd=0.0003
INFO     |      - reprojection_cost(83718): 570236.62500 (avg 3.40570)
INFO     |      accepted=True ATb_norm=4.33e+06 cost_prev=570236.6250 cost_new=410492.1562
INFO     |  step #2: cost=410492.1562 lambd=0.0001
INFO     |      - reprojection_cost(83718): 410492.15625 (avg 2.45164)
INFO     |      accepted=True ATb_norm=1.46e+06 cost_prev=410492.1562 cost_new=395957.6875
INFO     |  step #3: cost=395957.6875 lambd=0.0001
INFO     |      - reprojection_cost(83718): 395957.68750 (avg 2.36483)
INFO     |      accepted=True ATb_norm=6.26e+05 cost_prev=395957.6875 cost_new=391422.0625
INFO     |  step #4: cost=391422.0625 lambd=0.0000
INFO     |      - reprojection_cost(83718): 391422.06250 (avg 2.33774)
INFO     |      accepted=True ATb_norm=3.08e+05 cost_prev=391422.0625 cost_new=389443.2812
INFO     |  step #5: cost=389443.2812 lambd=0.0000
INFO     |      - reprojection_cost(83718): 389443.28125 (avg 2.32592)
INFO     |      accepted=True ATb_norm=1.72e+05 cost_prev=389443.2812 cost_new=388433.9375
INFO     |  step #6: cost=388433.9375 lambd=0.0000
INFO     |      - reprojection_cost(83718): 388433.93750 (avg 2.31990)
INFO     |      accepted=True ATb_norm=1.04e+05 cost_prev=388433.9375 cost_new=387857.8750
INFO     |  step #7: cost=387857.8750 lambd=0.0000
INFO     |      - reprojection_cost(83718): 387857.87500 (avg 2.31645)
INFO     |      accepted=True ATb_norm=6.78e+04 cost_prev=387857.8750 cost_new=387478.4375
INFO     |  step #8: cost=387478.4375 lambd=0.0000
INFO     |      - reprojection_cost(83718): 387478.43750 (avg 2.31419)
INFO     |      accepted=True ATb_norm=4.50e+04 cost_prev=387478.4375 cost_new=387225.2500
INFO     |  step #9: cost=387225.2500 lambd=0.0000
INFO     |      - reprojection_cost(83718): 387225.25000 (avg 2.31268)
INFO     |      accepted=True ATb_norm=3.34e+04 cost_prev=387225.2500 cost_new=387047.4375
INFO     |  step #10: cost=387047.4375 lambd=0.0000
INFO     |      - reprojection_cost(83718): 387047.43750 (avg 2.31161)
INFO     |      accepted=True ATb_norm=2.49e+04 cost_prev=387047.4375 cost_new=386909.7188
INFO     |  step #11: cost=386909.7188 lambd=0.0000
INFO     |      - reprojection_cost(83718): 386909.71875 (avg 2.31079)
INFO     |      accepted=True ATb_norm=1.98e+04 cost_prev=386909.7188 cost_new=386805.2188
INFO     |  step #12: cost=386805.2188 lambd=0.0000
INFO     |      - reprojection_cost(83718): 386805.21875 (avg 2.31017)
INFO     |      accepted=True ATb_norm=1.61e+04 cost_prev=386805.2188 cost_new=386729.2812
INFO     | Terminated @ iteration #13: cost=386729.2812 criteria=[0 0 1], term_deltas=2.0e-04,7.2e+01,9.5e-07 (solved in 0.0600 sec)

Reprojection error analysis

@jax.jit
def compute_reprojection_errors(
    vals: jaxls.VarValues,
    cam_idx: jax.Array,
    pt_idx: jax.Array,
    obs: jax.Array,
) -> jax.Array:
    """Compute reprojection errors for all observations.

    Args:
        vals: Variable values containing camera poses and 3D points
        cam_idx: Camera indices for each observation (n_obs,)
        pt_idx: Point indices for each observation (n_obs,)
        obs: 2D observations (n_obs, 2)

    Returns:
        Reprojection errors for each observation (n_obs,)
    """

    def single_error(
        c_idx: jax.Array, p_idx: jax.Array, observed: jax.Array
    ) -> jax.Array:
        pose = vals[jaxls.SE3Var(id=c_idx)]
        point = vals[Point3Var(id=p_idx)]
        projected = project_point(
            point,
            pose,
            focal_lengths[c_idx],
            distortion_k1[c_idx],
            distortion_k2[c_idx],
        )
        return jnp.linalg.norm(projected - observed)

    return jax.vmap(single_error)(cam_idx, pt_idx, obs)


errors_initial = compute_reprojection_errors(
    initial_vals, camera_indices, point_indices, observations
)
errors_final = compute_reprojection_errors(
    solution, camera_indices, point_indices, observations
)

print(
    f"Initial reprojection error: mean={float(errors_initial.mean()):.2f}px, median={float(jnp.median(errors_initial)):.2f}px"
)
print(
    f"Final reprojection error:   mean={float(errors_final.mean()):.2f}px, median={float(jnp.median(errors_final)):.2f}px"
)

Initial reprojection error: mean=31.24px, median=26.71px
Final reprojection error:   mean=2.28px, median=1.65px

Visualization

Show code cell source

Hide code cell source

import contextlib
import io
import numpy as np
import viser

# Compute per-point mean reprojection error for filtering.
point_errors = np.zeros(n_points)
point_counts = np.zeros(n_points)
for i, (pt_idx, err) in enumerate(zip(np.array(point_indices), np.array(errors_final))):
    point_errors[pt_idx] += err
    point_counts[pt_idx] += 1
point_mean_errors = point_errors / np.maximum(point_counts, 1)

# Filter to points with low reprojection error.
error_threshold = 1.0  # pixels
good_point_mask = point_mean_errors < error_threshold
good_point_indices = np.where(good_point_mask)[0]

points_init = np.array(initial_vals[Point3Var])
points_opt = np.array(solution[Point3Var])

# Subsample filtered points for visualization.
subsample = max(1, len(good_point_indices) // 30_000)
vis_indices = good_point_indices[::subsample]

print(f"Showing {len(vis_indices)} of {n_points} points (error < {error_threshold}px)")

# Create Viser server (suppress output).
with (
    contextlib.redirect_stdout(io.StringIO()),
    contextlib.redirect_stderr(io.StringIO()),
):
    server = viser.ViserServer(verbose=False)
server.scene.set_up_direction("+z")

# Offset for side-by-side views. Center between the two is at origin.
offset = 20.0

# Center the point clouds at origin.
points_center = points_opt[vis_indices].mean(axis=0)
label_height = points_opt[vis_indices][:, 2].max() - points_center[2] + 5.0

# Set initial camera position for a good view of both point clouds.
server.initial_camera.position = (0.0, -80.0, 30.0)
server.initial_camera.look_at = (0.0, 0.0, 0.0)

# Add labels (optimized on left, initial on right).
server.scene.add_label(
    "/optimized_label",
    text="Optimized",
    position=(-offset, 0.0, label_height),
    anchor="bottom-center",
)
server.scene.add_label(
    "/initial_label",
    text="Initial (with noise)",
    position=(offset, 0.0, label_height),
    anchor="bottom-center",
)

# Add optimized point cloud (left side).
server.scene.add_point_cloud(
    "/optimized/points",
    points=points_opt[vis_indices] - points_center + np.array([-offset, 0, 0]),
    colors=np.full((len(vis_indices), 3), [70, 130, 180], dtype=np.uint8),  # Steel blue
    point_size=0.08,
)

# Add initial point cloud (right side).
server.scene.add_point_cloud(
    "/initial/points",
    points=points_init[vis_indices] - points_center + np.array([offset, 0, 0]),
    colors=np.full((len(vis_indices), 3), 150, dtype=np.uint8),  # Gray
    point_size=0.08,
)

# Add camera frustums.
# BAL stores T_camera_world (world-to-camera transform).
# BAL cameras look down -Z, but Viser expects +Z, so we flip by 180° around X.
poses_init = initial_vals[jaxls.SE3Var]
poses_opt = solution[jaxls.SE3Var]

# 180° rotation around X-axis to flip from -Z to +Z forward direction.
flip_rotation = jaxlie.SO3.from_x_radians(np.pi)

for i in range(n_cameras):
    # Optimized camera (left side).
    pose_opt = jaxlie.SE3(wxyz_xyz=poses_opt.wxyz_xyz[i])
    T_world_camera_opt = pose_opt.inverse()
    cam_pos_opt = (
        np.array(T_world_camera_opt.translation())
        - points_center
        + np.array([-offset, 0, 0])
    )
    cam_rot_opt = T_world_camera_opt.rotation() @ flip_rotation

    server.scene.add_camera_frustum(
        f"/optimized/camera_{i}",
        fov=np.pi / 3,
        aspect=1.5,
        scale=0.5,
        wxyz=np.array(cam_rot_opt.wxyz),
        position=cam_pos_opt,
        color=(34, 139, 34),  # Forest green
    )

    # Initial camera (right side).
    pose_init = jaxlie.SE3(wxyz_xyz=poses_init.wxyz_xyz[i])
    T_world_camera_init = pose_init.inverse()
    cam_pos_init = (
        np.array(T_world_camera_init.translation())
        - points_center
        + np.array([offset, 0, 0])
    )
    cam_rot_init = T_world_camera_init.rotation() @ flip_rotation

    server.scene.add_camera_frustum(
        f"/initial/camera_{i}",
        fov=np.pi / 3,
        aspect=1.5,
        scale=0.5,
        wxyz=np.array(cam_rot_init.wxyz),
        position=cam_pos_init,
        color=(255, 99, 71),  # Tomato
    )

# Display inline in the notebook.
server.scene.show(height=500)

╭────── viser (listening *:8080) ───────╮
│             ╷                         │
│   HTTP      │ http://localhost:8080   │
│   Websocket │ ws://localhost:8080     │
│             ╵                         │
╰───────────────────────────────────────╯

Showing 5298 of 22106 points (error < 1.0px)

Show code cell source

Hide code cell source

import plotly.graph_objects as go
from IPython.display import HTML

# Reprojection error histogram.
fig_hist = go.Figure()

# Clip for visualization (outliers compress the histogram).
max_error = 50.0
errors_initial_clipped = jnp.clip(errors_initial, 0, max_error)
errors_final_clipped = jnp.clip(errors_final, 0, max_error)

fig_hist.add_trace(
    go.Histogram(
        x=errors_initial_clipped,
        name="Initial",
        opacity=0.7,
        nbinsx=50,
        marker_color="lightcoral",
    )
)

fig_hist.add_trace(
    go.Histogram(
        x=errors_final_clipped,
        name="Optimized",
        opacity=0.7,
        nbinsx=50,
        marker_color="steelblue",
    )
)

fig_hist.update_layout(
    title="Reprojection Error Distribution",
    xaxis=dict(title="Reprojection Error (pixels)", range=[0, max_error]),
    yaxis_title="Count",
    barmode="overlay",
    height=350,
    margin=dict(t=40, b=40, l=60, r=40),
    legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99),
)

HTML(fig_hist.to_html(full_html=False, include_plotlyjs="cdn"))

Bundle adjustment jointly refines camera poses and 3D point positions to minimize reprojection error. The Huber loss provides robustness to outliers, which is important for real-world data.

For solver configuration, see jaxls.TrustRegionConfig. For Lie group variables, see jaxls.SE3Var.