3D-вращения и пространственные преобразования с помощью RoMa - UPROGER

Затрудняетесь с кватернионами, векторами вращения, правилами правой руки и прочими вещами? Попробуйте RoMa: простую в использовании, стабильную и эффективную библиотеку для работы с вращениями и пространственными преобразованиями в PyTorch.

Немного контекста

Пространственные преобразования необходимы в физике, инженерии и компьютерном зрении. Ученые веками изучали 3D-геометрию и создали отличные математические инструменты для работы с такими преобразованиями, особенно с жесткими движениями и вращениями. Однако правильная реализация этих инструментов не является тривиальной.

Действительно, для разных задач часто приходится использовать разные соглашения или представления. Например, кватернионы отлично подходят для составления вращений; матрицы вращения идеальны для преобразования координат большого количества точек; векторы вращения компактны и полезны для представления небольших вращений, а углы Эйлера плохо подходят практически для всего, кроме пользовательского ввода.

Кроме того, превращение математики в код приводит к возникновению численных проблем, которые трудно заметить и которые вызывают малозаметные ошибки. Остается надеяться, что эти проблемы можно устранить, выбрав правильные алгоритмы или обработав особые случаи. Например, угол между двумя матрицами вращения может быть вычислен по разным формулам, которые математически эквивалентны, но дают радикально разные численные ошибки в 32-битной системе с плавающей точкой:

Представляем RoMa

RoMa призвана преодолеть эти препятствия. RoMa – это библиотека Python, совместимая с PyTorch версии 1.6 и выше. Она предоставляет простой в использовании, стабильный и эффективный набор инструментов для работы с вращениями, а также более общими пространственными преобразованиями.

Преобразования между представлениями вращения

RoMa предоставляет дифференцируемые процедуры для преобразования между различными представлениями вращения. Например, можно взять вектор вращения в 3D и преобразовать его в единичный кватернион следующим образом с помощью RoMa:

import torch, roma
rotvec = torch.randn(3) # 3D rotation vector
q = roma.rotvec_to_unitquat(rotvec) # unit quaternion, represented by a 4D tensor

Пакетированные данные

Для удобства функции в RoMa поддерживают произвольное число размерностей партии. Например, можно взять партию из 2×5 случайных векторов 3D вращения – представленных тензором 2x5x3 – и преобразовать ее в партию единичных кватернионов – представленных тензором 2x5x4 – используя тот же синтаксис, что и выше:

import torch, roma
rotvec = torch.randn(2,5,3)
q = roma.rotvec_to_unitquat(rotvec)

Регрессирующие вращения

Регрессия вращений с помощью нейронной сети нетривиальна, поскольку классические нейронные архитектуры выдают выходные сигналы, лежащие в евклидовом пространстве. В RoMa реализованы различные дифференцируемые функции для отображения таких выходов в пространство вращений, например, путем специальной ортонормировки Прокруста произвольной матрицы:

import torch, roma
M = my_fancy_neural_network(some_input) # Method returning an arbitrary 3x3 matrix
R = roma.special_procrustes(M) # Orthonormalizing M into a rotation matrix
assert roma.is_rotation_matrix(R, epsilon=1e-5)

Жесткие преобразования

Чтобы код был более читабельным, RoMa также включает утилиты для работы с нелинейными пространственными преобразованиями и, в частности, с жесткими движениями:

import torch, roma
# Rigid transformation parameterized by a rotation matrix and a translation vector
T1 = roma.Rigid(linear=roma.random_rotmat(), translation=torch.randn(3))
T2 = roma.Rigid(linear=roma.random_rotmat(), translation=torch.randn(3))

# Inverting and composing transformations
T = (T1.inverse() @ T2)
# Normalization to ensure that T is actually a rigid transformation.
T = T.normalize()
# Direct access to the translation part
T.translation += 0.5
# Transformation of points:
points = torch.randn(100,3)
# Adjusting the shape of T for proper broadcasting.
transformed_points = T[None].apply(points)
# Transformation of vectors:
vectors = torch.randn(10,20,3)
# Adjusting the shape of T for proper broadcasting.
transformed_vectors = T[None,None].linear_apply(vectors)
# Casting the transformation into an homogeneous 4x4 matrix.
M = T.to_homogeneous()

RoMa можно легко установить с помощью PIP, так что попробуйте: