postprocess_utils.py

from __future__ import annotations

from typing import Optional

import matplotlib
import numpy as np
import torch
import torch.nn.functional as F
from einops import rearrange
from plyfile import PlyData, PlyElement


def signed_log1p(x):
    """
    Computes log(1 + abs(x)) while keeping the original sign of x.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Transformed tensor with the same sign as x.
    """
    if isinstance(x, torch.Tensor):
        return torch.sign(x) * torch.log1p(torch.abs(x))
    elif isinstance(x, np.ndarray):
        return np.sign(x) * np.log1p(np.abs(x))
    else:
        raise TypeError("Input must be a torch.Tensor or numpy.ndarray")


def signed_log1p_inverse(x):
    """
    Computes the inverse of signed_log1p: x = sign(x) * (exp(abs(x)) - 1).

    Args:
        y (torch.Tensor): Input tensor (output of signed_log1p).

    Returns:
        torch.Tensor: Original tensor x.
    """
    if isinstance(x, torch.Tensor):
        return torch.sign(x) * (torch.exp(torch.abs(x)) - 1)
    elif isinstance(x, np.ndarray):
        return np.sign(x) * (np.exp(np.abs(x)) - 1)
    else:
        raise TypeError("Input must be a torch.Tensor or numpy.ndarray")


def colorize_depth(depth, cmap="Spectral"):
    min_d, max_d = (depth[depth > 0]).min(), (depth[depth > 0]).max()
    depth = (max_d - depth) / (max_d - min_d)

    cm = matplotlib.colormaps[cmap]
    depth = depth.clip(0, 1)
    depth = cm(depth, bytes=False)[..., 0:3]
    return depth


def save_ply(pointmap, image, output_file, downsample=20, mask=None):
    _, h, w, _ = pointmap.shape
    image = image[:, :h, :w]
    pointmap = pointmap[:, :h, :w]

    points = pointmap.reshape(-1, 3)  # (H*W, 3)
    colors = image.reshape(-1, 3)  # (H*W, 3)
    if mask is not None:
        points = points[mask.reshape(-1)]
        colors = colors[mask.reshape(-1)]

    indices = np.random.choice(
        colors.shape[0], int(colors.shape[0] / downsample), replace=False
    )
    points = points[indices]
    colors = colors[indices]

    vertices = []
    for p, c in zip(points, colors):
        vertex = (p[0], p[1], p[2], int(c[0]), int(c[1]), int(c[2]))
        vertices.append(vertex)

    vertex_dtype = np.dtype(
        [
            ("x", "f4"),
            ("y", "f4"),
            ("z", "f4"),
            ("red", "u1"),
            ("green", "u1"),
            ("blue", "u1"),
        ]
    )
    vertex_array = np.array(vertices, dtype=vertex_dtype)

    ply_element = PlyElement.describe(vertex_array, "vertex")
    PlyData([ply_element], text=True).write(output_file)


def fov_to_focal(fovx, fovy, h, w):
    focal_x = w * 0.5 / np.tan(fovx)
    focal_y = h * 0.5 / np.tan(fovy)
    focal = (focal_x + focal_y) / 2
    return focal


def get_rays(pose, h, w, focal=None, fovx=None, fovy=None):
    import torch.nn.functional as F

    pose = torch.from_numpy(pose).float()
    x, y = torch.meshgrid(
        torch.arange(w),
        torch.arange(h),
        indexing="xy",
    )
    x = x.flatten().unsqueeze(0).repeat(pose.shape[0], 1)
    y = y.flatten().unsqueeze(0).repeat(pose.shape[0], 1)

    cx = w * 0.5
    cy = h * 0.5
    intrinsics, focal = get_intrinsics(pose.shape[0], h, w, fovx, fovy, focal)

    if isinstance(focal, float):
        focal = np.array([focal])

    focal = torch.from_numpy(focal).float()
    camera_dirs = F.pad(
        torch.stack(
            [
                (x - cx + 0.5) / focal.unsqueeze(-1),
                (y - cy + 0.5) / focal.unsqueeze(-1),
            ],
            dim=-1,
        ),
        (0, 1),
        value=1.0,
    )  # [t, hw, 3]

    pose = pose.to(dtype=camera_dirs.dtype)
    rays_d = camera_dirs @ pose[:, :3, :3].transpose(1, 2)  # [t, hw, 3]

    rays_o = pose[:, :3, 3].unsqueeze(1).expand_as(rays_d)  # [hw, 3]

    rays_o = rays_o.view(pose.shape[0], h, w, 3)
    rays_d = rays_d.view(pose.shape[0], h, w, 3)

    return rays_o.float().numpy(), rays_d.float().numpy(), intrinsics


def get_intrinsics(batch_size, h, w, fovx=None, fovy=None, focal=None):
    if focal is None:
        focal_x = w * 0.5 / np.tan(fovx)
        focal_y = h * 0.5 / np.tan(fovy)
        focal = (focal_x + focal_y) / 2
    cx = w * 0.5
    cy = h * 0.5
    intrinsics = np.zeros((batch_size, 3, 3))
    intrinsics[:, 0, 0] = focal
    intrinsics[:, 1, 1] = focal
    intrinsics[:, 0, 2] = cx
    intrinsics[:, 1, 2] = cy
    intrinsics[:, 2, 2] = 1.0

    return intrinsics, focal


def save_pointmap(
    rgb,
    disparity,
    raymap,
    save_file,
    vae_downsample_scale=8,
    camera_pose=None,
    ray_o_scale_inv=1.0,
    max_depth=1e2,
    save_full_pcd_videos=False,
    smooth_camera=False,
    smooth_method="kalman",  # or simple
    **kwargs,
):
    """

    Args:
        rgb (numpy.ndarray): Shape of (t, h, w, 3), range [0, 1]
        disparity (numpy.ndarray): Shape of (t, h, w), range [0, 1]
        raymap (numpy.ndarray): Shape of (t, 6, h // 8, w // 8)
        ray_o_scale_inv (float, optional): A `ray_o` scale constant. Defaults to 10.
    """
    rgb = np.clip(rgb, 0, 1) * 255

    pointmap_dict = postprocess_pointmap(
        disparity,
        raymap,
        vae_downsample_scale,
        camera_pose,
        ray_o_scale_inv=ray_o_scale_inv,
        smooth_camera=smooth_camera,
        smooth_method=smooth_method,
        **kwargs,
    )

    save_ply(
        pointmap_dict["pointmap"],
        rgb,
        save_file,
        mask=(pointmap_dict["depth"] < max_depth),
    )

    if save_full_pcd_videos:
        pcd_dict = {
            "points": pointmap_dict["pointmap"],
            "colors": rgb,
            "intrinsics": pointmap_dict["intrinsics"],
            "poses": pointmap_dict["camera_pose"],
            "depths": pointmap_dict["depth"],
        }
        np.save(save_file.replace(".ply", "_pcd.npy"), pcd_dict)

    return pointmap_dict


def raymap_to_poses(
    raymap, camera_pose=None, ray_o_scale_inv=1.0, return_intrinsics=True
):
    ts = raymap.shape[0]
    if (not return_intrinsics) and (camera_pose is not None):
        return camera_pose, None, None

    raymap[:, 3:] = signed_log1p_inverse(raymap[:, 3:])

    # Extract ray origins and directions
    ray_o = (
        rearrange(raymap[:, 3:], "t c h w -> t h w c") * ray_o_scale_inv
    )  # [T, H, W, C]
    ray_d = rearrange(raymap[:, :3], "t c h w -> t h w c")  # [T, H, W, C]

    # Compute orientation and directions
    orient = ray_o.reshape(ts, -1, 3).mean(axis=1)  # T, 3
    image_orient = (ray_o + ray_d).reshape(ts, -1, 3).mean(axis=1)  # T, 3
    Focal = np.linalg.norm(image_orient - orient, axis=-1)  # T,
    Z_Dir = image_orient - orient  # T, 3

    # Compute the width (W) and field of view (FoV_x)
    W_Left = ray_d[:, :, :1, :].reshape(ts, -1, 3).mean(axis=1)
    W_Right = ray_d[:, :, -1:, :].reshape(ts, -1, 3).mean(axis=1)
    W = W_Right - W_Left
    W_real = (
        np.linalg.norm(np.cross(W, Z_Dir), axis=-1)
        / (raymap.shape[-1] - 1)
        * raymap.shape[-1]
    )
    Fov_x = np.arctan(W_real / (2 * Focal))

    # Compute the height (H) and field of view (FoV_y)
    H_Up = ray_d[:, :1, :, :].reshape(ts, -1, 3).mean(axis=1)
    H_Down = ray_d[:, -1:, :, :].reshape(ts, -1, 3).mean(axis=1)
    H = H_Up - H_Down
    H_real = (
        np.linalg.norm(np.cross(H, Z_Dir), axis=-1)
        / (raymap.shape[-2] - 1)
        * raymap.shape[-2]
    )
    Fov_y = np.arctan(H_real / (2 * Focal))

    # Compute X, Y, and Z directions for the camera
    X_Dir = W_Right - W_Left
    Y_Dir = np.cross(Z_Dir, X_Dir)
    X_Dir = np.cross(Y_Dir, Z_Dir)

    X_Dir /= np.linalg.norm(X_Dir, axis=-1, keepdims=True)
    Y_Dir /= np.linalg.norm(Y_Dir, axis=-1, keepdims=True)
    Z_Dir /= np.linalg.norm(Z_Dir, axis=-1, keepdims=True)

    # Create the camera-to-world (camera_pose) transformation matrix
    if camera_pose is None:
        camera_pose = np.zeros((ts, 4, 4))
        camera_pose[:, :3, 0] = X_Dir
        camera_pose[:, :3, 1] = Y_Dir
        camera_pose[:, :3, 2] = Z_Dir
        camera_pose[:, :3, 3] = orient
        camera_pose[:, 3, 3] = 1.0

    return camera_pose, Fov_x, Fov_y


def postprocess_pointmap(
    disparity,
    raymap,
    vae_downsample_scale=8,
    camera_pose=None,
    focal=None,
    ray_o_scale_inv=1.0,
    smooth_camera=False,
    smooth_method="simple",
    **kwargs,
):
    """

    Args:
        disparity (numpy.ndarray): Shape of (t, h, w), range [0, 1]
        raymap (numpy.ndarray): Shape of (t, 6, h // 8, w // 8)
        ray_o_scale_inv (float, optional): A `ray_o` scale constant. Defaults to 10.
    """
    depth = np.clip(1.0 / np.clip(disparity, 1e-3, 1), 0, 1e8)

    camera_pose, fov_x, fov_y = raymap_to_poses(
        raymap,
        camera_pose=camera_pose,
        ray_o_scale_inv=ray_o_scale_inv,
        return_intrinsics=(focal is not None),
    )
    if focal is None:
        focal = fov_to_focal(
            fov_x,
            fov_y,
            int(raymap.shape[2] * vae_downsample_scale),
            int(raymap.shape[3] * vae_downsample_scale),
        )

    if smooth_camera:
        # Check if sequence is static
        is_static, trans_diff, rot_diff = detect_static_sequence(camera_pose)

        if is_static:
            print(
                f"Detected static/near-static sequence (trans_diff={trans_diff:.6f}, rot_diff={rot_diff:.6f})"
            )
            # Apply stronger smoothing for static sequences
            camera_pose = adaptive_pose_smoothing(camera_pose, trans_diff, rot_diff)
        else:
            if smooth_method == "simple":
                camera_pose = smooth_poses(
                    camera_pose, window_size=5, method="gaussian"
                )
            elif smooth_method == "kalman":
                camera_pose = smooth_trajectory(camera_pose, window_size=5)

    ray_o, ray_d, intrinsics = get_rays(
        camera_pose,
        int(raymap.shape[2] * vae_downsample_scale),
        int(raymap.shape[3] * vae_downsample_scale),
        focal,
    )

    pointmap = depth[..., None] * ray_d + ray_o

    return {
        "pointmap": pointmap,
        "camera_pose": camera_pose,
        "intrinsics": intrinsics,
        "ray_o": ray_o,
        "ray_d": ray_d,
        "depth": depth,
    }


def detect_static_sequence(poses, threshold=0.01):
    """Detect if the camera sequence is static based on pose differences."""
    translations = poses[:, :3, 3]
    rotations = poses[:, :3, :3]

    # Compute translation differences
    trans_diff = np.linalg.norm(translations[1:] - translations[:-1], axis=1).mean()

    # Compute rotation differences (using matrix frobenius norm)
    rot_diff = np.linalg.norm(rotations[1:] - rotations[:-1], axis=(1, 2)).mean()

    return trans_diff < threshold and rot_diff < threshold, trans_diff, rot_diff


def adaptive_pose_smoothing(poses, trans_diff, rot_diff, base_window=5):
    """Apply adaptive smoothing based on motion magnitude."""
    # Increase window size for low motion sequences
    motion_magnitude = trans_diff + rot_diff
    adaptive_window = min(
        41, max(base_window, int(base_window * (0.1 / max(motion_magnitude, 1e-6))))
    )

    # Apply stronger smoothing for low motion
    poses_smooth = smooth_poses(poses, window_size=adaptive_window, method="gaussian")
    return poses_smooth


def get_pixel(H, W):
    # get 2D pixels (u, v) for image_a in cam_a pixel space
    u_a, v_a = np.meshgrid(np.arange(W), np.arange(H))
    # u_a = np.flip(u_a, axis=1)
    # v_a = np.flip(v_a, axis=0)
    pixels_a = np.stack(
        [u_a.flatten() + 0.5, v_a.flatten() + 0.5, np.ones_like(u_a.flatten())], axis=0
    )

    return pixels_a


def project(depth, intrinsic, pose):
    H, W = depth.shape
    pixel = get_pixel(H, W).astype(np.float32)
    points = (np.linalg.inv(intrinsic) @ pixel) * depth.reshape(-1)
    points = pose[:3, :4] @ np.concatenate(
        [points, np.ones((1, points.shape[1]))], axis=0
    )

    points = points.T.reshape(H, W, 3)

    return points


def depth_edge(
    depth: torch.Tensor,
    atol: float = None,
    rtol: float = None,
    kernel_size: int = 3,
    mask: Optional[torch.Tensor] = None,
) -> torch.BoolTensor:
    """
    Compute the edge mask of a depth map. The edge is defined as the pixels whose neighbors have a large difference in depth.

    Args:
        depth (torch.Tensor): shape (..., height, width), linear depth map
        atol (float): absolute tolerance
        rtol (float): relative tolerance

    Returns:
        edge (torch.Tensor): shape (..., height, width) of dtype torch.bool
    """
    is_numpy = isinstance(depth, np.ndarray)
    if is_numpy:
        depth = torch.from_numpy(depth)
    if isinstance(mask, np.ndarray):
        mask = torch.from_numpy(mask)

    shape = depth.shape
    depth = depth.reshape(-1, 1, *shape[-2:])
    if mask is not None:
        mask = mask.reshape(-1, 1, *shape[-2:])

    if mask is None:
        diff = F.max_pool2d(
            depth, kernel_size, stride=1, padding=kernel_size // 2
        ) + F.max_pool2d(-depth, kernel_size, stride=1, padding=kernel_size // 2)
    else:
        diff = F.max_pool2d(
            torch.where(mask, depth, -torch.inf),
            kernel_size,
            stride=1,
            padding=kernel_size // 2,
        ) + F.max_pool2d(
            torch.where(mask, -depth, -torch.inf),
            kernel_size,
            stride=1,
            padding=kernel_size // 2,
        )

    edge = torch.zeros_like(depth, dtype=torch.bool)
    if atol is not None:
        edge |= diff > atol
    if rtol is not None:
        edge |= (diff / depth).nan_to_num_() > rtol
    edge = edge.reshape(*shape)

    if is_numpy:
        return edge.numpy()
    return edge


@torch.jit.script
def align_rigid(
    p,
    q,
    weights,
):
    """Compute a rigid transformation that, when applied to p, minimizes the weighted
    squared distance between transformed points in p and points in q. See "Least-Squares
    Rigid Motion Using SVD" by Olga Sorkine-Hornung and Michael Rabinovich for more
    details (https://igl.ethz.ch/projects/ARAP/svd_rot.pdf).
    """

    device = p.device
    dtype = p.dtype
    batch, _, _ = p.shape

    # 1. Compute the centroids of both point sets.
    weights_normalized = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
    p_centroid = (weights_normalized[..., None] * p).sum(dim=-2)
    q_centroid = (weights_normalized[..., None] * q).sum(dim=-2)

    # 2. Compute the centered vectors.
    p_centered = p - p_centroid[..., None, :]
    q_centered = q - q_centroid[..., None, :]

    # 3. Compute the 3x3 covariance matrix.
    covariance = (q_centered * weights[..., None]).transpose(-1, -2) @ p_centered

    # 4. Compute the singular value decomposition and then the rotation.
    u, _, vt = torch.linalg.svd(covariance)
    s = torch.eye(3, dtype=dtype, device=device)
    s = s.expand((batch, 3, 3)).contiguous()
    s[..., 2, 2] = (u.det() * vt.det()).sign()
    rotation = u @ s @ vt

    # 5. Compute the optimal scale
    scale = (
        (torch.einsum("b i j, b k j -> b k i", rotation, p_centered) * q_centered).sum(
            -1
        )
        * weights
    ).sum(-1) / ((p_centered**2).sum(-1) * weights).sum(-1)
    # scale = (torch.einsum("b i j, b k j -> b k i", rotation, p_centered) * q_centered).sum([-1, -2]) / (p_centered**2).sum([-1, -2])

    # 6. Compute the optimal translation.
    translation = q_centroid - torch.einsum(
        "b i j, b j -> b i", rotation, p_centroid * scale[:, None]
    )

    return rotation, translation, scale


def align_camera_extrinsics(
    cameras_src: torch.Tensor,  # Bx3x4 tensor representing [R | t]
    cameras_tgt: torch.Tensor,  # Bx3x4 tensor representing [R | t]
    estimate_scale: bool = True,
    eps: float = 1e-9,
):
    """
    Align the source camera extrinsics to the target camera extrinsics.
    NOTE Assume OPENCV convention

    Args:
        cameras_src (torch.Tensor): Bx3x4 tensor representing [R | t] for source cameras.
        cameras_tgt (torch.Tensor): Bx3x4 tensor representing [R | t] for target cameras.
        estimate_scale (bool, optional): Whether to estimate the scale factor. Default is True.
        eps (float, optional): Small value to avoid division by zero. Default is 1e-9.

    Returns:
        align_t_R (torch.Tensor): 1x3x3 rotation matrix for alignment.
        align_t_T (torch.Tensor): 1x3 translation vector for alignment.
        align_t_s (float): Scaling factor for alignment.
    """

    R_src = cameras_src[:, :, :3]  # Extracting the rotation matrices from [R | t]
    R_tgt = cameras_tgt[:, :, :3]  # Extracting the rotation matrices from [R | t]

    RRcov = torch.bmm(R_tgt.transpose(2, 1), R_src).mean(0)
    U, _, V = torch.svd(RRcov)
    align_t_R = V @ U.t()

    T_src = cameras_src[:, :, 3]  # Extracting the translation vectors from [R | t]
    T_tgt = cameras_tgt[:, :, 3]  # Extracting the translation vectors from [R | t]

    A = torch.bmm(T_src[:, None], R_src)[:, 0]
    B = torch.bmm(T_tgt[:, None], R_src)[:, 0]

    Amu = A.mean(0, keepdim=True)
    Bmu = B.mean(0, keepdim=True)

    if estimate_scale and A.shape[0] > 1:
        # get the scaling component by matching covariances
        # of centered A and centered B
        Ac = A - Amu
        Bc = B - Bmu
        align_t_s = (Ac * Bc).mean() / (Ac**2).mean().clamp(eps)
    else:
        # set the scale to identity
        align_t_s = 1.0

    # get the translation as the difference between the means of A and B
    align_t_T = Bmu - align_t_s * Amu

    align_t_R = align_t_R[None]
    return align_t_R, align_t_T, align_t_s


def apply_transformation(
    cameras_src: torch.Tensor,  # Bx3x4 tensor representing [R | t]
    align_t_R: torch.Tensor,  # 1x3x3 rotation matrix
    align_t_T: torch.Tensor,  # 1x3 translation vector
    align_t_s: float,  # Scaling factor
    return_extri: bool = True,
) -> torch.Tensor:
    """
    Align and transform the source cameras using the provided rotation, translation, and scaling factors.
    NOTE Assume OPENCV convention

    Args:
        cameras_src (torch.Tensor): Bx3x4 tensor representing [R | t] for source cameras.
        align_t_R (torch.Tensor): 1x3x3 rotation matrix for alignment.
        align_t_T (torch.Tensor): 1x3 translation vector for alignment.
        align_t_s (float): Scaling factor for alignment.

    Returns:
        aligned_R (torch.Tensor): Bx3x3 tensor representing the aligned rotation matrices.
        aligned_T (torch.Tensor): Bx3 tensor representing the aligned translation vectors.
    """

    R_src = cameras_src[:, :, :3]
    T_src = cameras_src[:, :, 3]

    aligned_R = torch.bmm(R_src, align_t_R.expand(R_src.shape[0], 3, 3))

    # Apply the translation alignment to the source translations
    align_t_T_expanded = align_t_T[..., None].repeat(R_src.shape[0], 1, 1)
    transformed_T = torch.bmm(R_src, align_t_T_expanded)[..., 0]
    aligned_T = transformed_T + T_src * align_t_s

    if return_extri:
        extri = torch.cat([aligned_R, aligned_T.unsqueeze(-1)], dim=-1)
        return extri

    return aligned_R, aligned_T


def slerp(q1, q2, t):
    """Spherical Linear Interpolation between quaternions.
    Args:
        q1: (4,) first quaternion
        q2: (4,) second quaternion
        t: float between 0 and 1
    Returns:
        (4,) interpolated quaternion
    """
    # Compute the cosine of the angle between the two vectors
    dot = np.sum(q1 * q2)

    # If the dot product is negative, slerp won't take the shorter path
    # Fix by negating one of the input quaternions
    if dot < 0.0:
        q2 = -q2
        dot = -dot

    # Threshold for using linear interpolation instead of spherical
    DOT_THRESHOLD = 0.9995
    if dot > DOT_THRESHOLD:
        # If the inputs are too close for comfort, linearly interpolate
        # and normalize the result
        result = q1 + t * (q2 - q1)
        return result / np.linalg.norm(result)

    # Compute the angle between the quaternions
    theta_0 = np.arccos(dot)
    sin_theta_0 = np.sin(theta_0)

    # Compute interpolation factors
    theta = theta_0 * t
    sin_theta = np.sin(theta)

    s0 = np.cos(theta) - dot * sin_theta / sin_theta_0
    s1 = sin_theta / sin_theta_0

    return (s0 * q1) + (s1 * q2)


def interpolate_poses(pose1, pose2, weight):
    """Interpolate between two camera poses with weight.
    Args:
        pose1: (4, 4) first camera pose
        pose2: (4, 4) second camera pose
        weight: float between 0 and 1, weight for pose1 (1-weight for pose2)
    Returns:
        (4, 4) interpolated pose
    """
    from scipy.spatial.transform import Rotation as R

    # Extract rotations and translations
    R1 = R.from_matrix(pose1[:3, :3])
    R2 = R.from_matrix(pose2[:3, :3])
    t1 = pose1[:3, 3]
    t2 = pose2[:3, 3]

    # Get quaternions
    q1 = R1.as_quat()
    q2 = R2.as_quat()

    # Interpolate rotation using our slerp implementation
    q_interp = slerp(q1, q2, 1 - weight)  # 1-weight because weight is for pose1
    R_interp = R.from_quat(q_interp)

    # Linear interpolation for translation
    t_interp = weight * t1 + (1 - weight) * t2

    # Construct interpolated pose
    pose_interp = np.eye(4)
    pose_interp[:3, :3] = R_interp.as_matrix()
    pose_interp[:3, 3] = t_interp

    return pose_interp


def smooth_poses(poses, window_size=5, method="gaussian"):
    """Smooth camera poses temporally.
    Args:
        poses: (N, 4, 4) camera poses
        window_size: int, must be odd number
        method: str, 'gaussian' or 'savgol' or 'ma'
    Returns:
        (N, 4, 4) smoothed poses
    """
    from scipy.ndimage import gaussian_filter1d
    from scipy.signal import savgol_filter
    from scipy.spatial.transform import Rotation as R

    assert window_size % 2 == 1, "window_size must be odd"
    N = poses.shape[0]
    smoothed = np.zeros_like(poses)

    # Extract translations and quaternions
    translations = poses[:, :3, 3]
    rotations = R.from_matrix(poses[:, :3, :3])
    quats = rotations.as_quat()  # (N, 4)

    # Ensure consistent quaternion signs to prevent interpolation artifacts
    for i in range(1, N):
        if np.dot(quats[i], quats[i - 1]) < 0:
            quats[i] = -quats[i]

    # Smooth translations
    if method == "gaussian":
        sigma = window_size / 6.0  # approximately 99.7% of the weight within the window
        smoothed_trans = gaussian_filter1d(translations, sigma, axis=0, mode="nearest")
        smoothed_quats = gaussian_filter1d(quats, sigma, axis=0, mode="nearest")
    elif method == "savgol":
        # Savitzky-Golay filter: polynomial fitting
        poly_order = min(window_size - 1, 3)
        smoothed_trans = savgol_filter(
            translations, window_size, poly_order, axis=0, mode="nearest"
        )
        smoothed_quats = savgol_filter(
            quats, window_size, poly_order, axis=0, mode="nearest"
        )
    elif method == "ma":
        # Simple moving average
        kernel = np.ones(window_size) / window_size
        smoothed_trans = np.array(
            [np.convolve(translations[:, i], kernel, mode="same") for i in range(3)]
        ).T
        smoothed_quats = np.array(
            [np.convolve(quats[:, i], kernel, mode="same") for i in range(4)]
        ).T

    # Normalize quaternions
    smoothed_quats /= np.linalg.norm(smoothed_quats, axis=1, keepdims=True)

    # Reconstruct poses
    smoothed_rots = R.from_quat(smoothed_quats).as_matrix()

    for i in range(N):
        smoothed[i] = np.eye(4)
        smoothed[i, :3, :3] = smoothed_rots[i]
        smoothed[i, :3, 3] = smoothed_trans[i]

    return smoothed


def smooth_trajectory(poses, window_size=5):
    """Smooth camera trajectory using Kalman filter.
    Args:
        poses: (N, 4, 4) camera poses
        window_size: int, window size for initial smoothing
    Returns:
        (N, 4, 4) smoothed poses
    """
    from filterpy.kalman import KalmanFilter
    from scipy.spatial.transform import Rotation as R

    N = poses.shape[0]

    # Initialize Kalman filter for position and velocity
    kf = KalmanFilter(dim_x=6, dim_z=3)  # 3D position and velocity
    dt = 1.0  # assume uniform time steps

    # State transition matrix
    kf.F = np.array(
        [
            [1, 0, 0, dt, 0, 0],
            [0, 1, 0, 0, dt, 0],
            [0, 0, 1, 0, 0, dt],
            [0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 1],
        ]
    )

    # Measurement matrix
    kf.H = np.array([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0]])

    # Measurement noise
    kf.R *= 0.1

    # Process noise
    kf.Q *= 0.1

    # Initial state uncertainty
    kf.P *= 1.0

    # Extract translations and rotations
    translations = poses[:, :3, 3]
    rotations = R.from_matrix(poses[:, :3, :3])
    quats = rotations.as_quat()

    # First pass: simple smoothing for initial estimates
    smoothed = smooth_poses(poses, window_size, method="gaussian")
    smooth_trans = smoothed[:, :3, 3]

    # Second pass: Kalman filter for trajectory
    filtered_trans = np.zeros_like(translations)
    kf.x = np.zeros(6)
    kf.x[:3] = smooth_trans[0]

    filtered_trans[0] = smooth_trans[0]

    # Forward pass
    for i in range(1, N):
        kf.predict()
        kf.update(smooth_trans[i])
        filtered_trans[i] = kf.x[:3]

    # Backward smoothing for rotations using SLERP
    window_half = window_size // 2
    smoothed_quats = np.zeros_like(quats)

    for i in range(N):
        start_idx = max(0, i - window_half)
        end_idx = min(N, i + window_half + 1)
        weights = np.exp(
            -0.5 * ((np.arange(start_idx, end_idx) - i) / (window_half / 2)) ** 2
        )
        weights /= weights.sum()

        # Weighted average of nearby quaternions
        avg_quat = np.zeros(4)
        for j, w in zip(range(start_idx, end_idx), weights):
            if np.dot(quats[j], quats[i]) < 0:
                avg_quat += w * -quats[j]
            else:
                avg_quat += w * quats[j]
        smoothed_quats[i] = avg_quat / np.linalg.norm(avg_quat)

    # Reconstruct final smoothed poses
    final_smoothed = np.zeros_like(poses)
    smoothed_rots = R.from_quat(smoothed_quats).as_matrix()

    for i in range(N):
        final_smoothed[i] = np.eye(4)
        final_smoothed[i, :3, :3] = smoothed_rots[i]
        final_smoothed[i, :3, 3] = filtered_trans[i]

    return final_smoothed


def compute_scale(prediction, target, mask):
    if isinstance(prediction, np.ndarray):
        prediction = torch.from_numpy(prediction).float()
    if isinstance(target, np.ndarray):
        target = torch.from_numpy(target).float()
    if isinstance(mask, np.ndarray):
        mask = torch.from_numpy(mask).bool()

    numerator = torch.sum(mask * prediction * target, (1, 2))
    denominator = torch.sum(mask * prediction * prediction, (1, 2))

    scale = torch.zeros_like(numerator)

    valid = (denominator != 0).nonzero()

    scale[valid] = numerator[valid] / denominator[valid]

    return scale.item()


def get_raymap_from_camera_parameters(
    intrinsic,
    camera_pose,
    H,
    W,
    vae_downsample=8,
    align_corners=True,
):
    def get_raymap_from_trans2d(intrinsic, H, W):
        fu = intrinsic[:, 0, 0].unsqueeze(-1).unsqueeze(-1)
        fv = intrinsic[:, 1, 1].unsqueeze(-1).unsqueeze(-1)
        cu = intrinsic[:, 0, 2].unsqueeze(-1).unsqueeze(-1)
        cv = intrinsic[:, 1, 2].unsqueeze(-1).unsqueeze(-1)

        u, v = torch.meshgrid(torch.arange(W), torch.arange(H), indexing="xy")
        u = u.unsqueeze(0).repeat(intrinsic.shape[0], 1, 1).to(intrinsic.device)
        v = v.unsqueeze(0).repeat(intrinsic.shape[0], 1, 1).to(intrinsic.device)

        z_cam = torch.ones_like(u).to(intrinsic.device)
        x_cam = (u - cu) / fu
        y_cam = (v - cv) / fv
        addition_dim = torch.ones_like(u).to(intrinsic.device)
        return torch.stack((x_cam, y_cam, z_cam, addition_dim), dim=-1)

    raymap_cam = get_raymap_from_trans2d(intrinsic, H, W).to(camera_pose.device)

    T, raymap_cam_h, raymap_cam_w, _ = raymap_cam.shape
    raymap_cam = rearrange(raymap_cam, "t h w c -> t c (h w)")

    _camera_pose = camera_pose.clone()
    _camera_pose[:, :3, 3] = 0.0
    raymap_world = torch.bmm(_camera_pose, raymap_cam)
    raymap_world = rearrange(
        raymap_world, "t c (h w) -> t c h w", h=raymap_cam_h, w=raymap_cam_w
    )

    if vae_downsample != 1:
        raymap_world = F.interpolate(
            raymap_world,
            scale_factor=1 / vae_downsample,
            mode="bilinear",
            align_corners=align_corners,
        )
    raymap_world = raymap_world[:, :3]
    ray_o = torch.ones_like(raymap_world).to(raymap_world.device) * camera_pose[
        :, :3, 3
    ].unsqueeze(-1).unsqueeze(-1)

    raymap_world = torch.cat([raymap_world, ray_o], dim=1)
    return raymap_world


def camera_pose_to_raymap(
    camera_pose,
    intrinsic,
    ray_o_scale_factor: float = 10.0,
    dmax: float = 1.0,
    H: int = 480,
    W: int = 720,
    vae_downsample: int = 8,
    align_corners: bool = False,
) -> np.ndarray:
    """
    Convert camera pose to raymap.

    Args:
        camera_pose: (N, 4, 4) camera poses
        intrinsic: (N, 3, 3) intrinsics
        ray_o_scale_factor: A constant scale factor for ray_o to avoid too large translation values.
            Default to 10.0. If you use pre-trained AetherV1 model, you should always set it to 10.0.
        dmax: A constant scale factor for ray_d to avoid too large translation values.
            It should be equal to the maximum disparity value (before sqrt) of the sequence
            if you have ground truth disparity. Default to 1.0.
    Returns:
        (N, 6, H, W) raymap
    """
    is_numpy = isinstance(camera_pose, np.ndarray)
    if is_numpy:
        camera_pose = torch.from_numpy(camera_pose).float()
        intrinsic = torch.from_numpy(intrinsic).float()
    scale_factor = 1.0 / dmax
    camera_pose[:, :3, 3] = signed_log1p(
        camera_pose[:, :3, 3] / scale_factor * ray_o_scale_factor
    )
    raymap = get_raymap_from_camera_parameters(
        intrinsic,
        camera_pose,
        H,
        W,
        vae_downsample,
        align_corners,
    )
    if is_numpy:
        raymap = raymap.cpu().numpy()
    return raymap


def depth_to_disparity(depth, sqrt_disparity=True):
    """Convert depth to disparity.

    Args:
        depth: (N, H, W) depth map
        sqrt_disparity (bool, optional): Whether to take the square root of the disparity.
            Defaults to True.
    Returns:
        (N, H, W) disparity map
    """
    is_numpy = isinstance(depth, np.ndarray)
    if is_numpy:
        depth = torch.from_numpy(depth).float()
    disparity = 1.0 / depth
    valid_disparity = disparity[depth > 1e-6]
    dmax = valid_disparity.max()
    disparity = torch.clamp(disparity / dmax, min=0.0, max=1.0)

    if sqrt_disparity:
        disparity = torch.sqrt(disparity)

    if is_numpy:
        disparity = disparity.cpu().numpy()
    return disparity, dmax