Source code for blenderproc.python.utility.MathUtility

""" A collection of math functions. """

from typing import Union, List

import numpy as np
from mathutils import Matrix, Vector, Euler




def change_coordinate_frame_of_point(point: Union[np.ndarray, List[float], Vector],
                                     new_frame: List[str]) -> np.ndarray:
    """ Transforms the given point into another coordinate frame.

    Example: [1, 2, 3] and ["X", "-Z", "Y"] => [1, -3, 2]

    :param point: The point to convert in form of a np.ndarray, list or mathutils.Vector.
    :param new_frame: An array containing three elements, describing each axis of the new coordinate frame
                      based on the axes of the current frame. Available: ["X", "Y", "Z", "-X", "-Y", "-Z"].
    :return: The converted point also in form of a np.ndarray
    """
    assert len(new_frame) == 3, f"The specified coordinate frame has more or less than tree axes: {new_frame}"
    point = np.array(point)

    output = []
    for axis in new_frame:
        axis = axis.upper()

        if axis.endswith("X"):
            output.append(point[0])
        elif axis.endswith("Y"):
            output.append(point[1])
        elif axis.endswith("Z"):
            output.append(point[2])
        else:
            raise ValueError(f"Invalid axis: {axis}")

        if axis.startswith("-"):
            output[-1] *= -1

    return np.array(output)





def change_target_coordinate_frame_of_transformation_matrix(matrix: Union[np.ndarray, Matrix],
                                                            new_frame: List[str]) -> np.ndarray:
    """ Changes the coordinate frame the given transformation matrix is mapping to.

    Given a matrix $T_A^B$ that maps from A to B, this function can be used
    to change the axes of B into B' and therefore end up with $T_A^B'$.

    :param matrix: The matrix to convert in form of a np.ndarray or mathutils.Matrix
    :param new_frame: An array containing three elements, describing each axis of the new coordinate frame
                      based on the axes of the current frame. Available: ["X", "Y", "Z", "-X", "-Y", "-Z"].
    :return: The converted matrix is in form of a np.ndarray
    """
    tmat = MathUtility.build_coordinate_frame_changing_transformation_matrix(new_frame)

    # Apply transformation matrix
    output = np.matmul(tmat, matrix)
    return output





def change_source_coordinate_frame_of_transformation_matrix(matrix: Union[np.ndarray, Matrix],
                                                            new_frame: list) -> np.ndarray:
    """ Changes the coordinate frame the given transformation matrix is mapping from.

    Given a matrix $T_A^B$ that maps from A to B, this function can be used
    to change the axes of A into A' and therefore end up with $T_A'^B$.

    :param matrix: The matrix to convert in form of a np.ndarray or mathutils.Matrix
    :param new_frame: An array containing three elements, describing each axis of the new coordinate frame
                      based on the axes of the current frame. Available: ["X", "Y", "Z", "-X", "-Y", "-Z"].
    :return: The converted matrix is in form of a np.ndarray
    """
    tmat = MathUtility.build_coordinate_frame_changing_transformation_matrix(new_frame)
    tmat = np.linalg.inv(tmat)

    # Apply transformation matrix
    output = np.matmul(matrix, tmat)
    return output





def build_transformation_mat(translation: Union[np.ndarray, List[float], Vector],
                             rotation: Union[np.ndarray, List[List[float]], Matrix]) -> np.ndarray:
    """ Build a transformation matrix from translation and rotation parts.

    :param translation: A (3,) vector representing the translation part.
    :param rotation: A 3x3 rotation matrix or Euler angles of shape (3,).
    :return: The 4x4 transformation matrix.
    """
    translation = np.array(translation)
    rotation = np.array(rotation)

    mat = np.eye(4)
    if translation.shape[0] == 3:
        mat[:3, 3] = translation
    else:
        raise RuntimeError(f"Translation has invalid shape: {translation.shape}. Must be (3,) or (3,1) vector.")
    if rotation.shape == (3, 3):
        mat[:3, :3] = rotation
    elif rotation.shape[0] == 3:
        mat[:3, :3] = np.array(Euler(rotation).to_matrix())
    else:
        raise RuntimeError(f"Rotation has invalid shape: {rotation.shape}. Must be rotation matrix of shape "
                           f"(3,3) or Euler angles of shape (3,) or (3,1).")

    return mat





class MathUtility:
    """
    Math utility class
    """



    @staticmethod
    def build_coordinate_frame_changing_transformation_matrix(destination_frame: List[str]) -> np.ndarray:
        """ Builds a transformation matrix that switches the coordinate frame.

        :param destination_frame: An array containing three elements, describing each axis of the destination
                                  coordinate frame based on the axes of the source frame.
                                  Available: ["X", "Y", "Z", "-X", "-Y", "-Z"].
        :return: The transformation matrix
        """
        assert len(destination_frame) == 3, f"The specified coordinate frame has more or less than " \
                                            f"tree axes: {destination_frame}"

        # Build transformation matrix that maps the given matrix to the specified coordinate frame.
        tmat = np.zeros((4, 4))
        for i, axis in enumerate(destination_frame):
            axis = axis.upper()

            if axis.endswith("X"):
                tmat[i, 0] = 1
            elif axis.endswith("Y"):
                tmat[i, 1] = 1
            elif axis.endswith("Z"):
                tmat[i, 2] = 1
            else:
                raise Exception("Invalid axis: " + axis)

            if axis.startswith("-"):
                tmat[i] *= -1
        tmat[3, 3] = 1
        return tmat