Camera

def make_rt_inv(
    rotation: Float[torch.Tensor, "B 3"],
    translation: Float[torch.Tensor, "B 3"],
    orientation: str | None = "AP",
    isocenter: Float[torch.Tensor, "3"] | None = None,
) -> Float[torch.Tensor, "B 4 4"]:
    """Construct camera-to-world (extrinsic inverse) matrices.

    Computes `extrinsic_inv = pose @ reorient`, where `pose` encodes a
    ZXY Euler rotation and translation, and `reorient` is a fixed
    orientation-dependent transform.

    Args:
        rotation:    (B, 3) ZXY Euler angles in degrees.
        translation: (B, 3) Camera position in mm, relative to `isocenter`.
        orientation: Acquisition orientation — one of `"AP"`, `"PA"`, or `None`.
        isocenter:   Optional (3,) volume center in world coordinates.

    Returns:
        (B, 4, 4) camera-to-world matrices.
    """
    pose = convert(
        rotation,
        translation,
        "euler",
        convention="ZXY",
        isocenter=_default_isocenter(isocenter, rotation),
    )
    return pose @ _get_orientation_matrix(orientation, rotation.device, rotation.dtype)

def invert_rt_inv(
    extrinsic_inv: Float[torch.Tensor, "B 4 4"],
    orientation: str | None = "AP",
    isocenter: Float[torch.Tensor, "3"] | None = None,
) -> tuple[Float[torch.Tensor, "B 3"], Float[torch.Tensor, "B 3"]]:
    """Recover rotation and translation from camera-to-world matrices.

    Inverts the composition performed by [`make_rt_inv`](#nanodrr.camera.extrinsics.make_rt_inv). The
    `orientation` and `isocenter` arguments must match those used during
    construction to obtain correct results.

    Args:
        extrinsic_inv: (B, 4, 4) camera-to-world matrices.
        orientation:   Acquisition orientation — one of `"AP"`, `"PA"`, or `None`.
        isocenter:     Optional (3,) volume center in world coordinates.

    Returns:
        rotation:    (B, 3) ZXY Euler angles in degrees.
        translation: (B, 3) Camera position in mm, relative to `isocenter`.
    """
    device, dtype = extrinsic_inv.device, extrinsic_inv.dtype
    iso = _default_isocenter(isocenter, extrinsic_inv).view(1, 3)

    R_orient = _get_orientation_matrix(orientation, device, dtype)[:3, :3]

    R_pose = extrinsic_inv[..., :3, :3] @ R_orient.T
    t_pose = extrinsic_inv[..., :3, 3]

    translation = torch.einsum("bij,bj->bi", R_pose.transpose(-1, -2), t_pose - iso)
    rotation = rotmat_to_euler("ZXY", R_pose, degrees=True)

    return rotation, translation

def make_k_inv(
    sdd: float,
    delx: float,
    dely: float,
    x0: float,
    y0: float,
    height: int,
    width: int,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
) -> Float[torch.Tensor, "1 3 3"]:
    r"""Build the inverse intrinsic matrix $\mathbf K^{-1}$ for a cone-beam projector.

    $$
    \begin{align}
        f_x &= \frac{\mathrm{SDD}}{\Delta_x} &\quad c_x &= \frac{x_0}{\Delta_x} + \frac{W}{2} \\
        f_y &= \frac{\mathrm{SDD}}{\Delta_y} &\quad c_y &= \frac{y_0}{\Delta_y} + \frac{H}{2}
    \end{align}
    $$

    where `delx` $= \Delta_x$ and `dely` $= \Delta_y$.

    The returned matrix is the analytical inverse of the intrinsic matrix:

    $$
    \begin{equation}
        \mathbf K = \begin{bmatrix}
            f_x & 0 & c_x \\
            0 & f_y & c_y \\
            0 & 0 & 1
        \end{bmatrix} \,,
    \end{equation}
    $$

    Args:
        sdd: Source-to-detector distance (mm).
        delx: Pixel spacing in x (mm/px).
        dely: Pixel spacing in y (mm/px).
        x0: Principal-point offset from detector centre in x (mm).
        y0: Principal-point offset from detector centre in y (mm).
        height: Detector height in pixels.
        width: Detector width in pixels.
        dtype: Optional tensor dtype.
        device: Optional tensor device.

    Returns:
        (1, 3, 3) inverse intrinsic matrix.
    """
    fx = sdd / delx
    fy = sdd / dely
    cx = x0 / delx + width / 2.0
    cy = y0 / dely + height / 2.0

    return torch.tensor(
        [
            [
                [1.0 / fx, 0.0, -cx / fx],
                [0.0, 1.0 / fy, -cy / fy],
                [0.0, 0.0, 1.0],
            ]
        ],
        dtype=dtype,
        device=device,
    )

def resample(
    img: Float[torch.Tensor, "B C H W"],
    k_inv_old: Float[torch.Tensor, "B 3 3"],
    k_inv_new: Float[torch.Tensor, "B 3 3"],
) -> Float[torch.Tensor, "B C H W"]:
    r"""Resample an image from one camera's (inverse) intrinsic to another's.

    Each target pixel \(p'\) is mapped back to a source pixel via the homography

    $$
    \begin{equation}
        p = \mathbf H p' = \mathbf K_{\mathrm{old}}^{\phantom{-1}} \mathbf K_{\mathrm{new}}^{-1} p'
    \end{equation}
    $$

    and bilinearly interpolated. Out-of-bounds regions are filled with zeros.

    Args:
        img: Batch of images.
        k_inv_old: Inverse intrinsic matrices of the source images.
        k_inv_new: Inverse intrinsic matrices of the target images.

    Returns:
        Resampled images.
    """
    B, _, H, W = img.shape

    # Destination -> source homography per batch element
    H_mat = torch.linalg.inv(k_inv_old) @ k_inv_new

    # Build (H, W) grid of homogeneous destination pixel coords
    ys, xs = torch.meshgrid(
        torch.arange(H, dtype=img.dtype, device=img.device),
        torch.arange(W, dtype=img.dtype, device=img.device),
        indexing="ij",
    )
    ones = torch.ones_like(xs)
    coords = torch.stack([xs, ys, ones], dim=-1)

    # Apply per-batch homography
    coords_flat = coords.reshape(-1, 3).T.unsqueeze(0).expand(B, -1, -1)
    src = (H_mat @ coords_flat).permute(0, 2, 1).reshape(B, H, W, 3)

    # Perspective divide
    src_x = src[..., 0] / src[..., 2]
    src_y = src[..., 1] / src[..., 2]

    # Normalize to [-1, 1] for grid_sample
    grid_x = (src_x / (W - 1)) * 2 - 1
    grid_y = (src_y / (H - 1)) * 2 - 1
    grid = torch.stack([grid_x, grid_y], dim=-1)

    return F.grid_sample(img, grid, align_corners=True, padding_mode="zeros", mode="bilinear")