festifun/backend/festival4d/resolve.py

"""Annotation -> 3D point resolution (spec M8, integration phase).

Given a 2D bounding box drawn on one video at a local video time, recover a 3D scene point:

  1. Cast a ray from that camera through the bbox center, using the pose interpolated at
     ``t_video_s`` and the stored pinhole intrinsics.
  2. If another annotation for the **same event** exists from a **different video**,
     triangulate the two rays (reject near-parallel rays or a mutual gap > 0.5 scene units).
  3. Else take the nearest point-cloud point to the ray (within a 0.3-unit cylinder).
  4. Else return the ray at depth = distance from the camera to the point-cloud centroid.

This module is pure orchestration over lane B's frozen geometry primitives
(:func:`geometry.ray_from_pixel`, :func:`geometry.triangulate_rays`,
:func:`geometry.nearest_point_on_ray`, :func:`geometry.slerp_pose`) — it never reimplements
the pose math. It reads the point cloud via :func:`synthetic.read_ply`.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np

from festival4d import config, db, geometry
from festival4d.synthetic import read_ply

log = logging.getLogger("festival4d.resolve")

TRIANGULATION_MAX_GAP = 0.5   # scene units — reject looser two-view intersections (spec M8)
NEAREST_RADIUS = 0.3          # scene units — cylinder for the point-cloud fallback (spec M8)
PARALLEL_DOT = 0.9995         # |cos(angle)| above this => rays too parallel to triangulate


@dataclass
class Resolution:
    """Outcome of resolving one annotation to a 3D point."""

    point: list[float] | None
    method: str                 # 'triangulated' | 'nearest_point' | 'ray_depth' | 'no_geometry'
    gap: float | None = None    # mutual-approach distance when method == 'triangulated'


def _pose_tuple(p) -> tuple[list[float], list[float], tuple[float, float, float, float]]:
    return (
        [p.qw, p.qx, p.qy, p.qz],
        [p.tx, p.ty, p.tz],
        (p.fx, p.fy, p.cx, p.cy),
    )


def _pose_at(poses, t_video_s: float):
    """Interpolate a COLMAP pose ``(q, t, (fx,fy,cx,cy))`` at ``t_video_s``.

    ``poses`` is ascending by ``t_video_s`` (as returned by :func:`db.get_poses`). Clamps to
    the endpoints — no extrapolation, matching the frontend ``poseAt`` and spec M2.
    """
    if not poses:
        return None
    if t_video_s <= poses[0].t_video_s:
        return _pose_tuple(poses[0])
    if t_video_s >= poses[-1].t_video_s:
        return _pose_tuple(poses[-1])

    lo, hi = 0, len(poses) - 1
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if poses[mid].t_video_s <= t_video_s:
            lo = mid
        else:
            hi = mid
    a, b = poses[lo], poses[hi]
    span = b.t_video_s - a.t_video_s
    alpha = (t_video_s - a.t_video_s) / span if span > 1e-9 else 0.0
    q, t = geometry.slerp_pose(
        [a.qw, a.qx, a.qy, a.qz], [a.tx, a.ty, a.tz],
        [b.qw, b.qx, b.qy, b.qz], [b.tx, b.ty, b.tz],
        alpha,
    )
    # Intrinsics don't interpolate meaningfully across a fixed camera; use the earlier pose's.
    return list(q), list(t), (a.fx, a.fy, a.cx, a.cy)


def _ray_for_annotation(video_id: int, t_video_s: float, bbox):
    """World-space ray (origin, unit direction) through the bbox center, or ``None``.

    Returns ``None`` when the video has no poses (reconstruction absent for that camera).
    """
    video = db.get_video(video_id)
    if video is None:
        raise KeyError(f"no video with id={video_id}")
    pose = _pose_at(db.get_poses(video_id), t_video_s)
    if pose is None:
        return None
    q, t, (fx, fy, cx, cy) = pose
    x0, y0, x1, y1 = bbox
    px = (x0 + x1) / 2.0 * video.width
    py = (y0 + y1) / 2.0 * video.height
    return geometry.ray_from_pixel(q, t, fx, fy, cx, cy, px, py)


def _load_points() -> np.ndarray | None:
    if not config.POINTS_PLY.exists():
        return None
    try:
        pts, _colors = read_ply(config.POINTS_PLY)
    except Exception as exc:  # pragma: no cover - defensive
        log.warning("resolve: failed to read point cloud: %s", exc)
        return None
    return np.asarray(pts, dtype=np.float64).reshape(-1, 3)


def resolve_annotation(video_id: int, t_video_s: float, bbox, event_id: int | None = None) -> Resolution:
    """Resolve one annotation to a 3D point (see module docstring for the cascade)."""
    ray = _ray_for_annotation(video_id, t_video_s, bbox)
    if ray is None:
        return Resolution(point=None, method="no_geometry")
    origin, direction = ray

    # 1. Two-view triangulation against another annotation of the same event.
    if event_id is not None:
        for other in reversed(db.get_annotations()):  # most recent first
            if other.event_id != event_id or other.video_id == video_id:
                continue
            other_ray = _ray_for_annotation(
                other.video_id, other.t_video_s, [other.x0, other.y0, other.x1, other.y1]
            )
            if other_ray is None:
                continue
            o2, d2 = other_ray
            if abs(float(np.dot(direction, d2))) >= PARALLEL_DOT:
                continue  # near-parallel: no reliable intersection
            point, gap = geometry.triangulate_rays(origin, direction, o2, d2)
            if gap <= TRIANGULATION_MAX_GAP:
                return Resolution(
                    point=[float(v) for v in point], method="triangulated", gap=float(gap)
                )

    # 2/3. Point-cloud fallbacks.
    pts = _load_points()
    if pts is not None and len(pts):
        near = geometry.nearest_point_on_ray(origin, direction, pts, NEAREST_RADIUS)
        if near is not None:
            return Resolution(point=[float(v) for v in near], method="nearest_point")
        centroid = pts.mean(axis=0)
        depth = float(np.linalg.norm(centroid - origin))
        pt = origin + direction * depth
        return Resolution(point=[float(v) for v in pt], method="ray_depth")

    # 4. No point cloud at all: park at unit depth so a usable anchor still comes back.
    pt = origin + direction * 1.0
    return Resolution(point=[float(v) for v in pt], method="ray_depth")