File size: 4,855 Bytes
ea09995 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """MapNuRec — per-pixel feed-forward 3DGS, InstantNuRec-style, warm-started from
Depth-Anything-V2. Per input view i: DA-V2 predicts relative inverse-depth (disparity)
`disp`; a learned global affine maps it to metric depth `z = 1/(a*disp + b)` (a,b>0),
so DA-V2's strong relative geometry is inherited immediately and only the metric scale
is learned (anchored by the map-depth loss on ground). A small fresh per-pixel head on
[rgb, disp] predicts opacity / log-scale-mult / rotation; color = the source pixel RGB.
Each pixel is lifted to a world-space Gaussian; the union over views is rendered by gsplat.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
IMAGENET = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
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) # [V,H,W,3]
R, t = c2w[:, :3, :3], c2w[:, :3, 3]
return torch.einsum("vij,vhwj->vhwi", R, p_cam) + t[:, None, None, :]
class MapNuRec(nn.Module):
def __init__(self, depth_min=2.0, depth_max=80.0):
super().__init__()
from transformers import AutoModelForDepthEstimation
self.da = AutoModelForDepthEstimation.from_pretrained(
"depth-anything/Depth-Anything-V2-Small-hf") # DINOv2-S + DPT (warm-start)
# disparity -> metric depth affine z = 1/(softplus(a)*disp + softplus(b)).
# init from DA-V2 driving-disp stats (median~2.5): a=softplus(-3.8)=0.022,
# b=softplus(-4.6)=0.010 -> z(med)~15m, z(horizon,disp~0)~100m. Matches the map
# scale so the gamma-weighted map-depth anchor is active from step 0 (not zeroed).
self.aff = nn.Parameter(torch.tensor([-3.8, -4.6]))
self.dmin, self.dmax = depth_min, depth_max
# fresh per-pixel head on [rgb(3), disp(1)] -> opacity(1), log_scale_mult(1), rot(4)
self.head = nn.Sequential(
nn.Conv2d(4, 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)
# channel layout: [0]=opacity, [1:4]=log-scale (ANISOTROPIC), [4:8]=rot quat, [8]=spare
self.head[-1].bias.data[0] = 2.0 # opacity logit -> sigmoid~0.88
self.head[-1].bias.data[4] = 1.0 # rot quat w=1 (identity)
def disp(self, images):
"""images [V,3,H,W] in [0,1] -> per-pixel disparity [V,H,W] (larger=closer)."""
x = IMAGENET(images)
out = self.da(pixel_values=x).predicted_depth
if out.dim() == 4:
out = out[:, 0]
if out.shape[-2:] != images.shape[-2:]:
out = F.interpolate(out[:, None], size=images.shape[-2:], mode="bilinear",
align_corners=False)[:, 0]
return out.clamp(min=1e-3)
def forward(self, images, K, c2w):
"""images [V,3,H,W] (0..1), K [V,3,3], c2w [V,4,4] -> gaussian dict (world frame)."""
V, _, H, W = images.shape
disp = self.disp(images) # [V,H,W]
a, b = F.softplus(self.aff[0]), F.softplus(self.aff[1])
z = (1.0 / (a * disp + b)).clamp(self.dmin, self.dmax) # metric depth
h = self.head(torch.cat([images, disp[:, None]], dim=1)) # [V,9,H,W]
h = h.permute(0, 2, 3, 1) # [V,H,W,9]
opacity = torch.sigmoid(h[..., 0]) # [V,H,W]
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]
rgb = images.permute(0, 2, 3, 1) # color = source pixel
flat = lambda t, c: t.reshape(-1, c) if c > 1 else t.reshape(-1)
return 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)
|