tools.py

from typing import Tuple

import numpy as np

# ------------------------ General functions ------------------------
def e2p(u_e) -> np.ndarray:
    """
    Convert Euclidean to homogeneous coordinates
    :param u_e: d by n matrix; n Euclidean vectors of dimension d
    :return: d+1 by n matrix; n homogeneous vectors of dimension d+1
    """
    if len(u_e.shape) == 1:
        u_e = u_e.reshape(-1, 1)
    d, n = u_e.shape
    u_p = np.vstack((u_e, np.ones((1, n))))
    return u_p

def p2e(u_p) -> np.ndarray:
    """
    Convert homogeneous to Euclidean coordinates
    :param u_p: vector d+1 by 1; homogeneous vector of dimension d+1
    :return:
    """
    u_e = u_p[:-1] / u_p[-1]
    return u_e

def vlen(x) -> np.ndarray:
    """
    Compute Euclidean lengths of vectors
    :param x: d by n matrix; n vectors of dimension d
    :return: 1 by n row vector; Euclidean lengths of the vectors
    """
    return np.sqrt(np.sum(x ** 2, axis=0))

def sqc(x) -> np.ndarray:
    """
    Skew-symmetric matrix for cross-product
    :param x: vector 3×1
    :return: skew symmetric matrix (3×3) for cross product with x
    """
    x = x.reshape(-1)
    return np.array([
        [0, -x[2], x[1]],
        [x[2], 0, -x[0]],
        [-x[1], x[0], 0]
    ])


# ------------------------ Camera functions ------------------------
def PP2F(P1, P2) -> np.ndarray:
    """
    Compute the Fundamental matrix from two camera matrices
    :param P1: 3x4 camera matrix of the first camera
    :param P2: 3x4 camera matrix of the second camera
    :return: 3x3 Fundamental matrix
    """
    assert P1.shape == P2.shape == (3, 4), f"P1 shape is {P1.shape}, P2 shape is {P2.shape}, expected (3, 4)"
    Q1, q1 = P1[:, :-1], P1[:, -1]
    Q2, q2 = P2[:, :-1], P2[:, -1]
    Q1_Q2_inv = Q1 @ np.linalg.inv(Q2)
    F = Q1_Q2_inv.T @ sqc(q1 - Q1_Q2_inv @ q2)
    return F

def Pu2X(P1, P2, u1, u2) -> np.ndarray:
    """
    Triangulate the 3D point from two camera matrices and corresponding points.
    :param P1: 3x4 camera matrix of the first camera
    :param P2: 3x4 camera matrix of the second camera
    :param u1: 3xN homogeneous coordinates of points in the first image
    :param u2: 3xN homogeneous coordinates of points in the second image
    :return: 4xN 3D points in homogeneous coordinates (normalized)
    """
    assert P1.shape == P2.shape == (3, 4), f"P1 shape is {P1.shape}, P2 shape is {P2.shape}, expected (3, 4)"
    assert u1.shape[0] == u2.shape[0] == 3, f"u1p shape is {u1.shape}, u2p shape is {u2.shape}, expected (3, N)"

    n_points = u1.shape[1]
    X = np.zeros((4, n_points))  # To store the resulting 3D points in homogeneous coordinates

    # Normalize the points
    u1 /= u1[-1, :]
    u2 /= u2[-1, :]

    for i in range(n_points):
        D = np.vstack([
            u1[0, i] * P1[2, :] - P1[0, :],
            u1[1, i] * P1[2, :] - P1[1, :],
            u2[0, i] * P2[2, :] - P2[0, :],
            u2[1, i] * P2[2, :] - P2[1, :]
        ])

        # Conditioning matrix
        S = np.diag(1 / np.max(np.abs(D), axis=0))

        # Re-scale the problem and solve for X'
        _, _, Vt = np.linalg.svd(D @ S)
        X_prime = Vt[-1]  # Solution is the last row of V^T

        # Back-transform to get the original solution
        X[:, i] = S @ X_prime

    # Normalize the 3D points (make the last coordinate 1)
    X /= X[-1, :]

    return X

def Rt2E(R, t) -> np.ndarray:
    """
    Compute the Essential matrix from the rotation matrix and translation vector
    :param R: 3x3 rotation matrix
    :param t: 3x1 translation vector
    :return: 3x3 Essential matrix
    """
    assert R.shape == (3, 3), f"R shape is {R.shape}, expected (3, 3)"
    assert t.shape == (3, 1), f"t shape is {t.shape}, expected (3, 1)"

    return sqc(-t) @ R

def Eu2Rt(E, u1p, u2p) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute the rotation matrices and camera positions from the Essential matrix
    :param E: 3x3 Essential matrix
    :param u1p: 3x1 homogeneous coordinates of the point in the first image
    :param u2p: 3x1 homogeneous coordinates of the point in the second image
    :return R, t: Rotation matrix and translation vector
    """
    assert E.shape == (3, 3), f"E shape is {E.shape}, expected (3, 3)"
    assert u1p.shape[0] == u2p.shape[0] == 3, f"u1p shape is {u1p.shape}, expected (3, N)"

    # Matrix decomposition of E
    # Step 1: Compute SVD
    U, D, Vt = np.linalg.svd(E)

    # Step 2: Ensure that U and V are rotation matrices
    if np.linalg.det(U) < 0:
        U[:, -1] *= -1
    if np.linalg.det(Vt) < 0:
        Vt[-1, :] *= -1

    # Step 3: Compute the four possible rotation matrices
    P1 = np.eye(3, 4)
    for a in [-1, 1]:
        for b in [-1, 1]:
            W = np.array([
                [0, a, 0],
                [-a, 0, 0],
                [0, 0, 1]
            ])

            R = U @ W @ Vt
            t = (b * U[:, -1]).reshape(3, 1)

            # Check the triangulation
            P2 = np.hstack((R, t))
            X = Pu2X(P1, P2, u1p, u2p)

            # Keep R and t if the depth is positive for all points (chirality constraint)
            x = P1 @ X
            y = P2 @ X
            if np.all(x[2, :] > 0) and np.all(y[2, :] > 0):
                return R, t

    return [], np.zeros((3, 1))

def X_chirality_apical_angle(X, R1, t1, R2, t2, thr_a) -> np.ndarray:
    """
    Returns a boolean array. It indicates 3D points from X that lie in front of cameras P1 and P2.
    The angle between two vectors from X to each camera location must also be greater or equal to thr_a.
    :param X: 4xN 3D points in homogeneous coordinates
    :param R1: rotation matrix of the first camera (3x3)
    :param t1: translation vector of the first camera (3x1)
    :param R2: rotation matrix of the second camera (3x3)
    :param t2: translation vector of the second camera (3x1)
    :param thr_a: Threshold angle in degrees
    :return: Boolean array (1xN) indicating the points that satisfy the conditions
    """
    assert X.shape[0] == 4, f"X shape is {X.shape}, expected (4, N)"
    assert R1.shape == R2.shape == (3, 3), f"R1 shape is {R1.shape}, R2 shape is {R2.shape}, expected (3, 3)"
    if t1.shape != (3, 1):
        t1 = t1.reshape(3, 1)
    if t2.shape != (3, 1):
        t2 = t2.reshape(3, 1)

    # Camera positions
    C1 = -R1.T @ t1
    C2 = -R2.T @ t2

    # Vector from X to camera positions
    X_C1 = C1 - X[:3, :]
    X_C2 = C2 - X[:3, :]
    X_C1 /= vlen(X_C1)
    X_C2 /= vlen(X_C2)

    # Compute the apical angle
    apical_angle = np.arccos(np.sum(X_C1 * X_C2, axis=0)) * 180 / np.pi

    # Check the chirality and the apical angle
    return (X[2, :] > 0) & (apical_angle >= thr_a)


# ------------------------ Error function ------------------------
def err_F_sampson(F, u1, u2) -> np.ndarray:
    """
    Compute the Sampson error for a set of points
    :param F: Fundamental matrix
    :param u1: homogeneous coordinates of image 1
    :param u2: homogeneous coordinates of image 2
    :return: Sampson error
    """
    assert F.shape == (3, 3), f"F shape is {F.shape}, expected (3, 3)"
    assert u1.shape[0] == u2.shape[0] == 3, f"u1 shape is {u1.shape}, u2 shape is {u2.shape}, expected (3, N)"

    # # Compute the numerator (u2.T * F * u1)^2
    F_u1 = F @ u1  # 3xn matrix
    F_T_u2 = F.T @ u2  # 3xn matrix
    num = np.sum(u2 * F_u1, axis=0) ** 2

    # Compute the denominator (Fu1)^2 + (F^Tu2)^2
    denom = np.sum(F_u1[:2, :] ** 2, axis=0) + np.sum(F_T_u2[:2, :] ** 2, axis=0)

    # Sampson error (squared)
    errors = num / denom

    return errors

def u_correct_sampson(F, u1, u2) -> Tuple[np.ndarray, np.ndarray]:
    """
    Estimates sampson-corrected coordinates based on the given fundamental matrix F from image points in homogeneous coordinates
    :param F: Fundamental matrix
    :param u1: homogeneous coordinates of image 1
    :param u2: homogeneous coordinates of image 2
    :return nu1, nu2: Corrected points in homogeneous coordinates
    """
    assert F.shape == (3, 3), f"F shape is {F.shape}, expected (3, 3)"
    assert u1.shape[0] == u2.shape[0] == 3, f"u1 shape is {u1.shape}, u2 shape is {u2.shape}, expected (3, N)"

    S = np.array([[1, 0, 0], [0, 1, 0]])
    SF = S @ F
    SF_t = S @ F.T
    F_u1 = F @ u1  # 3xn matrix
    F_T_u2 = F.T @ u2  # 3xn matrix
    num = u2 * F_u1
    denom = np.sum(SF @ u1)**2 + np.sum(SF_t @ u2)**2

    correction = num / denom

    # Correct the points (update the first two rows)
    nu1 = u1 - (F_T_u2 * correction)
    nu2 = u2 - (F_u1 * correction)

    nu1[2, :] = 1
    nu2[2, :] = 1
    return nu1, nu2

def err_reproj(X, u, R, t, K):
    if X.shape[0] == 3:
        X = e2p(X)
    if u.shape[0] == 2:
        u = e2p(u)
    P = K @ np.hstack([R, t])
    X_proj = P @ X
    X_proj /= X_proj[-1, :]
    errors = np.linalg.norm(X_proj - u, axis=0)
    return errors


def get_colors(camera1, camera2, u1, u2):
    if u1.shape[0] == 3:
        u1 = p2e(u1)
    if u2.shape[0] == 3:
        u2 = p2e(u2)
    img1, img2 = camera1.image, camera2.image
    colors = np.zeros((3, u1.shape[1]))
    for i in range(u1.shape[1]):
        [x, y] = np.round(u1[:, i]).astype(int)
        [r, g, b] = img1[y, x]
        color1 = (r / 255.0, g / 255.0, b / 255.0)

        [x, y] = np.round(u2[:, i]).astype(int)
        [r, g, b] = img2[y, x]
        color2 = (r / 255.0, g / 255.0, b / 255.0)

        color = ((color1[0] + color2[0]) / 2,
                 (color1[1] + color2[1]) / 2,
                 (color1[2] + color2[2]) / 2)
        colors[:, i] = color

    return np.array(colors)