geometry.py

# Copyright 2021 Toyota Research Institute.  All rights reserved.
import logging

import cv2
import numpy as np
import torch
import torch.nn.functional as F

LOG = logging.getLogger(__name__)

PI = 3.14159265358979323846
EPS = 1e-7

def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
    """
    Returns torch.sqrt(torch.max(0, x))
    but with a zero subgradient where x is 0.
    """
    ret = torch.zeros_like(x)
    positive_mask = x > 0
    ret[positive_mask] = torch.sqrt(x[positive_mask])
    return ret

def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as rotation matrices to quaternions.

    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")

    batch_dim = matrix.shape[:-2]
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
        matrix.reshape(batch_dim + (9,)), dim=-1
    )

    q_abs = _sqrt_positive_part(
        torch.stack(
            [
                1.0 + m00 + m11 + m22,
                1.0 + m00 - m11 - m22,
                1.0 - m00 + m11 - m22,
                1.0 - m00 - m11 + m22,
            ],
            dim=-1,
        )
    )

    # we produce the desired quaternion multiplied by each of r, i, j, k
    quat_by_rijk = torch.stack(
        [
            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
        ],
        dim=-2,
    )

    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
    # the candidate won't be picked.
    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))

    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
    # forall i; we pick the best-conditioned one (with the largest denominator)

    return quat_candidates[
        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :  # pyre-ignore[16]
    ].reshape(batch_dim + (4,))

def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to rotation matrices.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))

def allocentric_to_egocentric(quat, proj_ctr, inv_intrinsics):
    """
    Parameters
    ----------
    quat: Tensor
        (N, 4). Batch of (allocentric) quaternions.

    proj_ctr: Tensor
        (N, 2). Projected centers. xy coordninates.

    inv_intrinsics: [type]
        (N, 3, 3). Inverted intrinsics.
    """
    R_obj_to_local = quaternion_to_matrix(quat)

    # ray == z-axis in local orientaion
    ray = unproject_points2d(proj_ctr, inv_intrinsics)
    z = ray / ray.norm(dim=1, keepdim=True)

    # gram-schmit process: local_y = global_y - global_y \dot local_z
    y = z.new_tensor([[0., 1., 0.]]) - z[:, 1:2] * z
    y = y / y.norm(dim=1, keepdim=True)
    x = torch.cross(y, z, dim=1)

    # local -> global
    R_local_to_global = torch.stack([x, y, z], dim=-1)

    # obj -> global
    R_obj_to_global = torch.bmm(R_local_to_global, R_obj_to_local)

    egocentric_quat = matrix_to_quaternion(R_obj_to_global)

    # Make sure it's unit norm.
    quat_norm = egocentric_quat.norm(dim=1, keepdim=True)
    if not torch.allclose(quat_norm, torch.as_tensor(1.), atol=1e-3):
        LOG.warning(
            f"Some of the input quaternions are not unit norm: min={quat_norm.min()}, max={quat_norm.max()}; therefore normalizing."
        )
        egocentric_quat = egocentric_quat / quat_norm.clamp(min=EPS)

    return egocentric_quat


def homogenize_points(xy):
    """
    Parameters
    ----------
    xy: Tensor
        xy coordinates. shape=(N, ..., 2)
        E.g., (N, 2) or (N, K, 2) or (N, H, W, 2)

    Returns
    -------
    Tensor:
        1. is appended to the last dimension. shape=(N, ..., 3)
        E.g, (N, 3) or (N, K, 3) or (N, H, W, 3).
    """
    # NOTE: this seems to work for arbitrary number of dimensions of input
    pad = torch.nn.ConstantPad1d(padding=(0, 1), value=1.)
    return pad(xy)


def project_points3d(Xw, K):
    _, C = Xw.shape
    assert C == 3
    uv, _ = cv2.projectPoints(
        Xw, np.zeros((3, 1), dtype=np.float32), np.zeros(3, dtype=np.float32), K, np.zeros(5, dtype=np.float32)
    )
    return uv.reshape(-1, 2)


def unproject_points2d(points2d, inv_K, scale=1.0):
    """
    Parameters
    ----------
    points2d: Tensor
        xy coordinates. shape=(N, ..., 2)
        E.g., (N, 2) or (N, K, 2) or (N, H, W, 2)

    inv_K: Tensor
        Inverted intrinsics; shape=(N, 3, 3)

    scale: float, default: 1.0
        Scaling factor.

    Returns
    -------
    Tensor:
        Unprojected 3D point. shape=(N, ..., 3)
        E.g., (N, 3) or (N, K, 3) or (N, H, W, 3)
    """
    points2d = homogenize_points(points2d)
    siz = points2d.size()
    points2d = points2d.view(-1, 3).unsqueeze(-1)  # (N, 3, 1)
    unprojected = torch.matmul(inv_K, points2d)  # (N, 3, 3) x (N, 3, 1) -> (N, 3, 1)
    unprojected = unprojected.view(siz)

    return unprojected * scale