GitHub Viewer
#include "tensor_math.hpp"
using namespace torch::indexing;
torch::Tensor quatToRotMat(const torch::Tensor &quat){
auto u = torch::unbind(torch::nn::functional::normalize(quat, torch::nn::functional::NormalizeFuncOptions().dim(-1)), -1);
torch::Tensor w = u[0];
torch::Tensor x = u[1];
torch::Tensor y = u[2];
torch::Tensor z = u[3];
return torch::stack({
torch::stack({
1.0 - 2.0 * (y.pow(2) + z.pow(2)),
2.0 * (x * y - w * z),
2.0 * (x * z + w * y)
}, -1),
torch::stack({
2.0 * (x * y + w * z),
1.0 - 2.0 * (x.pow(2) + z.pow(2)),
2.0 * (y * z - w * x)
}, -1),
torch::stack({
2.0 * (x * z - w * y),
2.0 * (y * z + w * x),
1.0 - 2.0 * (x.pow(2) + y.pow(2))
}, -1)
}, -2);
}
std::tuple autoScaleAndCenterPoses(const torch::Tensor &poses){
// Center at mean
torch::Tensor origins = poses.index({"...", Slice(None, 3), 3});
torch::Tensor center = torch::mean(origins, 0);
origins -= center;
// Scale
float f = 1.0f / torch::max(torch::abs(origins)).item();
origins *= f;
torch::Tensor transformedPoses = poses.clone();
transformedPoses.index_put_({"...", Slice(None, 3), 3}, origins);
return std::make_tuple(transformedPoses, center, f);
}
torch::Tensor rotationMatrix(const torch::Tensor &a, const torch::Tensor &b){
// Rotation matrix that rotates vector a to vector b
torch::Tensor a1 = a / a.norm();
torch::Tensor b1 = b / b.norm();
torch::Tensor v = torch::linalg_cross(a1, b1);
torch::Tensor c = torch::dot(a1, b1);
const float EPS = 1e-8;
if (c.item() < -1 + EPS){
torch::Tensor eps = (torch::rand(3) - 0.5f) * 0.01f;
return rotationMatrix(a1 + eps, b1);
}
torch::Tensor s = v.norm();
torch::Tensor skew = torch::zeros({3, 3}, torch::kFloat32);
skew[0][1] = -v[2];
skew[0][2] = v[1];
skew[1][0] = v[2];
skew[1][2] = -v[0];
skew[2][0] = -v[1];
skew[2][1] = v[0];
return torch::eye(3) + skew + torch::matmul(skew, skew * ((1 - c) / (s.pow(2) + EPS)));
}
torch::Tensor rodriguesToRotation(const torch::Tensor &rodrigues){
float theta = torch::linalg_vector_norm(rodrigues, 2, { -1 }, true, torch::kFloat32).item();
if (theta < FLOAT_EPS){
return torch::eye(3, torch::kFloat32);
}
torch::Tensor r = rodrigues / theta;
torch::Tensor ident = torch::eye(3, torch::kFloat32);
float a = r[0].item();
float b = r[1].item();
float c = r[2].item();
torch::Tensor rrT = torch::tensor({
{a * a, a * b, a * c},
{b * a, b * b, b * c},
{c * a, c * b, c * c}
}, torch::kFloat32);
torch::Tensor rCross = torch::tensor({
{0.0f, -c, b},
{c, 0.0f, -a},
{-b, a, 0.0f}
}, torch::kFloat32);
float cosTheta = std::cos(theta);
return cosTheta * ident + (1 - cosTheta) * rrT + std::sin(theta) * rCross;
}