src/depth_anything_3/model/utils/valid_conv_padding.py

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)