PathPointsTransformer/transformations.py

from typing import Optional, Tuple

import numpy

from numpy.linalg import svd


def centroid(points: numpy.ndarray) -> numpy.ndarray:
    return points.mean(axis=0)


def calculate_rotation_matrix(actual_points: numpy.ndarray,
                              expected_points: numpy.ndarray,
                              *,
                              calculation_points_number: Optional[int] = None) -> numpy.ndarray:
    """
    Вычисление матрицы поворота, который необходим для приближения текущих точек к ожидаемым, на основе алгоритма Кабша

    :param actual_points: массив текущих точек
    :param expected_points: массив ожидаемых точек
    :param calculation_points_number: число используемых для поворота точек
    :return: Матрица поворотов

    """
    sliced_actual_points = actual_points[:calculation_points_number] \
        if isinstance(calculation_points_number, int) \
        else actual_points
    sliced_expected_points = expected_points[:calculation_points_number] \
        if isinstance(calculation_points_number, int) \
        else expected_points

    # Для применения алгоритма Кабша необходимо,
    # чтобы центроиды обрабатываемых массивов точек совпадали с началом координат
    centred_actual_points = sliced_actual_points - centroid(sliced_actual_points)
    centred_expected_points = sliced_expected_points - centroid(sliced_expected_points)

    # Вычисляем матрицу кросс-ковариаций
    cross_covariance_matrix = numpy.dot(numpy.transpose(centred_actual_points), centred_expected_points)

    # Применяем к полученной матрице кросс-вариаций сингулярное разложение
    left_singular_vectors_matrix, singular_matrix, right_singular_vectors_matrix = svd(cross_covariance_matrix)

    # Некая нормализация полученного разложения (сам не знаю что здесь происходит, поэтому не трогаю)
    if (numpy.linalg.det(left_singular_vectors_matrix) * numpy.linalg.det(right_singular_vectors_matrix)) < 0.0:
        singular_matrix[-1] = -singular_matrix[-1]
        left_singular_vectors_matrix[:, -1] = -left_singular_vectors_matrix[:, -1]

    return numpy.dot(left_singular_vectors_matrix, right_singular_vectors_matrix)


def convert_rotation_matrix_to_angles(rotation_matrix: numpy.ndarray) -> Tuple[float, float, float]:
    """
    Вычисление углов поворота модели на основе применяемой к ней матрицы поворота

    :param rotation_matrix: матрица поворота
    :return: Кортеж из трёх углов - угол поворота вокруг оси X, угол поворота вокруг оси Y, угол поворота вокруг оси Z

    """
    alpha = numpy.arccos(rotation_matrix[2, 2] / (1 - rotation_matrix[0, 2] ** 2) ** 0.5)
    beta = numpy.arcsin(rotation_matrix[0, 2])
    gamma = numpy.arccos(rotation_matrix[0, 0] / (1 - rotation_matrix[0, 2] ** 2) ** 0.5)
    return alpha, beta, gamma