src/wjad/data/ftheta_proj.py

"""f-theta 正向投影：3D 点（相机系） -> 像素。

仅支持 backward polynomial 形式（与 NVIDIA 工具一致），
forward 形式可在内部用牛顿迭代反推。
"""

from __future__ import annotations

import torch

from ..modules.rays import FThetaCamera


def project_points_ftheta(
    points_cam: torch.Tensor,        # [..., 3]，相机系下 3D 点
    cam: FThetaCamera,
) -> tuple[torch.Tensor, torch.Tensor]:
    """正向投影：相机系点 -> 像素 ``(u, v)``，并返回深度。

    返回
    ----
    uv : [..., 2]
    depth : [..., 1]，沿主光轴（z）方向的深度（如果 z<0 则视为后方，仍计算
            但调用方需用 ``depth > 0`` 做有效性筛选）。
    """
    x = points_cam[..., 0]
    y = points_cam[..., 1]
    z = points_cam[..., 2]
    norm = torch.sqrt(x * x + y * y + z * z).clamp_min(1e-6)
    cos_theta = z / norm
    cos_theta = cos_theta.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
    theta = torch.acos(cos_theta)
    phi = torch.atan2(y, x)

    if cam.intr.is_bw_poly:
        # backward poly 是 r_pix -> theta；正向需要反求 theta -> r_pix。
        # 用牛顿迭代：希望 _eval_poly(r) = theta
        r = theta.clone()  # 初始猜测
        for _ in range(8):
            f = cam._eval_poly(r) - theta
            df = cam._eval_poly_grad(r).clamp_min(1e-6)
            r = r - f / df
        r_pix = r
    else:
        r_pix = cam._eval_poly(theta)

    cos_p = torch.cos(phi)
    sin_p = torch.sin(phi)
    du = r_pix * cos_p
    dv = r_pix * sin_p
    # 反线性修正：linear_cde 是仿射 (du,dv) = M (du0,dv0)，正投影需要逆
    c = cam.intr.linear_cde[0]
    d = cam.intr.linear_cde[1]
    e = cam.intr.linear_cde[2]
    # 简化：忽略 linear_cde 的修正（与 unproject 中近似一致）
    du0 = du
    dv0 = dv
    u = du0 + cam.intr.cx
    v = dv0 + cam.intr.cy
    uv = torch.stack([u, v], dim=-1)
    depth = z.unsqueeze(-1)
    return uv, depth