import math import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # Legacy edge-copy padding (kept for reference / ablations). # # Copies whole-pixel edges from the 4 adjacent faces and averages the 4 # corners. Exact only for continuous view-invariant signals at pad=1; for # wider pads or face-local signals (z-depth) it accumulates distortion # rapidly. `cube_resample_pad` below supersedes this as the default. # --------------------------------------------------------------------------- orderings = [ [0, 1, 3, 4, 5], [1, 2, 0, 4, 5], [2, 3, 1, 4, 5], [3, 0, 2, 4, 5], [4, 1, 3, 2, 0], [5, 1, 3, 0, 2], ] rotations = [ [0, 0, 0, 0, 0], [0, 0, 0,-1, 1], [0, 0, 0, 2, 2], [0, 0, 0, 1,-1], [0, 1,-1, 2, 0], [0,-1, 1, 0, 2] ] def _take_right(face, rot): if rot == 0: return face[..., :, 0] elif rot == 1: return face[..., 0, :].flip(-1) elif rot == 2: return face[..., :, -1].flip(-1) elif rot == -1: return face[..., -1, :] def _take_left(face, rot): if rot == 0: return face[..., :, -1] elif rot == 1: return face[..., -1, :].flip(-1) elif rot == 2: return face[..., :, 0].flip(-1) elif rot == -1: return face[..., 0, :] def _take_top(face, rot): if rot == 0: return face[..., -1, :] elif rot == 1: return face[..., :, 0] elif rot == 2: return face[..., 0, :].flip(-1) elif rot == -1: return face[..., :, -1].flip(-1) def _take_bottom(face, rot): if rot == 0: return face[..., 0, :] elif rot == 1: return face[..., :, -1] elif rot == 2: return face[..., -1, :].flip(-1) elif rot == -1: return face[..., :, 0].flip(-1) def valid_pad_conv_fn(x): assert x.ndim == 4 N, C, H, W = x.shape # Seam-aware edge copy needs all 6 neighbouring faces. Fall back to plain # reflect padding when the input is a sub-cubemap (e.g. single equatorial # face per sample in one_face_mode mode). if N % 6 != 0: return F.pad(x, [1, 1, 1, 1], mode='reflect') B = N // 6 # Reshape to (B, 6, C, H, W) to handle batches of cubemaps x_reshaped = x.view(B, 6, C, H, W) y = x.new_empty(B, 6, C, H+2, W+2) y[..., 1:-1, 1:-1] = x_reshaped for i in range(6): r_idx, l_idx, t_idx, b_idx = orderings[i][1:5] r_rot, l_rot, t_rot, b_rot = rotations[i][1:5] r_edge = _take_right (x_reshaped[:, r_idx], r_rot) l_edge = _take_left (x_reshaped[:, l_idx], l_rot) t_edge = _take_top (x_reshaped[:, t_idx], t_rot) b_edge = _take_bottom(x_reshaped[:, b_idx], b_rot) y[:, i, :, 1:-1, 0 ] = l_edge y[:, i, :, 1:-1, -1 ] = r_edge y[:, i, :, 0, 1:-1] = t_edge y[:, i, :, -1, 1:-1] = b_edge y[:, i, :, 0, 0 ] = 0.5*(y[:, i, :, 0, 1] + y[:, i, :, 1, 0]) y[:, i, :, 0, -1 ] = 0.5*(y[:, i, :, 0, -2] + y[:, i, :, 1, -1]) y[:, i, :, -1, 0 ] = 0.5*(y[:, i, :, -2, 0] + y[:, i, :, -1, 1]) y[:, i, :, -1,-1 ] = 0.5*(y[:, i, :, -2, -1] + y[:, i, :, -1, -2]) # Flatten back to (N, C, H+2, W+2) return y.view(N, C, H+2, W+2) # --------------------------------------------------------------------------- # Spherical-resample padding (DreamCube-style). # # The simple `valid_pad_conv_fn` above copies edge pixels from adjacent faces. # That is exact for continuous scalar fields whose value depends only on the # 3-D point (e.g. euclidean depth, RGB radiance), but wrong when the signal is # expressed *in each face's local frame* — notably z-depth, which measures the # distance along the face's own viewing axis. A neighbour face's z-depth at # the shared edge is generally a different number than the current face would # observe, because the neighbour projects the same 3-D point against a # different forward axis. # # DreamCube's fix: for each padded pixel build a 3-D ray from the face's own # intrinsics, find which neighbour face it hits, and grid-sample that face at # the exact sub-pixel location. We reuse the user's own intrinsics/extrinsics # convention so no axis-flip bookkeeping is needed. # --------------------------------------------------------------------------- _CUBE_EXTRINSICS_CACHE = None # extrinsics are resolution-independent _CUBE_PAD_GRID_CACHE = {} # keyed by (H, W, padding, device, dtype) def _get_default_extrinsics(device): """Return (6, 3, 3) world→camera rotation matrices for the canonical cube. Matches `src.utils.geometry_utils.get_cubemap_intrinsics_extrinsics`: Front +Z, Right +X, Back -Z, Left -X, Top -Y (y-down world), Bottom +Y. """ global _CUBE_EXTRINSICS_CACHE if _CUBE_EXTRINSICS_CACHE is None: face_configs = torch.tensor([ (0.0, 0.0), # Front (-90.0, 0.0), # Right (180.0, 0.0), # Back (90.0, 0.0), # Left (0.0, -90.0), # Top (0.0, 90.0), # Bottom ], dtype=torch.float64) y = torch.deg2rad(face_configs[:, 0]) p = torch.deg2rad(face_configs[:, 1]) cy, sy = torch.cos(y), torch.sin(y) cp, sp = torch.cos(p), torch.sin(p) Ry = torch.zeros(6, 3, 3, dtype=torch.float64) Ry[:, 0, 0] = cy; Ry[:, 0, 2] = sy Ry[:, 1, 1] = 1.0 Ry[:, 2, 0] = -sy; Ry[:, 2, 2] = cy Rx = torch.zeros(6, 3, 3, dtype=torch.float64) Rx[:, 0, 0] = 1.0 Rx[:, 1, 1] = cp; Rx[:, 1, 2] = -sp Rx[:, 2, 1] = sp; Rx[:, 2, 2] = cp R = torch.bmm(Rx, Ry) # Snap near-zero entries caused by float sin/cos of +-90. R[R.abs() < 1e-10] = 0.0 _CUBE_EXTRINSICS_CACHE = R.to(torch.float32) return _CUBE_EXTRINSICS_CACHE.to(device) def _make_intrinsics(H, W, fov_deg, device): assert H == W, f"cube faces must be square, got ({H}, {W})" f = H / (2.0 * math.tan(math.radians(fov_deg / 2.0))) K = torch.tensor([[f, 0, W / 2.0], [0, f, H / 2.0], [0, 0, 1.0]], dtype=torch.float32, device=device) return K.unsqueeze(0).expand(6, -1, -1) def _build_cube_pad_grid(H, W, padding, fov_deg, device): """Precompute the 3-D grid-sample grid and mask for cubemap resample padding. Returns: grid: (6, H_pad, W_pad, 3) float32, coords (u_norm, v_norm, face_z_norm) ready for F.grid_sample on a volume of shape (B, C, 6, H, W). mask: (H_pad, W_pad) bool, True on the padded border. """ key = (H, W, padding, fov_deg, str(device)) cached = _CUBE_PAD_GRID_CACHE.get(key) if cached is not None: return cached P = padding H_pad, W_pad = H + 2 * P, W + 2 * P R_all = _get_default_extrinsics(device) # (6, 3, 3), world→cam K_all = _make_intrinsics(H, W, fov_deg, device) # (6, 3, 3) # Padded-pixel centres in the original face coordinate system. # Pixel k is centred at k+0.5 (image-coord convention matching the user's # intrinsics where cx = W/2, so ray_x at centre(k) = (k+0.5-cx)/fx and the # face boundaries land at u = 0 and u = W, symmetric around cx). # Using centres (not edges) is what prevents argmax ties when a padded # ray lies exactly on a cube corner. v_pix, u_pix = torch.meshgrid( torch.arange(H_pad, device=device, dtype=torch.float32) + 0.5 - P, torch.arange(W_pad, device=device, dtype=torch.float32) + 0.5 - P, indexing='ij', ) ones = torch.ones_like(u_pix) # Ray in each source face's camera frame, then lifted into the world. ray_world_list = [] for i in range(6): fx = K_all[i, 0, 0]; fy = K_all[i, 1, 1] cx = K_all[i, 0, 2]; cy = K_all[i, 1, 2] dx = (u_pix - cx) / fx dy = (v_pix - cy) / fy ray_cam = torch.stack([dx, dy, ones], dim=-1) # (H_pad, W_pad, 3) # world←cam: ray_w = R^T @ ray_c → einsum 'ji,hwj->hwi' ray_world = torch.einsum('ji,hwj->hwi', R_all[i], ray_cam) ray_world_list.append(ray_world) ray_world = torch.stack(ray_world_list, dim=0) # (6, H_pad, W_pad, 3) # Pick the target face for each pixel: the one whose forward axis has the # largest projection with the ray. Forward_j in world coords = R_j[2, :] # (third row of the world→cam rotation, equivalently the cam z-axis # expressed in world). forwards = R_all[:, 2, :] # (6, 3) dots = torch.einsum('kc,shwc->shwk', forwards, ray_world) # (6_src, H_pad, W_pad, 6_tgt) face_j = torch.argmax(dots, dim=-1) # (6, H_pad, W_pad) # Transform each ray back to its target face's camera frame, then project. R_j = R_all[face_j] # (6, H_pad, W_pad, 3, 3) K_j = K_all[face_j] # (6, H_pad, W_pad, 3, 3) ray_cam_j = torch.einsum('shwab,shwb->shwa', R_j, ray_world) # (6, H_pad, W_pad, 3) z = ray_cam_j[..., 2:3].clamp(min=1e-6) pixel_j = torch.einsum('shwab,shwb->shwa', K_j, ray_cam_j / z) u_j = pixel_j[..., 0] v_j = pixel_j[..., 1] # u_j / v_j are in image-centre coords (pixel k centre at k+0.5), so the # align_corners=False normalisation is simply 2*u/W - 1 — the +1 in the # numerator is already absorbed into u_j's +0.5 offset. u_norm = 2.0 * u_j / W - 1.0 v_norm = 2.0 * v_j / H - 1.0 face_z_norm = (2.0 * face_j.to(torch.float32) + 1.0) / 6.0 - 1.0 grid = torch.stack([u_norm, v_norm, face_z_norm], dim=-1) # (6, H_pad, W_pad, 3) mask = torch.ones(H_pad, W_pad, dtype=torch.bool, device=device) mask[P:-P, P:-P] = False _CUBE_PAD_GRID_CACHE[key] = (grid, mask) return grid, mask def cube_resample_pad(x, padding, fov_deg=90.0): """DreamCube-style spherical resample padding. Args: x: (N, C, H, W) with N % 6 == 0 — cube faces stacked on the batch dim in the user's canonical order (Front, Right, Back, Left, Top, Bottom). padding: number of pixels to pad on every side. fov_deg: per-face FOV in degrees (90 for a standard cubemap). Returns: (N, C, H + 2P, W + 2P) — interior is the original content, border is filled by bilinearly resampling the geometrically-correct neighbour. """ assert x.ndim == 4, f"expected 4-D (N, C, H, W), got {tuple(x.shape)}" N, C, H, W = x.shape P = int(padding) if P <= 0: return x # Seam-aware padding only makes sense when we have the full 6-face cubemap # (needs neighbour faces to stitch from). Single-face / sub-cubemap inputs # (e.g. ScannetPano in one_face_mode mode) fall back to plain reflect # padding, which matches the "base" branch of the full-cubemap path for the # interior and handles the borders without any neighbour lookup. if N % 6 != 0: return F.pad(x, [P] * 4, mode='reflect') B = N // 6 H_pad, W_pad = H + 2 * P, W + 2 * P device, dtype = x.device, x.dtype grid, mask = _build_cube_pad_grid(H, W, P, fov_deg, device) # 3-D grid_sample on the 6-slice volume: (B, C, D=6, H, W) with grid # (B, D_out=6, H_pad, W_pad, 3). Using exact per-face z-centres and # padding_mode='border' means bilinear interpolation collapses onto the # selected face slice, while the border mode clamps u/v coords that # (numerically) overshoot by 1 ulp. x_vol = x.view(B, 6, C, H, W).permute(0, 2, 1, 3, 4).contiguous() grid_b = grid.to(torch.float32).unsqueeze(0).expand(B, -1, -1, -1, -1) sampled = F.grid_sample( x_vol.to(torch.float32), grid_b, mode='bilinear', padding_mode='border', align_corners=False, ) # (B, C, 6, H_pad, W_pad) sampled = sampled.permute(0, 2, 1, 3, 4) # (B, 6, C, H_pad, W_pad) # Reflect-pad the original as a baseline so the interior stays identical. base = F.pad(x, [P] * 4, mode='reflect').view(B, 6, C, H_pad, W_pad) out = torch.where(mask.view(1, 1, 1, H_pad, W_pad), sampled.to(dtype), base) return out.reshape(N, C, H_pad, W_pad) def make_cube_resample_pad_fn(padding=1, fov_deg=90.0): """Factory: build a pad_fn compatible with PaddedConv2d for a fixed padding.""" def _fn(x): return cube_resample_pad(x, padding=padding, fov_deg=fov_deg) return _fn # --------------------------------------------------------------------------- class PaddedConv2d(nn.Conv2d): def __init__(self, *args, pad_fn=None, **kwargs): kwargs = dict(kwargs) kwargs["padding"] = 0 super().__init__(*args, **kwargs) self.pad_fn = pad_fn def forward(self, x): x = self.pad_fn(x) return F.conv2d( x, self.weight, self.bias, stride=self.stride, padding=0, dilation=self.dilation, groups=self.groups ) @classmethod def from_existing(cls, conv: nn.Conv2d, pad_fn): new = cls( conv.in_channels, conv.out_channels, conv.kernel_size, stride=conv.stride, padding=0, dilation=conv.dilation, groups=conv.groups, bias=(conv.bias is not None), padding_mode="zeros", pad_fn=pad_fn ) new.weight = conv.weight if conv.bias is not None: new.bias = conv.bias return new def set_valid_pad_conv(module: nn.Module, fov_deg: float = 90.0): """Replace every kernel>1/padding>0 Conv2d with a seam-aware padded variant. Uses DreamCube-style spherical resample padding (`cube_resample_pad`). This is more accurate than the legacy edge-copy `valid_pad_conv_fn` (especially for pad>1 and face-local signals such as z-depth) and, somewhat counter-intuitively, also faster on GPU because a single fused grid_sample replaces the per-face Python loop. The legacy pad remains in this file for reference; call it directly if you need to ablate. Args: module: root module (walked recursively). fov_deg: per-face FOV in degrees (90 for a standard cubemap). """ for name, child in list(module.named_children()): if isinstance(child, nn.Conv2d): if child.kernel_size != (1, 1) and child.padding != (0, 0): P = int(child.padding[0]) pad_fn = make_cube_resample_pad_fn(padding=P, fov_deg=fov_deg) setattr(module, name, PaddedConv2d.from_existing(child, pad_fn)) else: set_valid_pad_conv(child, fov_deg=fov_deg)