src/utils/geometry_utils.py

import math
import torch
import numpy as np
from pytorch360convert import e2c, c2e, e2p


# Standard 'stack' face order: ['Front', 'Right', 'Back', 'Left', 'Up', 'Down'].
# (h_deg, v_deg) for e2p that reproduce the same faces as e2c at fov=90.
_CUBE_FACE_ANGLES = (
    (0.0, 0.0),     # Front
    (90.0, 0.0),    # Right
    (180.0, 0.0),   # Back
    (-90.0, 0.0),   # Left
    (0.0, 90.0),    # Up
    (0.0, -90.0),   # Down
)


def _fov_crop_size(face_w: int, fov_deg: float) -> int:
    """Pixel size of the central 90-deg crop inside a face_w-pixel face at fov_deg."""
    if fov_deg <= 90.0:
        return face_w
    return int(round(face_w / math.tan(math.radians(fov_deg / 2.0))))


def _e2p_stack(erp_tensor, face_w, fov_deg, mode):
    """Build a (6, C, face_w, face_w) cubemap stack from a single ERP using e2p."""
    faces = [
        e2p(erp_tensor, fov_deg=fov_deg, h_deg=h, v_deg=v,
            out_hw=(face_w, face_w), mode=mode)
        for h, v in _CUBE_FACE_ANGLES
    ]
    return torch.stack(faces, dim=0)


def _crop_overlap(cube_stack, face_w, fov_deg):
    """Center-crop the central 90-deg region from faces rendered at fov_deg.

    Discards the FOV>90 overlap so faces tile the sphere without gaps
    and can be fed to c2e (which assumes 90-deg faces).
    """
    crop = _fov_crop_size(face_w, fov_deg)
    if crop == face_w:
        return cube_stack
    start = (face_w - crop) // 2
    end = start + crop
    return cube_stack[..., start:end, start:end]


def erp_to_cubemap(erp_tensor, face_w=504, fov=90.0, cube_format="stack", mode="bilinear", **kwargs):
    """Convert ERP to cubemap stack.

    fov=90 uses the fast e2c path.  fov>90 uses e2p per face so neighbouring
    faces overlap — useful for cubemap-format losses that fight seams.
    """
    if fov == 90.0:
        if erp_tensor.ndim == 3:
            return e2c(erp_tensor, face_w=face_w, cube_format=cube_format, mode=mode, **kwargs)
        if erp_tensor.ndim == 4:
            return torch.stack([
                e2c(erp_tensor[b], face_w=face_w, cube_format=cube_format, mode=mode, **kwargs)
                for b in range(erp_tensor.shape[0])
            ], dim=0)
        raise ValueError(
            f"Unsupported erp_tensor shape {tuple(erp_tensor.shape)}. Expected (C, H, W) or (B, C, H, W)."
        )

    # fov > 90: use e2p per face (only 'stack' format supported)
    if cube_format != "stack":
        raise NotImplementedError(
            f"erp_to_cubemap with fov != 90 only supports cube_format='stack', got {cube_format!r}"
        )
    if erp_tensor.ndim == 3:
        return _e2p_stack(erp_tensor, face_w, fov, mode)
    if erp_tensor.ndim == 4:
        return torch.stack([
            _e2p_stack(erp_tensor[b], face_w, fov, mode)
            for b in range(erp_tensor.shape[0])
        ], dim=0)
    raise ValueError(
        f"Unsupported erp_tensor shape {tuple(erp_tensor.shape)}. Expected (C, H, W) or (B, C, H, W)."
    )


def cubemap_to_erp(cube_tensor, erp_h=1024, erp_w=2048, fov=90.0, cube_format="stack", mode="bilinear", **kwargs):
    """Convert cubemap stack to ERP.

    For fov>90 the overlap region of each face is center-cropped to 90 deg
    before stitching via c2e, so each ERP pixel comes from exactly one face.
    """
    if fov != 90.0:
        if cube_format != "stack":
            raise NotImplementedError(
                f"cubemap_to_erp with fov != 90 only supports cube_format='stack', got {cube_format!r}"
            )
        face_w = cube_tensor.shape[-1]
        cube_tensor = _crop_overlap(cube_tensor, face_w, fov)

    if cube_tensor.ndim == 4:
        if cube_tensor.shape[0] != 6:
            raise ValueError(
                f"Expected cube_tensor shape (6, C, h, w) for 4D input, got {tuple(cube_tensor.shape)}"
            )
        return c2e(cube_tensor, h=erp_h, w=erp_w, cube_format=cube_format, mode=mode, **kwargs)

    if cube_tensor.ndim == 5:
        if cube_tensor.shape[1] != 6:
            raise ValueError(
                f"Expected cube_tensor shape (B, 6, C, h, w) for 5D input, got {tuple(cube_tensor.shape)}"
            )
        erp_batch = []
        for b in range(cube_tensor.shape[0]):
            erp_batch.append(
                c2e(cube_tensor[b], h=erp_h, w=erp_w, cube_format=cube_format, mode=mode, **kwargs)
            )
        return torch.stack(erp_batch, dim=0)

    raise ValueError(
        f"Unsupported cube_tensor shape {tuple(cube_tensor.shape)}. Expected (6, C, h, w) or (B, 6, C, h, w)."
    )

def unit_normals(n, eps = 1e-6):
    assert n.dim() >= 3 and n.size(-3) == 3, "normals must have channel=3 at dim -3"
    denom = torch.clamp(torch.linalg.norm(n, dim=-3, keepdim=True), min=eps)
    return n / denom


def z_depth_to_euclidean(z_depth, intrinsics):
    # S may be any subset of cubemap faces — see euclidean_to_z_depth for the
    # same relaxation rationale (identical K per face).
    squeeze_batch = False
    if z_depth.ndim == 4:
        if z_depth.shape[1] != 1:
            raise ValueError(
                f"Expected z_depth shape (S, 1, H, W) for 4D input, got {tuple(z_depth.shape)}"
            )
        z_depth = z_depth.unsqueeze(0)
        squeeze_batch = True
    elif z_depth.ndim == 5:
        if z_depth.shape[2] != 1:
            raise ValueError(
                f"Expected z_depth shape (B, S, 1, H, W) for 5D input, got {tuple(z_depth.shape)}"
            )
    else:
        raise ValueError(
            f"Unsupported z_depth shape {tuple(z_depth.shape)}. Expected (6, 1, H, W) or (B, 6, 1, H, W)."
        )

    if intrinsics.ndim == 3:
        intrinsics = intrinsics[0]

    fx = intrinsics[0, 0]
    fy = intrinsics[1, 1]
    cx = intrinsics[0, 2]
    cy = intrinsics[1, 2]

    _, _, _, H, W = z_depth.shape
    u = torch.arange(W, device=z_depth.device, dtype=z_depth.dtype)
    v = torch.arange(H, device=z_depth.device, dtype=z_depth.dtype)
    vv, uu = torch.meshgrid(v, u, indexing="ij")

    x_norm = (uu - cx) / fx
    y_norm = (vv - cy) / fy
    ray_length_multiplier = torch.sqrt(x_norm * x_norm + y_norm * y_norm + 1.0)
    ray_length_multiplier = ray_length_multiplier.unsqueeze(0).unsqueeze(0).unsqueeze(0)

    out = z_depth * ray_length_multiplier
    return out.squeeze(0) if squeeze_batch else out


def compute_edge_mask(depth, abs_thresh = 0.1, rel_thresh = 0.1):
    assert depth.ndim == 2
    depth = depth.astype(np.float32, copy=False)

    valid = depth > 0
    eps = 1e-6

    edge = np.zeros_like(valid, dtype=bool)

    d1 = depth[:, :-1]
    d2 = depth[:, 1:]
    v_pair = valid[:, :-1] & valid[:, 1:]

    diff = np.abs(d1 - d2)
    rel = diff / (np.minimum(d1, d2) + eps)

    edge_pair = v_pair & (diff > abs_thresh) & (rel > rel_thresh)

    edge[:, :-1] |= edge_pair
    edge[:, 1:] |= edge_pair

    d1 = depth[:-1, :]
    d2 = depth[1:, :]
    v_pair = valid[:-1, :] & valid[1:, :]

    diff = np.abs(d1 - d2)
    rel = diff / (np.minimum(d1, d2) + eps)

    edge_pair = v_pair & (diff > abs_thresh) & (rel > rel_thresh)

    edge[:-1, :] |= edge_pair
    edge[1:, :]  |= edge_pair

    keep = valid & (~edge)
    return keep


_ERP_DIRS_CACHE: dict = {}


def _get_erp_dirs(H: int, W: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Cached (3, H, W) unit direction vectors for ERP → 3D unprojection.

    Convention: x = cos(phi)*cos(theta),  y = sin(phi),  z = cos(phi)*sin(theta)
    """
    key = (H, W, device, dtype)
    if key not in _ERP_DIRS_CACHE:
        u = (torch.arange(W, device=device, dtype=dtype) + 0.5) / W   # (W,)
        v = (torch.arange(H, device=device, dtype=dtype) + 0.5) / H   # (H,)
        theta = u * (2.0 * torch.pi) - torch.pi    # (W,) azimuth
        phi   = (0.5 - v) * torch.pi               # (H,) elevation

        cos_phi = torch.cos(phi)   # (H,)
        sin_phi = torch.sin(phi)   # (H,)
        cos_the = torch.cos(theta) # (W,)
        sin_the = torch.sin(theta) # (W,)

        dir_x = cos_phi[:, None] * cos_the[None, :]   # (H, W)
        dir_y = sin_phi[:, None].expand(H, W)          # (H, W)
        dir_z = cos_phi[:, None] * sin_the[None, :]   # (H, W)

        _ERP_DIRS_CACHE[key] = torch.stack([dir_x, dir_y, dir_z], dim=0)  # (3, H, W)
    return _ERP_DIRS_CACHE[key]


def erp_to_pointcloud(depth: torch.Tensor) -> torch.Tensor:
    """Differentiable ERP depth → XYZ point map.

    The direction grid is cached per (H, W, device, dtype) so calls inside a
    training loop pay only the cost of a single broadcast multiply.

    Args:
        depth: float tensor  (H, W) | (B, H, W) | (B, 1, H, W)

    Returns:
        xyz: float tensor, same dtype/device
             (H, W)      → (3, H, W)
             (B, H, W)   → (B, 3, H, W)
             (B, 1, H, W) → (B, 3, H, W)
    """
    ndim = depth.dim()
    if ndim == 2:
        H, W = depth.shape
        dirs = _get_erp_dirs(H, W, depth.device, depth.dtype)   # (3, H, W)
        return depth.unsqueeze(0) * dirs                         # (3, H, W)
    if ndim == 3:
        depth = depth.unsqueeze(1)                               # (B, 1, H, W)
    # depth is (B, 1, H, W) from here
    B, _, H, W = depth.shape
    dirs = _get_erp_dirs(H, W, depth.device, depth.dtype)        # (3, H, W)
    return depth * dirs.unsqueeze(0)                             # (B, 3, H, W)


def get_cubemap_intrinsics_extrinsics(image_size=512, fov=90.0):
    """Compute cubemap camera intrinsics and extrinsics for 6 faces.

    Args:
        image_size: face resolution in pixels (square faces assumed).
        fov: per-face field of view in degrees (>=90).  At fov=90 the
             focal length equals image_size/2 (standard cubemap).

    Returns:
        (extrinsics, intrinsics) — both float32 tensors on CPU.
        intrinsics:  (6, 3, 3)
        extrinsics:  (6, 4, 4)
    """
    S = float(image_size)
    f = S / (2.0 * math.tan(math.radians(fov / 2.0)))
    cx, cy = S / 2.0, S / 2.0

    K = torch.tensor([
        [f, 0, cx],
        [0, f, cy],
        [0, 0, 1]
    ], dtype=torch.float32)
    intrinsics = K.unsqueeze(0).expand(6, -1, -1).clone()

    face_configs = [
        (0, 0),    # Front  (Look +Z)
        (-90, 0),  # Right  (Look +X)
        (180, 0),  # Back   (Look -Z)
        (90, 0),   # Left   (Look -X)
        (0, -90),  # Top    (Look +Y)
        (0, 90)    # Bottom (Look -Y)
    ]

    face_configs = torch.tensor(face_configs, dtype=torch.float32)
    y = torch.deg2rad(face_configs[:, 0])
    p = torch.deg2rad(face_configs[:, 1])

    cy_r = torch.cos(y)
    sy = torch.sin(y)
    cp = torch.cos(p)
    sp = torch.sin(p)

    Ry = torch.zeros((6, 3, 3), dtype=torch.float32)
    Ry[:, 0, 0] = cy_r
    Ry[:, 0, 2] = sy
    Ry[:, 1, 1] = 1.0
    Ry[:, 2, 0] = -sy
    Ry[:, 2, 2] = cy_r

    Rx = torch.zeros((6, 3, 3), dtype=torch.float32)
    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)

    extrinsics = torch.zeros((6, 4, 4), dtype=torch.float32)
    extrinsics[:, :3, :3] = R
    extrinsics[:, 3, 3] = 1.0

    return extrinsics, intrinsics