geometry

Multi-view geometry primitives in JAX, float64, jit- and grad-friendly.

A world point X maps to camera coordinates as R @ X + tvec, where R is R(rvec) for an ordinary camera and the improper diag(-1, 1, 1) @ R(rvec) for a mirrored one (ADR-0053, pose_rmat_one). Intrinsics intr are [fx, fy, cx, cy] or [f, cx, cy] (intr_to_kmat); distortion dists follow OpenCV's order [k1, k2, p1, p2, k3, k4, k5, k6, s1, s2, s3, s4]; pmat is K @ [R | t]. A function takes its operand first, then rvecs, tvecs, intrs, dists. The batched functions are jitted vmaps of the *_one variants, which compose with jax.jacfwd for bundle adjustment; the arrays are tiny, so JAX runs on the CPU.

A pixel coordinate is in the raw frame of the video the view reads: full resolution, the whole frame, shared by every view of that video (ADR-0055). A point -- a label, a prediction, cx and cy -- sits at a pixel center: pixel i spans [i - 0.5, i + 0.5], so a W-wide frame spans [-0.5, W - 0.5] and centers on (W - 1) / 2. A window bound -- a crop, a detector's box -- is a slice bound, image[y0:y1, x0:x1], on the pixel edges, so the same frame spans [0, W].

backproject_ray_one

backproject_ray_one(
    pixel: Float[Array, 2],
    rvec: Float[Array, 3],
    tvec: Float[Array, 3],
    intr: Float[Array, P],
    dist: Float[Array, K],
    mirrored: Bool[Array, ""],
) -> tuple[Float[Array, 3], Float[Array, 3]]

Back-project an image pixel to its viewing ray in world coordinates.

Inverts project_full_one: the ray origin + s * direction (s >= 0) is the set of world points that project to pixel. Returns (origin, direction), the camera center -R.T @ tvec and the unnormalized world direction; pair it with closest_point_on_ray to move a 3D point onto the ray.

Source code in src/deeperfly/geometry.py

def backproject_ray_one(
    pixel: Float[Array, "2"],
    rvec: Float[Array, "3"],
    tvec: Float[Array, "3"],
    intr: Float[Array, "P"],
    dist: Float[Array, "K"],
    mirrored: Bool[Array, ""],
) -> tuple[Float[Array, "3"], Float[Array, "3"]]:
    """Back-project an image pixel to its viewing ray in world coordinates.

    Inverts [project_full_one][]: the ray `origin + s * direction` (`s >= 0`) is the
    set of world points that project to `pixel`. Returns `(origin, direction)`, the
    camera center `-R.T @ tvec` and the unnormalized world direction; pair it with
    [closest_point_on_ray][] to move a 3D point onto the ray.
    """
    fx, fy = intr[0], intr[-3]
    cx, cy = intr[-2], intr[-1]
    xy_dist = jnp.stack([(pixel[0] - cx) / fx, (pixel[1] - cy) / fy])
    xy = undistort_one(xy_dist, dist)
    dir_cam = jnp.stack([xy[0], xy[1], jnp.ones(())])
    rmat = pose_rmat_one(rvec, mirrored)
    direction = rmat.T @ dir_cam
    origin = -rmat.T @ tvec
    return origin, direction

closest_point_on_ray

closest_point_on_ray(
    origin: Float[ArrayLike, 3],
    direction: Float[ArrayLike, 3],
    target: Float[ArrayLike, 3],
) -> Float[Array, 3]

The point on a ray nearest target, its orthogonal projection onto the line.

The 3D-correction drag moves a point to where its dragged pixel's ray passes closest to its old position, so it reprojects onto the pixel and moves the least.

Source code in src/deeperfly/geometry.py

def closest_point_on_ray(
    origin: Float[ArrayLike, "3"],
    direction: Float[ArrayLike, "3"],
    target: Float[ArrayLike, "3"],
) -> Float[Array, "3"]:
    """The point on a ray nearest `target`, its orthogonal projection onto the line.

    The 3D-correction drag moves a point to where its dragged pixel's ray passes
    closest to its old position, so it reprojects onto the pixel and moves the least.
    """
    origin, direction, target = map(jnp.asarray, (origin, direction, target))
    s = jnp.dot(target - origin, direction) / jnp.dot(direction, direction)
    return origin + s * direction

distort

distort(
    pts2d: Float[Array, "V *pts 2"],
    dists: Float[Array, "V K"],
) -> Float[Array, "V *pts 2"]

Apply OpenCV-style radial, tangential and thin-prism distortion.

For normalized (x, y) and r^2 = x^2 + y^2, with the coefficients [k1, k2, p1, p2, k3, k4, k5, k6, s1, s2, s3, s4] (zero beyond K, as in cv2.projectPoints):

x_d = x * (1 + k1 r^2 + k2 r^4 + k3 r^6)
        / (1 + k4 r^2 + k5 r^4 + k6 r^6)
      + 2 p1 x y + p2 (r^2 + 2 x^2) + s1 r^2 + s2 r^4
y_d = y * (...) / (...)
      + p1 (r^2 + 2 y^2) + 2 p2 x y + s3 r^2 + s4 r^4

Source code in src/deeperfly/geometry.py

@jax.jit
def distort(
    pts2d: Float[Array, "V *pts 2"],
    dists: Float[Array, "V K"],
) -> Float[Array, "V *pts 2"]:
    """Apply OpenCV-style radial, tangential and thin-prism distortion.

    For normalized `(x, y)` and `r^2 = x^2 + y^2`, with the coefficients
    `[k1, k2, p1, p2, k3, k4, k5, k6, s1, s2, s3, s4]` (zero beyond `K`, as in
    `cv2.projectPoints`):

    ```text
    x_d = x * (1 + k1 r^2 + k2 r^4 + k3 r^6)
            / (1 + k4 r^2 + k5 r^4 + k6 r^6)
          + 2 p1 x y + p2 (r^2 + 2 x^2) + s1 r^2 + s2 r^4
    y_d = y * (...) / (...)
          + p1 (r^2 + 2 y^2) + 2 p2 x y + s3 r^2 + s4 r^4
    ```
    """
    if dists.shape[-1] == 0:
        return pts2d
    v = pts2d.shape[0]
    flat = pts2d.reshape(v, -1, 2)
    out = jax.vmap(jax.vmap(distort_one, in_axes=(0, None)), in_axes=(0, 0))(
        flat, dists
    )
    return out.reshape(pts2d.shape)

distort_one

distort_one(
    xy: Float[Array, 2], dist: Float[Array, K]
) -> Float[Array, 2]

distort for one normalized 2D point; dist holds 0 to 12 coefficients.

Source code in src/deeperfly/geometry.py

def distort_one(
    xy: Float[Array, "2"],
    dist: Float[Array, "K"],
) -> Float[Array, "2"]:
    """[distort][] for one normalized 2D point; `dist` holds 0 to 12 coefficients."""
    n = dist.shape[-1]
    if n == 0:
        return xy
    x, y = xy[0], xy[1]
    x2, y2 = x * x, y * y
    r2 = x2 + y2

    num = 1.0 + dist[0] * r2
    r4 = r6 = jnp.zeros(())
    if n >= 2:
        r4 = r2 * r2
        num = num + dist[1] * r4
    if n >= 5:
        r6 = r4 * r2
        num = num + dist[4] * r6
    den = 1.0
    if n >= 6:
        den = den + dist[5] * r2
    if n >= 7:
        den = den + dist[6] * r4
    if n >= 8:
        den = den + dist[7] * r6
    mult = num / den

    add_x = jnp.zeros(())
    add_y = jnp.zeros(())
    xy_prod = x * y
    if n >= 3:
        add_x = 2 * dist[2] * xy_prod
        add_y = dist[2] * (r2 + 2 * y2)
    if n >= 4:
        add_x = add_x + dist[3] * (r2 + 2 * x2)
        add_y = add_y + 2 * dist[3] * xy_prod
    if n >= 9:
        add_x = add_x + dist[8] * r2
    if n >= 10:
        add_x = add_x + dist[9] * r4
    if n >= 11:
        add_y = add_y + dist[10] * r2
    if n >= 12:
        add_y = add_y + dist[11] * r4

    return jnp.stack([x * mult + add_x, y * mult + add_y])

intr_to_kmat

intr_to_kmat(
    intr: Float[Array, "*batch P"],
) -> Float[Array, "*batch 3 3"]

3x3 intrinsic matrices from packed [fx, fy, cx, cy] or [f, cx, cy].

Source code in src/deeperfly/geometry.py

@jax.jit
def intr_to_kmat(
    intr: Float[Array, "*batch P"],
) -> Float[Array, "*batch 3 3"]:
    """3x3 intrinsic matrices from packed `[fx, fy, cx, cy]` or `[f, cx, cy]`."""
    kmat = jnp.zeros((*intr.shape[:-1], 3, 3))
    kmat = kmat.at[..., 0, 0].set(intr[..., 0])
    kmat = kmat.at[..., 1, 1].set(intr[..., -3])
    kmat = kmat.at[..., :2, 2].set(intr[..., -2:])
    kmat = kmat.at[..., 2, 2].set(1.0)
    return kmat

pmats

pmats(
    rvecs: Float[Array, "*batch 3"],
    tvecs: Float[Array, "*batch 3"],
    intrs: Float[Array, "*batch P"],
    mirrored: Bool[Array, "*batch"],
) -> Float[Array, "*batch 3 4"]

Projection matrices K @ [R | t], the pinhole model without distortion.

The one place the package builds them, so no caller can forget the mirror.

Source code in src/deeperfly/geometry.py

@jax.jit
def pmats(
    rvecs: Float[Array, "*batch 3"],
    tvecs: Float[Array, "*batch 3"],
    intrs: Float[Array, "*batch P"],
    mirrored: Bool[Array, "*batch"],
) -> Float[Array, "*batch 3 4"]:
    """Projection matrices `K @ [R | t]`, the pinhole model without distortion.

    The one place the package builds them, so no caller can forget the mirror.
    """
    rtmat = jnp.concatenate((pose_rmat(rvecs, mirrored), tvecs[..., None]), axis=-1)
    return intr_to_kmat(intrs) @ rtmat

pose_rmat

pose_rmat(
    rvecs: Float[Array, "*batch 3"],
    mirrored: Bool[Array, "*batch"],
) -> Float[Array, "*batch 3 3"]

Effective rotations of poses that may be improper; see pose_rmat_one.

Source code in src/deeperfly/geometry.py

@jax.jit
def pose_rmat(
    rvecs: Float[Array, "*batch 3"],
    mirrored: Bool[Array, "*batch"],
) -> Float[Array, "*batch 3 3"]:
    """Effective rotations of poses that may be improper; see [pose_rmat_one][]."""
    flat = jax.vmap(pose_rmat_one)(rvecs.reshape(-1, 3), mirrored.reshape(-1))
    return flat.reshape(*rvecs.shape[:-1], 3, 3)

pose_rmat_one

pose_rmat_one(
    rvec: Float[ArrayLike, 3], mirrored: Bool[ArrayLike, ""]
) -> Float[Array, "3 3"]

The rotation of a pose that may be improper: S @ R(rvec) when mirrored.

S = diag(-1, 1, 1). A view through a prism sees the world reflected, which no axis-angle vector carries, so it is a proper rvec plus a flag (ADR-0053); S @ R is still orthogonal, so back-projection needs no special case.

Source code in src/deeperfly/geometry.py

def pose_rmat_one(
    rvec: Float[ArrayLike, "3"], mirrored: Bool[ArrayLike, ""]
) -> Float[Array, "3 3"]:
    """The rotation of a pose that may be improper: `S @ R(rvec)` when `mirrored`.

    `S = diag(-1, 1, 1)`. A view through a prism sees the world reflected, which no
    axis-angle vector carries, so it is a proper `rvec` plus a flag (ADR-0053); `S @ R`
    is still orthogonal, so back-projection needs no special case.
    """
    rmat = rvec_to_rmat_one(jnp.asarray(rvec))
    return rmat.at[0].multiply(jnp.where(mirrored, -1.0, 1.0))

project_full

project_full(
    pts3d: Float[Array, "*pts 3"],
    rvecs: Float[Array, "V 3"],
    tvecs: Float[Array, "V 3"],
    intrs: Float[Array, "V P"] | Float[Array, P],
    dists: Float[Array, "V K"] | Float[Array, K],
    mirrored: Bool[Array, V],
) -> Float[Array, "V *pts 2"]

Project 3D world points to 2D image points through full camera models.

X_cam = R X + tvec, divided by depth, distorted by distort and mapped through the intrinsics. intrs and dists may be one row shared by every view; mirrored is per view, never shared.

Source code in src/deeperfly/geometry.py

@jax.jit
def project_full(
    pts3d: Float[Array, "*pts 3"],
    rvecs: Float[Array, "V 3"],
    tvecs: Float[Array, "V 3"],
    intrs: Float[Array, "V P"] | Float[Array, "P"],
    dists: Float[Array, "V K"] | Float[Array, "K"],
    mirrored: Bool[Array, "V"],
) -> Float[Array, "V *pts 2"]:
    """Project 3D world points to 2D image points through full camera models.

    `X_cam = R X + tvec`, divided by depth, distorted by [distort][] and mapped through
    the intrinsics. `intrs` and `dists` may be one row shared by every view;
    `mirrored` is per view, never shared.
    """
    v = rvecs.shape[0]
    pts_shape = pts3d.shape[:-1]
    pts_flat = pts3d.reshape(-1, 3)
    if intrs.ndim == 1:
        intrs = jnp.broadcast_to(intrs, (v, intrs.shape[0]))
    if dists.ndim == 1:
        dists = jnp.broadcast_to(dists, (v, dists.shape[0]))
    project_v = jax.vmap(
        jax.vmap(project_full_one, in_axes=(0, None, None, None, None, None)),
        in_axes=(None, 0, 0, 0, 0, 0),
    )
    out = project_v(pts_flat, rvecs, tvecs, intrs, dists, mirrored)
    return out.reshape(v, *pts_shape, 2)

project_full_one

project_full_one(
    pt3d: Float[Array, 3],
    rvec: Float[Array, 3],
    tvec: Float[Array, 3],
    intr: Float[Array, P],
    dist: Float[Array, K],
    mirrored: Bool[Array, ""],
) -> Float[Array, 2]

Project one 3D point through one camera, composable with vmap and jacfwd.

mirrored comes last because it is never differentiated, and is required: a dropped mirror is silent on an unmirrored rig and wrong on a prism one.

Source code in src/deeperfly/geometry.py

def project_full_one(
    pt3d: Float[Array, "3"],
    rvec: Float[Array, "3"],
    tvec: Float[Array, "3"],
    intr: Float[Array, "P"],
    dist: Float[Array, "K"],
    mirrored: Bool[Array, ""],
) -> Float[Array, "2"]:
    """Project one 3D point through one camera, composable with `vmap` and `jacfwd`.

    `mirrored` comes last because it is never differentiated, and is required: a
    dropped mirror is silent on an unmirrored rig and wrong on a prism one.
    """
    rmat = pose_rmat_one(rvec, mirrored)
    p_cam = rmat @ pt3d + tvec
    xy = p_cam[:2] / p_cam[2]
    xy = distort_one(xy, dist)
    fx, fy = intr[0], intr[-3]
    cx, cy = intr[-2], intr[-1]
    return jnp.stack([fx * xy[0] + cx, fy * xy[1] + cy])

project_pmat

project_pmat(
    pts3d: Float[Array, "*pts 3"],
    pmats: Float[Array, "*views 3 4"],
) -> Float[Array, "*views *pts 2"]

Project 3D points to 2D through 3x4 projection matrices, (*views, *pts, 2).

Source code in src/deeperfly/geometry.py

@jax.jit
def project_pmat(
    pts3d: Float[Array, "*pts 3"],
    pmats: Float[Array, "*views 3 4"],
) -> Float[Array, "*views *pts 2"]:
    """Project 3D points to 2D through 3x4 projection matrices, `(*views, *pts, 2)`."""
    output_shape = (*pmats.shape[:-2], *pts3d.shape[:-1], 2)
    pmats_flat = pmats.reshape(-1, 3, 4)
    pts_flat = pts3d.reshape(-1, 3).T
    pts2dh = pmats_flat[:, :, :3] @ pts_flat + pmats_flat[:, :, 3:]
    pts2d = (pts2dh[:, :2] / pts2dh[:, 2:]).transpose(0, 2, 1)
    return pts2d.reshape(output_shape)

rmat_to_rvec

rmat_to_rvec(
    rmat: Float[Array, "*batch 3 3"],
) -> Float[Array, "*batch 3"]

Rotation matrices to axis-angle vectors, as OpenCV's Rodrigues does.

The axis comes from R - R^T, which vanishes at theta = pi; near it the axis is read off (R + I) / 2, with the signs from the off-diagonal entries.

Source code in src/deeperfly/geometry.py

@jax.jit
def rmat_to_rvec(
    rmat: Float[Array, "*batch 3 3"],
) -> Float[Array, "*batch 3"]:
    """Rotation matrices to axis-angle vectors, as OpenCV's `Rodrigues` does.

    The axis comes from `R - R^T`, which vanishes at `theta = pi`; near it the axis is
    read off `(R + I) / 2`, with the signs from the off-diagonal entries.
    """
    flat = rmat.reshape(-1, 3, 3)
    out = jax.vmap(rmat_to_rvec_one)(flat)
    return out.reshape(*rmat.shape[:-2], 3)

rmat_to_rvec_one

rmat_to_rvec_one(
    rmat: Float[Array, "3 3"],
) -> Float[Array, 3]

Rotation matrix to axis-angle vector, for one rotation.

Source code in src/deeperfly/geometry.py

def rmat_to_rvec_one(rmat: Float[Array, "3 3"]) -> Float[Array, "3"]:
    """Rotation matrix to axis-angle vector, for one rotation."""
    r00, r01, r02 = rmat[0, 0], rmat[0, 1], rmat[0, 2]
    r10, r11, r12 = rmat[1, 0], rmat[1, 1], rmat[1, 2]
    r20, r21, r22 = rmat[2, 0], rmat[2, 1], rmat[2, 2]

    rho = jnp.array([r21 - r12, r02 - r20, r10 - r01])
    sin_theta = jnp.linalg.norm(rho) / 2
    cos_theta = jnp.clip((r00 + r11 + r22 - 1) / 2, -1.0, 1.0)
    theta = jnp.arccos(cos_theta)
    near_pi = sin_theta < _NEAR_PI_SIN_THRESH

    rvec = rho * (theta / jnp.where(near_pi, 1.0, 2.0 * sin_theta))

    ax = jnp.sqrt(jnp.maximum((r00 + 1.0) / 2, 0.0))
    ay = jnp.sqrt(jnp.maximum((r11 + 1.0) / 2, 0.0)) * jnp.where(r01 < 0, -1.0, 1.0)
    az = jnp.sqrt(jnp.maximum((r22 + 1.0) / 2, 0.0)) * jnp.where(r02 < 0, -1.0, 1.0)
    flip = (
        (jnp.abs(ax) < jnp.abs(ay))
        & (jnp.abs(ax) < jnp.abs(az))
        & ((r12 > 0) != (ay * az > 0))
    )
    az = jnp.where(flip, -az, az)
    axis = jnp.array([ax, ay, az])
    axis_norm = jnp.linalg.norm(axis)
    rvec_pi = axis * (theta / jnp.where(axis_norm == 0, 1.0, axis_norm))
    rvec_degenerate = jnp.where(cos_theta > 0, 0.0, rvec_pi)

    return jnp.where(near_pi, rvec_degenerate, rvec)

rvec_to_rmat

rvec_to_rmat(
    rvec: Float[Array, "*batch 3"],
) -> Float[Array, "*batch 3 3"]

Axis-angle rotation vectors to rotation matrices (Rodrigues).

R = I + a W + b W^2 with W = skew(rvec), a = sin(t) / t and b = (1 - cos(t)) / t^2, taken from their Taylor series for small t, so the result stays orthogonal to machine precision and there is no 0/0 at t = 0.

Source code in src/deeperfly/geometry.py

@jax.jit
def rvec_to_rmat(
    rvec: Float[Array, "*batch 3"],
) -> Float[Array, "*batch 3 3"]:
    """Axis-angle rotation vectors to rotation matrices (Rodrigues).

    `R = I + a W + b W^2` with `W = skew(rvec)`, `a = sin(t) / t` and
    `b = (1 - cos(t)) / t^2`, taken from their Taylor series for small `t`, so the
    result stays orthogonal to machine precision and there is no `0/0` at `t = 0`.
    """
    flat = rvec.reshape(-1, 3)
    out = jax.vmap(rvec_to_rmat_one)(flat)
    return out.reshape(*rvec.shape[:-1], 3, 3)

rvec_to_rmat_one

rvec_to_rmat_one(
    rvec: Float[Array, 3],
) -> Float[Array, "3 3"]

Rodrigues' rotation for one rotation vector (axis times angle in radians).

Source code in src/deeperfly/geometry.py

def rvec_to_rmat_one(rvec: Float[Array, "3"]) -> Float[Array, "3 3"]:
    """Rodrigues' rotation for one rotation vector (axis times angle in radians)."""
    theta_sq = jnp.dot(rvec, rvec)
    small = theta_sq < _SMALL_THETA_SQ
    # Guard the sqrt too: its derivative is infinite at `theta_sq = 0`, and reverse-mode
    # AD multiplies that by the unused branch's zero cotangent into a NaN the `where`
    # keeps, which the BA Jacobian hits for any camera at `rvec = 0`.
    theta = jnp.sqrt(jnp.where(small, 1.0, theta_sq))
    sinc_theta = jnp.where(
        small,
        1 - theta_sq / 6 + theta_sq**2 / 120,
        jnp.sin(theta) / theta,
    )
    cosc_theta = jnp.where(
        small,
        0.5 - theta_sq / 24 + theta_sq**2 / 720,
        (1 - jnp.cos(theta)) / jnp.where(small, 1.0, theta_sq),
    )
    rx, ry, rz = rvec[0], rvec[1], rvec[2]
    skew = jnp.array(
        [
            [0.0, -rz, ry],
            [rz, 0.0, -rx],
            [-ry, rx, 0.0],
        ]
    )
    outer = jnp.outer(rvec, rvec)
    return (
        (1 - cosc_theta * theta_sq) * jnp.eye(3)
        + sinc_theta * skew
        + cosc_theta * outer
    )

triangulate_dlt

triangulate_dlt(
    pts2d: Float[Array, "V *pts 2"],
    pmats: Float[Array, "V 3 4"],
    weights: Float[Array, V * pts] | None = None,
    method: str = "invit",
) -> Float[Array, "*pts 3"]

Triangulate 3D points by direct linear transformation (DLT).

The homogeneous point is the null vector of the rows [x * p3 - p1; y * p3 - p2] of every view, the smallest eigenvector of A^T A. A point with fewer than two valid views, or degenerate views, is NaN. weights scale each view's rows by sqrt(weight), as bundle adjustment weighs confidences; a zero, negative or non-finite weight drops the view. method is "invit", inverse iteration with an explicit adjugate, or "eigh", slower but robust on degenerate systems; they agree to ~1e-12 where the triangulation is well posed (ADR-0043).

Source code in src/deeperfly/geometry.py

@partial(jax.jit, static_argnames="method")
def triangulate_dlt(
    pts2d: Float[Array, "V *pts 2"],
    pmats: Float[Array, "V 3 4"],
    weights: Float[Array, "V *pts"] | None = None,
    method: str = "invit",
) -> Float[Array, "*pts 3"]:
    """Triangulate 3D points by direct linear transformation (DLT).

    The homogeneous point is the null vector of the rows `[x * p3 - p1; y * p3 - p2]`
    of every view, the smallest eigenvector of `A^T A`. A point with fewer than two
    valid views, or degenerate views, is NaN. `weights` scale each view's rows by
    `sqrt(weight)`, as bundle adjustment weighs confidences; a zero, negative or
    non-finite weight drops the view. `method` is `"invit"`, inverse iteration with an
    explicit adjugate, or `"eigh"`, slower but robust on degenerate systems; they agree
    to ~1e-12 where the triangulation is well posed (ADR-0043).
    """
    ata, enough = _dlt_normal_eqs(pts2d, pmats, weights)
    if method == "eigh":
        eigvals, eigvecs = jnp.linalg.eigh(ata)
        vec = eigvecs[..., :, 0]
        # A near-zero second-smallest eigenvalue is the degenerate rig the adjugate
        # flags for "invit"; NaN it so both methods agree.
        tr = eigvals.sum(axis=-1)
        degenerate = eigvals[..., 1] < _DLT_DEGENERATE_TOL * jnp.where(tr > 0, tr, 1.0)
    else:
        vec, _, degenerate = _smallest_eigvec_invit(ata)
    return _dehomogenize(vec, enough & ~degenerate)

triangulate_dlt_flagged

triangulate_dlt_flagged(
    pts2d: Float[Array, "V *pts 2"],
    pmats: Float[Array, "V 3 4"],
    weights: Float[Array, V * pts] | None = None,
) -> tuple[Float[Array, "*pts 3"], Bool[Array, "*pts"]]

Fast DLT triangulation, plus a mask of the points to re-solve with eigh.

triangulate_dlt with method="invit", and the illcond mask of solved points too ill-conditioned (condition number above ~1e11) to trust; NaN points are not in it, since no solver rescues them.

Source code in src/deeperfly/geometry.py

@jax.jit
def triangulate_dlt_flagged(
    pts2d: Float[Array, "V *pts 2"],
    pmats: Float[Array, "V 3 4"],
    weights: Float[Array, "V *pts"] | None = None,
) -> tuple[Float[Array, "*pts 3"], Bool[Array, "*pts"]]:
    """Fast DLT triangulation, plus a mask of the points to re-solve with `eigh`.

    [triangulate_dlt][deeperfly.geometry.triangulate_dlt] with `method="invit"`, and
    the `illcond` mask of solved points too ill-conditioned (condition number above
    ~1e11) to trust; NaN points are not in it, since no solver rescues them.
    """
    ata, enough = _dlt_normal_eqs(pts2d, pmats, weights)
    vec, illcond, degenerate = _smallest_eigvec_invit(ata)
    solvable = enough & ~degenerate
    return _dehomogenize(vec, solvable), illcond & solvable

undistort_one

undistort_one(
    xy_dist: Float[Array, 2], dist: Float[Array, K]
) -> Float[Array, 2]

Invert distort_one: a distorted normalized coordinate to an undistorted one.

OpenCV's fixed-point iteration (cv2.undistortPoints), starting from xy_dist, which has the intrinsics already removed. It turns a clicked pixel into a back-projection direction (backproject_ray_one); where the iteration diverges, outside the lens model's invertible region, the result falls back to xy_dist.

Source code in src/deeperfly/geometry.py

def undistort_one(
    xy_dist: Float[Array, "2"],
    dist: Float[Array, "K"],
) -> Float[Array, "2"]:
    """Invert [distort_one][]: a distorted normalized coordinate to an undistorted one.

    OpenCV's fixed-point iteration (`cv2.undistortPoints`), starting from `xy_dist`,
    which has the intrinsics already removed. It turns a clicked pixel into a
    back-projection direction ([backproject_ray_one][]); where the iteration diverges,
    outside the lens model's invertible region, the result falls back to `xy_dist`.
    """
    n = dist.shape[-1]
    if n == 0:
        return xy_dist
    xd, yd = xy_dist[0], xy_dist[1]
    x, y = xd, yd
    for _ in range(_UNDISTORT_ITERS):
        x2, y2 = x * x, y * y
        r2 = x2 + y2
        r4 = r2 * r2
        r6 = r4 * r2

        num = 1.0 + dist[0] * r2
        if n >= 2:
            num = num + dist[1] * r4
        if n >= 5:
            num = num + dist[4] * r6
        den = 1.0
        if n >= 6:
            den = den + dist[5] * r2
        if n >= 7:
            den = den + dist[6] * r4
        if n >= 8:
            den = den + dist[7] * r6
        icdist = den / num

        xy_prod = x * y
        add_x = jnp.zeros(())
        add_y = jnp.zeros(())
        if n >= 3:
            add_x = 2 * dist[2] * xy_prod
            add_y = dist[2] * (r2 + 2 * y2)
        if n >= 4:
            add_x = add_x + dist[3] * (r2 + 2 * x2)
            add_y = add_y + 2 * dist[3] * xy_prod
        if n >= 9:
            add_x = add_x + dist[8] * r2
        if n >= 10:
            add_x = add_x + dist[9] * r4
        if n >= 11:
            add_y = add_y + dist[10] * r2
        if n >= 12:
            add_y = add_y + dist[11] * r4

        x = (xd - add_x) * icdist
        y = (yd - add_y) * icdist

    # Outside the lens model's invertible region (a point near or behind the camera)
    # the iteration diverges: keep it only where it inverts better than no correction.
    result = jnp.stack([x, y])
    resid_iter = distort_one(result, dist) - xy_dist
    resid_none = distort_one(xy_dist, dist) - xy_dist
    good = jnp.all(jnp.isfinite(result)) & (
        jnp.dot(resid_iter, resid_iter) <= jnp.dot(resid_none, resid_none)
    )
    return jnp.where(good, result, xy_dist)