src/fpgm/geometry/camera.py

"""Pinhole camera: projection/unprojection paired with an explicit resolution.

The h5 ``intrinsic`` is valid at the scene-flow annotation resolution, which is not
generally the mp4 frame resolution. :meth:`Camera.rescaled` is the single
grep-able choke point for resolution changes -- :meth:`Camera.project` and
:meth:`Camera.unproject` never take a resolution argument, so it is structurally
impossible to project at the wrong scale without going through ``rescaled`` first.
"""

from __future__ import annotations

import logging
import warnings

import numpy as np

from fpgm.geometry.transforms import assert_valid_se3, invert_se3, transform_points
from fpgm.types import CameraIntrinsics

logger = logging.getLogger(__name__)


class Camera:
    """A pinhole camera: intrinsics plus a world-to-camera SE3 extrinsic."""

    def __init__(self, intrinsics: CameraIntrinsics, world_to_cam: np.ndarray) -> None:
        """Build a camera.

        Args:
            intrinsics: Pinhole intrinsics, valid at ``intrinsics.width/height``.
            world_to_cam: ``(4, 4)`` SE3 matrix mapping world points to camera frame.

        Raises:
            GeometryError: if ``world_to_cam`` is not a valid SE3 matrix -- this
                catches an accidental cam2world/world2cam swap at construction time
                rather than letting it manifest as silently-wrong depths later.
        """
        world_to_cam = np.asarray(world_to_cam, dtype=np.float64)
        assert_valid_se3(world_to_cam)
        self.K = intrinsics
        self._world_to_cam = world_to_cam
        self._cam_to_world = invert_se3(world_to_cam)

    @classmethod
    def from_pointworld(
        cls, intrinsic_3x3: np.ndarray, extrinsic_4x4: np.ndarray, width: int, height: int
    ) -> "Camera":
        """Build a :class:`Camera` from raw PointWorld-DROID h5 arrays.

        Args:
            intrinsic_3x3: ``(3, 3)`` K matrix, valid at ``(width, height)``.
            extrinsic_4x4: ``(4, 4)`` world-to-camera SE3 matrix (PointWorld stores
                ``extrinsic = inv(cam2world)``).
            width: Annotation-resolution pixel width the intrinsic is valid at.
            height: Annotation-resolution pixel height the intrinsic is valid at.
        """
        intrinsics = CameraIntrinsics.from_matrix(intrinsic_3x3, width=width, height=height)
        return cls(intrinsics, extrinsic_4x4)

    def rescaled(self, width: int, height: int) -> "Camera":
        """Return a new :class:`Camera` with intrinsics rescaled to ``(width, height)``.

        The extrinsic is unaffected -- only pixel-space quantities change with
        resolution. This is the only sanctioned way to change the resolution a
        camera operates at.
        """
        return Camera(self.K.scaled(width, height), self._world_to_cam)

    def world_to_cam(self, points_world: np.ndarray) -> np.ndarray:
        """Transform points ``(..., 3)`` from world frame to camera frame."""
        return transform_points(self._world_to_cam, points_world)

    def cam_to_world(self, points_cam: np.ndarray) -> np.ndarray:
        """Transform points ``(..., 3)`` from camera frame to world frame."""
        return transform_points(self._cam_to_world, points_cam)

    def project_cam(self, points_cam: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Project camera-frame points ``(..., 3)`` to pixels, intrinsic only.

        Returns:
            A tuple ``(uv, depth)`` of shapes ``(..., 2)`` and ``(...,)``. Points
            behind the camera (``depth <= 0``) are returned as ordinary data for the
            caller to mask -- this method never raises on them, since a track that
            temporarily goes behind the camera plane is a normal occurrence, not an
            error.
        """
        points_cam = np.asarray(points_cam, dtype=np.float64)
        z = points_cam[..., 2]
        safe_z = np.where(z == 0.0, np.finfo(np.float64).eps, z)
        u = self.K.fx * points_cam[..., 0] / safe_z + self.K.cx
        v = self.K.fy * points_cam[..., 1] / safe_z + self.K.cy
        uv = np.stack([u, v], axis=-1)
        return uv, z

    def project(self, points_world: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Project world-frame points ``(..., 3)`` to pixels via the extrinsic.

        Returns:
            A tuple ``(uv, depth)`` of shapes ``(..., 2)`` and ``(...,)``, where
            ``depth`` is ``Zc`` (the camera-frame Z coordinate). See
            :meth:`project_cam` for the behind-camera convention.
        """
        return self.project_cam(self.world_to_cam(points_world))

    def unproject_to_cam(self, uv: np.ndarray, depth: np.ndarray) -> np.ndarray:
        """Back-project pixels ``(..., 2)`` + depth ``(...,)`` to camera-frame points."""
        uv = np.asarray(uv, dtype=np.float64)
        depth = np.asarray(depth, dtype=np.float64)
        x = (uv[..., 0] - self.K.cx) * depth / self.K.fx
        y = (uv[..., 1] - self.K.cy) * depth / self.K.fy
        return np.stack([x, y, depth], axis=-1)

    def unproject(self, uv: np.ndarray, depth: np.ndarray) -> np.ndarray:
        """Back-project pixels ``(..., 2)`` + depth ``(...,)`` to world-frame points."""
        return self.cam_to_world(self.unproject_to_cam(uv, depth))

    def sanity_check_resolution(self, rgb_shape: tuple[int, ...]) -> None:
        """Warn loudly if ``rgb_shape`` disagrees with the principal point / claimed size.

        A well-formed pinhole intrinsic has its principal point near the image
        centre, i.e. ``2*cx ~= width`` and ``2*cy ~= height``. If that disagrees
        badly with ``rgb_shape``, the resolution this intrinsic is assumed to be
        valid at is likely wrong -- e.g. it was read at mp4 resolution instead of
        annotation resolution, or vice versa.

        Args:
            rgb_shape: ``(H, W, ...)`` shape of an actual image array to check against.
        """
        img_h, img_w = rgb_shape[0], rgb_shape[1]
        if img_w != self.K.width or img_h != self.K.height:
            warnings.warn(
                f"Camera resolution ({self.K.width}x{self.K.height}) does not match "
                f"image shape ({img_w}x{img_h}); did you forget to call .rescaled()?",
                stacklevel=2,
            )
        implied_w, implied_h = 2.0 * self.K.cx, 2.0 * self.K.cy
        rel_err_w = abs(implied_w - self.K.width) / max(self.K.width, 1)
        rel_err_h = abs(implied_h - self.K.height) / max(self.K.height, 1)
        if rel_err_w > 0.25 or rel_err_h > 0.25:
            warnings.warn(
                f"Principal point (2*cx={implied_w:.1f}, 2*cy={implied_h:.1f}) disagrees "
                f"badly with claimed resolution ({self.K.width}x{self.K.height}); the "
                "resolution this intrinsic is assumed valid at may be wrong.",
                stacklevel=2,
            )