senyilmaz-ComfyUI-Reframe360/projection.py

"""
Reframe360 — projection core.

This module is the single source of truth for the camera model.
js/reframe360.js contains a GLSL port of `build_direction_map` and
`trace_to_equirect`. If you change the math here, change it there too,
or the interactive preview will stop matching the rendered output.

Camera model
------------
An equirectangular image is mapped to the inside of a UNIT sphere (R = 1).
A pinhole/fisheye camera sits inside it, offset along ITS OWN Z axis by
perspective_warp, and looks down that same axis. Rays are cast from the
camera, intersected with the sphere, and the hit point is looked up in the
equirect texture.

The offset is camera-local: rotate the camera and the offset direction
rotates with it. Both the rays and the offset go through the same rotation
matrix.

Note on scale: sphere radius is fixed at 1.0 on purpose. Only the RATIO
of camera offset to radius affects the image, so a "sphere diameter"
control would be a redundant knob. perspective_warp IS that ratio.
"""

import math
import torch
import torch.nn.functional as F

# Real sensor sizes in mm (width, height).
SENSOR_PRESETS = {
    "Full Frame (36 x 24)": (36.0, 24.0),
    "Super35 (24.89 x 18.66)": (24.89, 18.66),
    "ARRI ALEXA 35 (27.99 x 19.22)": (27.99, 19.22),
    "ARRI ALEXA LF (36.70 x 25.54)": (36.70, 25.54),
    "RED KOMODO 6K (27.03 x 14.26)": (27.03, 14.26),
    "APS-C (23.6 x 15.6)": (23.6, 15.6),
    "Micro Four Thirds (17.3 x 13)": (17.3, 13.0),
    "Super16 (12.52 x 7.41)": (12.52, 7.41),
    "1 inch (13.2 x 8.8)": (13.2, 8.8),
    "iPhone main (9.8 x 7.3)": (9.8, 7.3),
    "Custom": (36.0, 24.0),
}

LENS_PROJECTIONS = ["rectilinear", "stereographic", "equisolid", "equidistant"]

ASPECT_PRESETS = {
    "Custom": None,
    "16:9": 16 / 9,
    "2.39:1 (scope)": 2.39,
    "1.85:1 (flat)": 1.85,
    "4:3": 4 / 3,
    "1:1": 1.0,
    "9:16 (vertical)": 9 / 16,
    "4:5 (social)": 4 / 5,
}


def resolve_sensor(sensor_w, sensor_h, sensor_fit, out_w, out_h):
    """
    Adjust the sensor to match the output aspect ratio.

    Without this, changing the output aspect would silently change the
    effective focal length — the classic "why doesn't 50mm look like 50mm"
    bug. Mirrors Blender's sensor_fit behaviour.

    Returns the effective (width, height) in mm.
    """
    out_aspect = out_w / out_h

    if sensor_fit == "auto":
        # Fit whichever output dimension is larger.
        fit = "horizontal" if out_w >= out_h else "vertical"
    else:
        fit = sensor_fit

    if fit == "horizontal":
        return sensor_w, sensor_w / out_aspect
    else:
        return sensor_h * out_aspect, sensor_h


def rotation_matrix(yaw_deg, pitch_deg, roll_deg, device, dtype):
    """
    Camera orientation. Order: roll (Z), then pitch (X), then yaw (Y).
    R = Ry(yaw) @ Rx(pitch) @ Rz(roll)

    yaw   — pan. Positive looks right.
    pitch — tilt. Positive looks UP, as in C4D / Blender / Nuke.
    roll  — dutch. Positive rolls clockwise.
    """
    y = math.radians(yaw_deg)
    p = math.radians(pitch_deg)
    r = math.radians(roll_deg)

    cy, sy = math.cos(y), math.sin(y)
    cp, sp = math.cos(p), math.sin(p)
    cr, sr = math.cos(r), math.sin(r)

    Ry = torch.tensor([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]], device=device, dtype=dtype)
    Rx = torch.tensor([[1.0, 0.0, 0.0], [0.0, cp, sp], [0.0, -sp, cp]], device=device, dtype=dtype)
    Rz = torch.tensor([[cr, -sr, 0.0], [sr, cr, 0.0], [0.0, 0.0, 1.0]], device=device, dtype=dtype)

    return Ry @ Rx @ Rz


def build_direction_map(out_w, out_h, sensor_w, sensor_h, focal, lens, device, dtype):
    """
    For every output pixel, work out which way the ray leaves the camera.

    Sensor position -> ray angle from the optical axis (theta), via the
    chosen lens projection. Camera looks down +Z.

    Returns (dirs, valid):
      dirs  — (H, W, 3) unit direction vectors in camera space
      valid — (H, W) bool, False where the lens can't form an image
    """
    # Pixel centres, in mm on the sensor plane.
    xs = (torch.arange(out_w, device=device, dtype=dtype) + 0.5) / out_w - 0.5
    ys = (torch.arange(out_h, device=device, dtype=dtype) + 0.5) / out_h - 0.5

    x_mm = xs * sensor_w
    y_mm = -ys * sensor_h  # image Y is down, world Y is up

    gy, gx = torch.meshgrid(y_mm, x_mm, indexing="ij")

    r = torch.sqrt(gx * gx + gy * gy)  # radial distance from sensor centre
    phi = torch.atan2(gy, gx)          # angle around the optical axis

    valid = torch.ones_like(r, dtype=torch.bool)

    if lens == "rectilinear":
        # theta = atan(r/f). Blows up approaching 180 deg FOV — that's
        # physics, not a bug. Use a fisheye projection for very short focals.
        theta = torch.atan2(r, torch.full_like(r, focal))
    elif lens == "stereographic":
        theta = 2.0 * torch.atan2(r, torch.full_like(r, 2.0 * focal))
    elif lens == "equisolid":
        # r = 2f*sin(theta/2)  ->  theta = 2*asin(r/2f). Undefined past r = 2f.
        s = r / (2.0 * focal)
        valid = s <= 1.0
        theta = 2.0 * torch.asin(s.clamp(max=1.0))
    elif lens == "equidistant":
        # r = f*theta. Wraps past the back of the sphere beyond theta = pi.
        theta = r / focal
        valid = theta <= math.pi
        theta = theta.clamp(max=math.pi)
    else:
        raise ValueError(f"unknown lens projection: {lens}")

    st = torch.sin(theta)
    dirs = torch.stack([st * torch.cos(phi), st * torch.sin(phi), torch.cos(theta)], dim=-1)
    return dirs, valid


def trace_to_equirect(dirs, offset):
    """
    Cast rays from `offset` and intersect the unit sphere, then convert the
    hit points to equirectangular UV.

    `offset` is the camera position in WORLD space, already rotated. It must
    be passed in rotated — an offset fixed to world Z would slide the camera
    sideways whenever you yaw away from front.

    Solving |O + t*d|^2 = 1 gives t^2 + 2t(O.d) + (|O|^2 - 1) = 0.
    Rotation is orthonormal so |O| == |perspective_warp| < 1, meaning the
    camera is always inside the sphere: the discriminant is always positive
    and exactly one root is positive. No misses, no edge cases.

    Returns (u, v) each in [0, 1].
    """
    ox, oy, oz = offset[0], offset[1], offset[2]

    b = dirs[..., 0] * ox + dirs[..., 1] * oy + dirs[..., 2] * oz   # O . d
    c = ox * ox + oy * oy + oz * oz - 1.0                           # |O|^2 - 1, always negative
    t = -b + torch.sqrt(b * b - c)                                  # the positive root

    px = t * dirs[..., 0] + ox
    py = t * dirs[..., 1] + oy
    pz = t * dirs[..., 2] + oz

    lon = torch.atan2(px, pz)                  # -pi .. pi
    lat = torch.asin(py.clamp(-1.0, 1.0))      # -pi/2 .. pi/2

    u = lon / (2.0 * math.pi) + 0.5
    v = 0.5 - lat / math.pi
    return u, v


def _wrap_pad(equi, pad):
    """
    Pad the equirect horizontally with wrapped pixels.

    grid_sample has no 'wrap' padding mode, so the 360-degree seam would
    show as a hard edge. Padding by hand fixes it.
    """
    left = equi[..., -pad:]
    right = equi[..., :pad]
    return torch.cat([left, equi, right], dim=-1)


def reframe(
    equi,                 # (B, C, H, W) float tensor, equirectangular
    out_w, out_h,
    sensor_w, sensor_h, sensor_fit,
    focal,
    lens,
    yaw, pitch, roll,
    warp,
    interpolation="bicubic",
):
    """
    Render one perspective view out of an equirectangular image.

    Everything runs as a single grid_sample, so batches (i.e. video)
    cost no more per frame than a still.
    """
    device, dtype = equi.device, equi.dtype

    eff_w, eff_h = resolve_sensor(sensor_w, sensor_h, sensor_fit, out_w, out_h)

    dirs, valid = build_direction_map(out_w, out_h, eff_w, eff_h, focal, lens, device, dtype)

    R = rotation_matrix(yaw, pitch, roll, device, dtype)
    dirs = torch.einsum("ij,hwj->hwi", R, dirs)

    # The camera slides along ITS OWN Z axis, not the world's. Rotate the
    # offset by the same matrix as the rays so it tracks where you're looking.
    offset = R @ torch.tensor([0.0, 0.0, float(warp)], device=device, dtype=dtype)

    u, v = trace_to_equirect(dirs, offset)

    _, _, H, W = equi.shape
    pad = 8
    padded = _wrap_pad(equi, pad)

    # Map u into the padded image, then to grid_sample's [-1, 1] space.
    u_px = u * W + pad
    gx = (u_px / (W + 2 * pad)) * 2.0 - 1.0
    gy = v * 2.0 - 1.0

    grid = torch.stack([gx, gy], dim=-1).unsqueeze(0).expand(equi.shape[0], -1, -1, -1)

    out = F.grid_sample(
        padded, grid,
        mode=interpolation,
        padding_mode="border",   # clamps at the poles
        align_corners=False,
    )

    if not valid.all():
        out = out * valid.to(dtype).view(1, 1, out_h, out_w)

    return out.clamp(0.0, 1.0)


def fov_degrees(sensor_w, sensor_h, focal, lens="rectilinear"):
    """Report the field of view, for the UI readout."""
    def angle(half_extent):
        if lens == "rectilinear":
            return math.degrees(2.0 * math.atan(half_extent / focal))
        if lens == "stereographic":
            return math.degrees(4.0 * math.atan(half_extent / (2.0 * focal)))
        if lens == "equisolid":
            s = min(half_extent / (2.0 * focal), 1.0)
            return math.degrees(4.0 * math.asin(s))
        if lens == "equidistant":
            return math.degrees(min(2.0 * half_extent / focal, 2.0 * math.pi))
        return 0.0

    return angle(sensor_w / 2.0), angle(sensor_h / 2.0)