Upload folder using huggingface_hub

mapvggt/__init__.py

@@ -0,0 +1,3 @@
+from mapvggt.model import MapVGGT, lift_to_world
+
+__all__ = ["MapVGGT", "lift_to_world"]

mapvggt/crosscolor.py

@@ -0,0 +1,57 @@
+"""CrossColor: multi-view color fusion for per-pixel Gaussians.
+
+Each Gaussian is born from one input pixel (view i) and currently carries that source
+pixel's RGB (frozen). That single color is wrong for many target views (exposure/AWB
+differences across the rig, grazing source angle) and causes seams where Gaussians from
+different views overlap. CrossColor instead fuses the colors from ALL input views that
+observe the Gaussian: reproject the Gaussian's world position into every input view,
+sample, keep views where it is in-frustum AND depth-consistent (not occluded), and fuse.
+
+Step-1 here is a PARAMETER-FREE fusion (visibility/occlusion-weighted, no learning) to
+validate the hypothesis on an existing checkpoint before training a learned head.
+"""
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+
+def reproject_sample(means_w, img_j, K_j, c2w_j, z_j, occ_tol=0.10):
+    """Project world points means_w [N,3] into view j; sample img_j color + visibility.
+    Returns color [N,3], visible [N] bool. Visible = in-frustum, in front, and the
+    point's camera-z agrees with view j's predicted depth at that pixel (occlusion test)."""
+    N = means_w.shape[0]
+    H, W = img_j.shape[-2:]
+    R = c2w_j[:3, :3]; t = c2w_j[:3, 3]
+    p_cam = (means_w - t) @ R                                  # world->cam (R^T (x-t))
+    zc = p_cam[:, 2]
+    uv = p_cam[:, :2] / zc.clamp(min=1e-4).unsqueeze(-1)
+    u = K_j[0, 0] * uv[:, 0] + K_j[0, 2]
+    v = K_j[1, 1] * uv[:, 1] + K_j[1, 2]
+    inb = (zc > 1e-3) & (u >= 0) & (u <= W - 1) & (v >= 0) & (v <= H - 1)
+    # bilinear sample color + view-j depth at (u,v)
+    gx = (u / (W - 1)) * 2 - 1; gy = (v / (H - 1)) * 2 - 1
+    grid = torch.stack([gx, gy], dim=-1)[None, None]          # [1,1,N,2]
+    col = F.grid_sample(img_j[None], grid, align_corners=True, mode="bilinear")[0, :, 0].t()  # [N,3]
+    zj = F.grid_sample(z_j[None, None], grid, align_corners=True, mode="bilinear")[0, 0, 0]    # [N]
+    consistent = (zc - zj).abs() <= occ_tol * zc.clamp(min=1e-3)
+    return col, inb & consistent
+
+
+def fuse_colors(means_w, imgs, K, c2w, z, source_view, source_color, occ_tol=0.10, source_boost=1.0):
+    """means_w [N,3]; imgs [V,3,H,W]; K [V,3,3]; c2w [V,4,4]; z [V,H,W] per-view depth;
+    source_view [N] long; source_color [N,3] = the Gaussian's own (always-valid) source
+    pixel color. Seeds the fused color with the source color at weight source_boost (so a
+    Gaussian is never left colorless), then ADDS other views where in-frustum + depth-
+    consistent. Returns fused [N,3], n_extra_views [N] (how many OTHER views contributed)."""
+    V = imgs.shape[0]
+    N = means_w.shape[0]
+    acc = source_color * source_boost
+    wsum = torch.full((N,), float(source_boost), device=means_w.device)
+    extra = torch.zeros(N, device=means_w.device)
+    for j in range(V):
+        col, vis = reproject_sample(means_w, imgs[j], K[j], c2w[j], z[j], occ_tol)
+        vis = vis & (source_view != j)                       # source already seeded
+        w = vis.float()
+        acc += w.unsqueeze(-1) * col
+        wsum += w; extra += w
+    return acc / wsum.clamp(min=1e-6).unsqueeze(-1), extra

mapvggt/heads.py

@@ -0,0 +1,113 @@
+"""MapGS contribution heads for MapVGGT (VGGT-Omega backbone, frozen): external query
+heads that realize full MapTokenGS parity without surgery on VGGT internals.
+
+  * MapAnchorHead  -- contrib ① MAGT: HD-map-anchored Gaussian tokens whose centers are
+    bounded residuals  mu = anchor + s(t)*tanh(Delta)  around HD-map anchors, conditioned
+    on a pooled VGGT scene feature. Static, frame-independent.
+  * DynActorHead   -- PointForward scene-graph dynamics: per-instance object-canonical
+    Gaussians placed into each target frame by tracked box poses, with the lifespan
+    opacity envelope. Frame-dependent (placed at render time).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from mapgs.model.blocks import fourier_encode
+from mapgs.geometry.transforms import rotmat_to_quat, quat_multiply as quat_mul
+
+N_ANCHOR_TYPES = 3  # ground / lane / boundary
+
+
+class MapAnchorHead(nn.Module):
+    """anchors -> ng Gaussians each, centers = anchor + s(t)*tanh(Delta) (MAGT)."""
+    def __init__(self, c_ctx=2048, dim=256, ng=2, n_freq=10, pos_scale=50.0, base=0.5):
+        super().__init__()
+        self.ng, self.n_freq, self.pos_scale, self.base = ng, n_freq, pos_scale, base
+        pe_dim = 3 * (2 * n_freq + 1)
+        self.anchor_mlp = nn.Sequential(
+            nn.Linear(pe_dim + N_ANCHOR_TYPES + 3 + 1, dim), nn.GELU(), nn.Linear(dim, dim))
+        self.ctx_proj = nn.Linear(c_ctx, dim)
+        self.dec = nn.Linear(dim, ng * 14)              # per gaussian: pos3,logscale3,rot4,op1,rgb3
+        nn.init.zeros_(self.dec.weight); nn.init.zeros_(self.dec.bias)
+
+    def forward(self, anchor_pos, anchor_type, anchor_normal, ctx, s_t):
+        """anchor_pos [B,Na,3], type [B,Na], normal [B,Na,3], ctx [B,c_ctx]. Returns gaussian dict."""
+        B, Na, _ = anchor_pos.shape
+        pe = fourier_encode(anchor_pos / self.pos_scale, self.n_freq)
+        onehot = torch.zeros(B, Na, N_ANCHOR_TYPES, device=anchor_pos.device, dtype=pe.dtype)
+        onehot.scatter_(2, anchor_type.long().unsqueeze(-1).clamp(0, N_ANCHOR_TYPES - 1), 1.0)
+        feat = torch.cat([pe, onehot, anchor_normal, anchor_pos[..., 2:3]], dim=-1)
+        tok = self.anchor_mlp(feat) + self.ctx_proj(ctx)[:, None, :]      # [B,Na,dim]
+        o = self.dec(tok).view(B, Na, self.ng, 14)
+        anc = anchor_pos[:, :, None, :]                                  # [B,Na,1,3]
+        pos = anc + s_t * torch.tanh(o[..., 0:3])                        # bounded residual (MAGT)
+        scale = self.base * torch.exp(o[..., 3:6].clamp(-3, 3))
+        quat = F.normalize(o[..., 6:10] + o.new_tensor([1., 0, 0, 0]), dim=-1)
+        opacity = torch.sigmoid(o[..., 10] - 6.0)        # start ~off (sigmoid(-6)~0.002); sparsity-gated
+        rgb = torch.sigmoid(o[..., 11:14])
+        flat = lambda t, c: t.reshape(-1, c) if c > 1 else t.reshape(-1)
+        return dict(means=flat(pos, 3), scales=flat(scale, 3), quats=flat(quat, 4),
+                    opacities=flat(opacity, 1), colors=flat(rgb, 3))
+
+
+class DynActorHead(nn.Module):
+    """Learned object-canonical Gaussians per instance slot, placed per frame by box poses."""
+    def __init__(self, c_ctx=2048, dim=256, max_inst=8, ng=32):
+        super().__init__()
+        self.max_inst, self.ng = max_inst, ng
+        self.query = nn.Parameter(torch.randn(max_inst, dim) * 0.02)
+        self.ctx_proj = nn.Linear(c_ctx, dim)
+        self.dec = nn.Linear(dim, ng * 14)
+        nn.init.zeros_(self.dec.weight); nn.init.zeros_(self.dec.bias)
+
+    def canonical(self, ctx):
+        """Returns per-instance canonical gaussians (object-local unit cube). ctx [B,c_ctx]."""
+        tok = self.query + self.ctx_proj(ctx)                            # [max_inst,dim] (B=1)
+        o = self.dec(tok).view(self.max_inst, self.ng, 14)
+        return dict(
+            local=torch.tanh(o[..., 0:3]),                               # [-1,1] object cube
+            logscale=o[..., 3:6].clamp(-3, 3), rot=o[..., 6:10],
+            op=torch.sigmoid(o[..., 10] - 2.0), rgb=torch.sigmoid(o[..., 11:14]))
+
+    def place(self, canon, frame, box_c, box_R, canon_idx, box_valid, radius, base=0.5,
+              sigma=2.0, gain=1.0):
+        """Place all instances' canonical gaussians into `frame`. Returns gaussian dict
+        (world frame) with lifespan-tempered opacity. Mirrors place_dynamics_14d."""
+        I = box_c.shape[0]
+        n = min(self.max_inst, I)
+        if n == 0:
+            z = torch.zeros(0, device=box_c.device)
+            return dict(means=z.view(0, 3), scales=z.view(0, 3), quats=z.view(0, 4),
+                        opacities=z, colors=z.view(0, 3))
+        means, scales, quats, ops, rgbs = [], [], [], [], []
+        for i in range(n):
+            if not bool(box_valid[i, frame]):
+                continue
+            R = box_R[i, frame]; c = box_c[i, frame]
+            p_local = radius[i] * canon["local"][i]                      # [ng,3] metric object
+            p_world = torch.einsum("ij,nj->ni", R, p_local) + c
+            p_world = torch.nan_to_num(p_world).clamp(-200, 200)
+            q = quat_mul(rotmat_to_quat(R)[None].expand(self.ng, 4),
+                         F.normalize(canon["rot"][i] + canon["rot"].new_tensor([1., 0, 0, 0]), dim=-1))
+            df = (float(frame) - float(canon_idx[i])) / (sigma + 1.0)
+            life = gain * torch.exp(torch.tensor(-0.5 * df * df, device=c.device))
+            means.append(p_world)
+            scales.append(base * torch.exp(canon["logscale"][i]))
+            quats.append(q); ops.append(canon["op"][i] * life); rgbs.append(canon["rgb"][i])
+        if not means:
+            z = torch.zeros(0, device=box_c.device)
+            return dict(means=z.view(0, 3), scales=z.view(0, 3), quats=z.view(0, 4),
+                        opacities=z, colors=z.view(0, 3))
+        return dict(means=torch.cat(means), scales=torch.cat(scales), quats=torch.cat(quats),
+                    opacities=torch.cat(ops), colors=torch.cat(rgbs))
+
+
+def cat_gaussians(*dicts):
+    """Union of gaussian dicts over the 5 render fields (skip empties; ignore extras like depth)."""
+    keys = ("means", "scales", "quats", "opacities", "colors")
+    ds = [d for d in dicts if d is not None and d["means"].shape[0] > 0]
+    if not ds:
+        return dicts[0]
+    return {k: torch.cat([d[k] for d in ds], dim=0) for k in keys}

mapvggt/model.py

@@ -0,0 +1,124 @@
+"""MapVGGT -- per-pixel feed-forward 3DGS warm-started from VGGT-Omega (1B-512),
+replacing the TokenGS / DA-V2 backbones. VGGT-Omega is a multi-view ViT that predicts
+per-pixel METRIC depth in each camera frame (no affine needed, unlike MapNuRec's DA-V2).
+We feed the N context views together (cross-view attention), take its metric depth, and
+attach a fresh per-pixel Gaussian head on [rgb, log-depth, depth-conf] -> opacity /
+log-scale / rotation; color = source pixel RGB. Each pixel is lifted to a world-space
+Gaussian using the KNOWN driving camera poses (we do not use VGGT's predicted poses),
+and the union over views is rendered by gsplat. Map-grounding (map-depth / free-space /
+extrap / vert losses) is applied in the trainer, exactly as in MapTokenGS / MapNuRec.
+"""
+from __future__ import annotations
+
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# VGGT-Omega repo + base checkpoint locations are environment-overridable so the package
+# works on any clone. Set VGGT_OMEGA_REPO (path to the facebookresearch/vggt-omega clone)
+# and MAPVGGT_VGGT_CKPT (path to vggt_omega_1b_512.pt, obtained from its FAIR-licensed HF repo).
+_VGGT_REPO = os.environ.get("VGGT_OMEGA_REPO", "/mnt/william/_vggt_omega_repo")
+if _VGGT_REPO and _VGGT_REPO not in sys.path:
+    sys.path.insert(0, _VGGT_REPO)
+VGGT_CKPT = os.environ.get("MAPVGGT_VGGT_CKPT", "/mnt/william/vggt_omega_ckpt/vggt_omega_1b_512.pt")
+
+
+def lift_to_world(z, K, c2w):
+    """z [V,H,W] metric depth (along optical axis) -> world points [V,H,W,3] (OpenCV +z fwd)."""
+    V, H, W = z.shape
+    dev = z.device
+    ys, xs = torch.meshgrid(torch.arange(H, device=dev), torch.arange(W, device=dev), indexing="ij")
+    u, v = (xs + 0.5).float(), (ys + 0.5).float()
+    fx, fy = K[:, 0, 0, None, None], K[:, 1, 1, None, None]
+    cx, cy = K[:, 0, 2, None, None], K[:, 1, 2, None, None]
+    x = (u - cx) / fx * z
+    y = (v - cy) / fy * z
+    p_cam = torch.stack([x, y, z], dim=-1)
+    R, t = c2w[:, :3, :3], c2w[:, :3, 3]
+    return torch.einsum("vij,vhwj->vhwi", R, p_cam) + t[:, None, None, :]
+
+
+class MapVGGT(nn.Module):
+    def __init__(self, depth_min=1.5, depth_max=120.0, finetune_backbone=False, ckpt=VGGT_CKPT,
+                 with_map=False, with_dyn=False, s_max=2.0):
+        super().__init__()
+        from vggt_omega.models import VGGTOmega
+        self.vggt = VGGTOmega()
+        sd = torch.load(ckpt, map_location="cpu")
+        self.vggt.load_state_dict(sd)
+        self.finetune_backbone = finetune_backbone
+        if not finetune_backbone:
+            self.vggt.eval()
+            for p in self.vggt.parameters():
+                p.requires_grad_(False)
+        self.dmin, self.dmax = depth_min, depth_max
+        # MapGS contribution heads (frozen-VGGT external query heads)
+        self.with_map, self.with_dyn = with_map, with_dyn
+        self.s_max = s_max; self.cur_s = s_max          # tempering radius for MAGT residual
+        if with_map:
+            from mapvggt.heads import MapAnchorHead
+            self.map_head = MapAnchorHead(c_ctx=2 * 1024)
+        if with_dyn:
+            from mapvggt.heads import DynActorHead
+            self.dyn_head = DynActorHead(c_ctx=2 * 1024)
+        # fresh per-pixel head on [rgb(3), log-depth(1), conf(1)] -> opacity(1), log-scale(3), rot(4)
+        self.head = nn.Sequential(
+            nn.Conv2d(5, 64, 3, padding=1), nn.GELU(),
+            nn.Conv2d(64, 64, 3, padding=1), nn.GELU(),
+            nn.Conv2d(64, 9, 1))
+        nn.init.zeros_(self.head[-1].weight); nn.init.zeros_(self.head[-1].bias)
+        # [0]=opacity, [1:4]=log-scale (anisotropic), [4:8]=rot quat (wxyz), [8]=spare
+        self.head[-1].bias.data[0] = 2.0          # opacity logit -> sigmoid~0.88
+        self.head[-1].bias.data[4] = 1.0          # identity quaternion w=1
+
+    def _run_vggt(self, images):
+        """-> depth [V,H,W], conf [V,H,W], ctx [2048] (pooled scene feature for the heads)."""
+        cm = torch.no_grad() if not self.finetune_backbone else torch.enable_grad()
+        with cm:
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                pred = self.vggt(images[None])           # auto-unsqueeze to [1,V,3,H,W]
+        depth = pred["depth"][0, ..., 0].float()         # [V,H,W]
+        conf = pred.get("depth_conf")
+        conf = conf[0].float() if conf is not None else torch.ones_like(depth)
+        ctx = pred["camera_and_register_tokens"][0].float().mean(dim=(0, 1))  # [2048]
+        return depth.clamp(self.dmin, self.dmax), conf, ctx
+
+    def vggt_depth(self, images):
+        d, c, _ = self._run_vggt(images)
+        return d, c
+
+    def forward(self, images, K, c2w, anchor_pos=None, anchor_type=None, anchor_normal=None):
+        """images [V,3,H,W] (0..1), K [V,3,3], c2w [V,4,4] -> gaussian dict (world frame).
+        If with_map and anchors given, unions the MAGT map-anchored gaussians. Stashes the
+        pooled ctx feature for the dynamic head (placed per-frame at render time)."""
+        V, _, H, W = images.shape
+        z, conf, ctx = self._run_vggt(images)            # [V,H,W] metric, detached if frozen
+        self._ctx = ctx
+        logd = torch.log(z.clamp(self.dmin, self.dmax))
+        cn = (conf / conf.amax(dim=(-2, -1), keepdim=True).clamp_min(1e-6))
+        h = self.head(torch.cat([images, logd[:, None], cn[:, None]], dim=1))  # [V,9,H,W]
+        h = h.permute(0, 2, 3, 1)                         # [V,H,W,9]
+        opacity = torch.sigmoid(h[..., 0])
+        base = (z / K[:, 0, 0, None, None]).clamp(min=1e-4)            # pixel footprint at depth
+        scale = base[..., None] * torch.exp(h[..., 1:4].clamp(-3, 3))  # [V,H,W,3] anisotropic
+        quat = F.normalize(h[..., 4:8] + torch.tensor([1.0, 0, 0, 0], device=images.device), dim=-1)
+        xyz = lift_to_world(z, K, c2w)                   # [V,H,W,3] world
+        rgb = images.permute(0, 2, 3, 1)                 # color = source pixel
+        static = dict(means=xyz.reshape(-1, 3), scales=scale.reshape(-1, 3),
+                      quats=quat.reshape(-1, 4), opacities=opacity.reshape(-1),
+                      colors=rgb.reshape(-1, 3), depth=z)
+        if self.with_map and anchor_pos is not None:
+            from mapvggt.heads import cat_gaussians
+            mg = self.map_head(anchor_pos, anchor_type, anchor_normal, ctx[None], self.cur_s)
+            self._map_opacity = mg["opacities"]          # for the sparsity penalty (gate extras OFF by default)
+            u = cat_gaussians(static, mg)
+            u["depth"] = z                               # keep per-pixel depth for map-depth loss
+            return u
+        self._map_opacity = None
+        return static
+
+    def dyn_canonical(self):
+        """Per-instance canonical dynamic gaussians from the stashed ctx (call after forward)."""
+        return self.dyn_head.canonical(self._ctx)

mapvggt/refine.py

@@ -0,0 +1,56 @@
+"""Render-then-refine 2D UNet decoder (PointForward-style image-space refinement).
+
+The splat renders RGB + depth + alpha; a small 2D UNet corrects splatting artifacts
+(seams between per-view gaussians, sub-pixel misalignment, low-alpha haze, small holes)
+that a direct-RGB splat bakes in permanently. Output is a residual on the rendered RGB,
+zero-initialized so the model starts exactly at the un-refined baseline and only improves.
+Trained end-to-end (gradients flow through gsplat to the gaussians/backbone too).
+"""
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _cbr(i, o):
+    return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.GroupNorm(8, o), nn.SiLU())
+
+
+class RefineUNet(nn.Module):
+    def __init__(self, in_ch=5, base=48):
+        super().__init__()
+        self.e1 = _cbr(in_ch, base)
+        self.e2 = _cbr(base, base * 2)
+        self.e3 = _cbr(base * 2, base * 4)
+        self.pool = nn.MaxPool2d(2)
+        self.bott = _cbr(base * 4, base * 4)
+        self.u3 = _cbr(base * 4 + base * 4, base * 2)
+        self.u2 = _cbr(base * 2 + base * 2, base)
+        self.u1 = _cbr(base + base, base)
+        self.out = nn.Conv2d(base, 3, 1)
+        nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias)   # start == rendered RGB
+
+    def _up(self, x, ref):
+        return F.interpolate(x, size=ref.shape[-2:], mode="bilinear", align_corners=False)
+
+    def forward(self, rgb, depth, alpha):
+        """rgb [S,3,H,W] in [0,1]; depth [S,H,W]; alpha [S,H,W]. Returns refined rgb [S,3,H,W]."""
+        H, W = rgb.shape[-2:]
+        dn = (depth.clamp(0, 100) / 50.0).unsqueeze(1)          # normalized depth channel
+        a = alpha.clamp(0, 1).unsqueeze(1)
+        x = torch.cat([rgb, dn, a], dim=1)                      # [S,5,H,W]
+        # pad to a multiple of 8 (3 pooling stages) so any H,W works; crop back at the end
+        ph, pw = (8 - H % 8) % 8, (8 - W % 8) % 8
+        if ph or pw:
+            x = F.pad(x, (0, pw, 0, ph), mode="replicate")
+            rgb = F.pad(rgb, (0, pw, 0, ph), mode="replicate")
+        e1 = self.e1(x)
+        e2 = self.e2(self.pool(e1))
+        e3 = self.e3(self.pool(e2))
+        b = self.bott(self.pool(e3))
+        d3 = self.u3(torch.cat([self._up(b, e3), e3], 1))
+        d2 = self.u2(torch.cat([self._up(d3, e2), e2], 1))
+        d1 = self.u1(torch.cat([self._up(d2, e1), e1], 1))
+        res = self.out(d1)
+        out = (rgb + res).clamp(0, 1)
+        return out[..., :H, :W]

mapvggt/uncertainty.py

@@ -0,0 +1,49 @@
+"""Idea B -- Uncertainty-conditioned Gaussians.
+
+VGGT-Omega emits a per-pixel depth confidence. Depth uncertainty is a 1-D uncertainty
+ALONG THE VIEWING RAY. We shape each per-pixel Gaussian's covariance accordingly:
+orient one axis along the world-space ray from the source camera and set its extent to
+the depth standard deviation sigma_z (large for far / low-confidence pixels), while the
+two perpendicular axes keep the pixel footprint. An uncertain far Gaussian thus renders
+as a soft blur over the depth range it could occupy, instead of a hard point at a wrong
+depth (a ghost). Confident pixels -> sigma_z~0 -> ordinary near-isotropic splat.
+"""
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+from mapgs.geometry.transforms import rotmat_to_quat
+
+
+def ray_frame_quat(ray_dir):
+    """Per-row rotation (as wxyz quat) whose LOCAL Z axis = ray_dir [N,3] (unit)."""
+    z = F.normalize(ray_dir, dim=-1)
+    up = torch.tensor([0.0, 0.0, 1.0], device=z.device).expand_as(z)
+    # if ray ~ parallel to world-up, switch reference to avoid degenerate cross product
+    deg = (z * up).sum(-1).abs() > 0.99
+    up = torch.where(deg.unsqueeze(-1), torch.tensor([0.0, 1.0, 0.0], device=z.device), up)
+    x = F.normalize(torch.cross(up, z, dim=-1), dim=-1)
+    y = torch.cross(z, x, dim=-1)
+    R = torch.stack([x, y, z], dim=-1)            # columns = local axes in world (local->world)
+    return rotmat_to_quat(R)                       # [N,4] wxyz
+
+
+def uncertainty_to_sigma(z, conf, gamma, conf_ref=None):
+    """z [N] metric depth, conf [N] VGGT confidence (>=1, higher=more certain).
+    Returns along-ray sigma_z [N] in meters. u in [0,1] = relative uncertainty."""
+    cref = conf_ref if conf_ref is not None else conf.amax().clamp(min=1e-3)
+    u = (1.0 - conf / cref).clamp(0.0, 1.0)        # 0 at most-confident pixel, ->1 uncertain
+    return gamma * z * u
+
+
+def ray_uncertainty_cov(means, z, conf, cam_centers, base_perp, gamma, conf_ref=None,
+                        sigma_max_frac=0.5):
+    """means [N,3] world centers; z [N]; conf [N]; cam_centers [N,3] source camera centers;
+    base_perp [N] perpendicular (footprint) scale. Returns scales [N,3], quats [N,4]:
+    covariance elongated along the ray by sigma_z (capped at sigma_max_frac*z)."""
+    ray = means - cam_centers
+    quat = ray_frame_quat(ray)
+    sig = uncertainty_to_sigma(z, conf, gamma, conf_ref).clamp(max=sigma_max_frac * z)
+    s_along = base_perp + sig
+    scales = torch.stack([base_perp, base_perp, s_along], dim=-1)
+    return scales, quat