Untitled

import math
import numpy as np
from typing import Any, Dict, Tuple


def _normalize_rotation(R: np.ndarray) -> np.ndarray:
    # Orthonormalize with SVD (robust)
    U, _, Vt = np.linalg.svd(R)
    Rn = U @ Vt
    if np.linalg.det(Rn) < 0:
        U[:, -1] *= -1
        Rn = U @ Vt
    return Rn


def _quat_from_rotmat(R: np.ndarray) -> Tuple[float, float, float, float]:
    """
    Rotation matrix (3x3) -> quaternion (x,y,z,w)
    """
    R = _normalize_rotation(R)
    tr = R[0, 0] + R[1, 1] + R[2, 2]

    if tr > 0.0:
        S = math.sqrt(tr + 1.0) * 2.0
        w = 0.25 * S
        x = (R[2, 1] - R[1, 2]) / S
        y = (R[0, 2] - R[2, 0]) / S
        z = (R[1, 0] - R[0, 1]) / S
    elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
        S = math.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2.0
        w = (R[2, 1] - R[1, 2]) / S
        x = 0.25 * S
        y = (R[0, 1] + R[1, 0]) / S
        z = (R[0, 2] + R[2, 0]) / S
    elif R[1, 1] > R[2, 2]:
        S = math.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2.0
        w = (R[0, 2] - R[2, 0]) / S
        x = (R[0, 1] + R[1, 0]) / S
        y = 0.25 * S
        z = (R[1, 2] + R[2, 1]) / S
    else:
        S = math.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2.0
        w = (R[1, 0] - R[0, 1]) / S
        x = (R[0, 2] + R[2, 0]) / S
        y = (R[1, 2] + R[2, 1]) / S
        z = 0.25 * S

    return float(x), float(y), float(z), float(w)


def _convert_R_vision_to_babylon(R: np.ndarray, transpose_R: bool = False) -> np.ndarray:
    """
    Convert a rotation matrix from a vision-like convention (image y down)
    to Babylon-like (y up) using F = diag(1, -1, 1).

    R_bab = F * R * F

    transpose_R=True if later you notice it's inverted (world<->cam convention mismatch).
    """
    F = np.array([[1.0, 0.0, 0.0],
                  [0.0, -1.0, 0.0],
                  [0.0, 0.0, 1.0]], dtype=np.float64)

    if transpose_R:
        R = R.T

    Rb = F @ R @ F
    return _normalize_rotation(Rb)


def to_babylon_dict(
    estimate: Dict[str, Any],
    camera_height: float = 1.6,
    camera_distance: float = 4.0,
    transpose_R: bool = False,
) -> Dict[str, Any]:
    """
    Transform your estimate dict into a Babylon-ready dict.
    """
    cam = estimate["camera"]

    fov = float(cam["fov_y"])  # radians already
    R_align = np.asarray(cam["r_align"], dtype=np.float64)
    if R_align.shape != (3, 3):
        raise ValueError(f"camera.r_align must be 3x3, got {R_align.shape}")

    R_bab = _convert_R_vision_to_babylon(R_align, transpose_R=transpose_R)
    qx, qy, qz, qw = _quat_from_rotmat(R_bab)

    out = {
        "image": estimate.get("image", None),
        "camera": {
            "fov": fov,
            "position": [0.0, float(camera_height), float(-camera_distance)],
            "rotationQuaternion": [qx, qy, qz, qw],
        },
        # Optionnel : garder ces infos si tu veux debug/placer une grille plus tard
        "debug": {
            "fx": float(cam["fx"]),
            "fy": float(cam["fy"]),
            "cx": float(cam["cx"]),
            "cy": float(cam["cy"]),
            "plane": estimate.get("plane", None),
        },
    }

    # Nettoyage des None (si pas d'image/plane)
    if out["image"] is None:
        del out["image"]
    if out["debug"]["plane"] is None:
        del out["debug"]["plane"]

    return out