367 lines
13 KiB
Python
367 lines
13 KiB
Python
import torch
|
|
import numpy as np
|
|
from torch.nn import functional as F
|
|
from einops.einops import rearrange
|
|
|
|
|
|
def cam2pixel(cam_coord, f, c):
|
|
x = cam_coord[:, 0] / cam_coord[:, 2] * f[0] + c[0]
|
|
y = cam_coord[:, 1] / cam_coord[:, 2] * f[1] + c[1]
|
|
z = cam_coord[:, 2]
|
|
return np.stack((x, y, z), 1)
|
|
|
|
|
|
def pixel2cam(pixel_coord, f, c):
|
|
x = (pixel_coord[:, 0] - c[0]) / f[0] * pixel_coord[:, 2]
|
|
y = (pixel_coord[:, 1] - c[1]) / f[1] * pixel_coord[:, 2]
|
|
z = pixel_coord[:, 2]
|
|
return np.stack((x, y, z), 1)
|
|
|
|
|
|
def world2cam(world_coord, R, t):
|
|
cam_coord = np.dot(R, world_coord.transpose(1, 0)).transpose(1, 0) + t.reshape(1, 3)
|
|
return cam_coord
|
|
|
|
|
|
def cam2world(cam_coord, R, t):
|
|
world_coord = np.dot(np.linalg.inv(R), (cam_coord - t.reshape(1, 3)).transpose(1, 0)).transpose(1, 0)
|
|
return world_coord
|
|
|
|
|
|
def rigid_transform_3D(A, B):
|
|
n, dim = A.shape
|
|
centroid_A = np.mean(A, axis=0)
|
|
centroid_B = np.mean(B, axis=0)
|
|
H = np.dot(np.transpose(A - centroid_A), B - centroid_B) / n
|
|
U, s, V = np.linalg.svd(H)
|
|
R = np.dot(np.transpose(V), np.transpose(U))
|
|
if np.linalg.det(R) < 0:
|
|
s[-1] = -s[-1]
|
|
V[2] = -V[2]
|
|
R = np.dot(np.transpose(V), np.transpose(U))
|
|
|
|
varP = np.var(A, axis=0).sum()
|
|
c = 1 / varP * np.sum(s)
|
|
|
|
t = -np.dot(c * R, np.transpose(centroid_A)) + np.transpose(centroid_B)
|
|
return c, R, t
|
|
|
|
|
|
def rigid_align(A, B):
|
|
c, R, t = rigid_transform_3D(A, B)
|
|
A2 = np.transpose(np.dot(c * R, np.transpose(A))) + t
|
|
return A2
|
|
|
|
|
|
def transform_joint_to_other_db(src_joint, src_name, dst_name):
|
|
src_joint_num = len(src_name)
|
|
dst_joint_num = len(dst_name)
|
|
|
|
new_joint = np.zeros(((dst_joint_num,) + src_joint.shape[1:]), dtype=np.float32)
|
|
for src_idx in range(len(src_name)):
|
|
name = src_name[src_idx]
|
|
if name in dst_name:
|
|
dst_idx = dst_name.index(name)
|
|
new_joint[dst_idx] = src_joint[src_idx]
|
|
|
|
return new_joint
|
|
|
|
|
|
def rotation_matrix_to_angle_axis(rotation_matrix):
|
|
"""Convert 3x4 rotation matrix to Rodrigues vector
|
|
|
|
Args:
|
|
rotation_matrix (Tensor): rotation matrix.
|
|
|
|
Returns:
|
|
Tensor: Rodrigues vector transformation.
|
|
|
|
Shape:
|
|
- Input: :math:`(N, 3, 4)`
|
|
- Output: :math:`(N, 3)`
|
|
|
|
Example:
|
|
>>> input = torch.rand(2, 3, 4) # Nx4x4
|
|
>>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3
|
|
"""
|
|
# todo add check that matrix is a valid rotation matrix
|
|
quaternion = rotation_matrix_to_quaternion(rotation_matrix)
|
|
return quaternion_to_angle_axis(quaternion)
|
|
|
|
def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
|
|
"""Convert quaternion vector to angle axis of rotation.
|
|
|
|
Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
|
|
|
|
Args:
|
|
quaternion (torch.Tensor): tensor with quaternions.
|
|
|
|
Return:
|
|
torch.Tensor: tensor with angle axis of rotation.
|
|
|
|
Shape:
|
|
- Input: :math:`(*, 4)` where `*` means, any number of dimensions
|
|
- Output: :math:`(*, 3)`
|
|
|
|
Example:
|
|
>>> quaternion = torch.rand(2, 4) # Nx4
|
|
>>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3
|
|
"""
|
|
if not torch.is_tensor(quaternion):
|
|
raise TypeError("Input type is not a torch.Tensor. Got {}".format(
|
|
type(quaternion)))
|
|
|
|
if not quaternion.shape[-1] == 4:
|
|
raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
|
|
.format(quaternion.shape))
|
|
# unpack input and compute conversion
|
|
q1: torch.Tensor = quaternion[..., 1]
|
|
q2: torch.Tensor = quaternion[..., 2]
|
|
q3: torch.Tensor = quaternion[..., 3]
|
|
sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3
|
|
|
|
sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
|
|
cos_theta: torch.Tensor = quaternion[..., 0]
|
|
two_theta: torch.Tensor = 2.0 * torch.where(
|
|
cos_theta < 0.0,
|
|
torch.atan2(-sin_theta, -cos_theta),
|
|
torch.atan2(sin_theta, cos_theta))
|
|
|
|
k_pos: torch.Tensor = two_theta / sin_theta
|
|
k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
|
|
k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
|
|
|
|
angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
|
|
angle_axis[..., 0] += q1 * k
|
|
angle_axis[..., 1] += q2 * k
|
|
angle_axis[..., 2] += q3 * k
|
|
return angle_axis
|
|
|
|
def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
|
|
"""Convert 3x4 rotation matrix to 4d quaternion vector
|
|
|
|
This algorithm is based on algorithm described in
|
|
https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201
|
|
|
|
Args:
|
|
rotation_matrix (Tensor): the rotation matrix to convert.
|
|
|
|
Return:
|
|
Tensor: the rotation in quaternion
|
|
|
|
Shape:
|
|
- Input: :math:`(N, 3, 4)`
|
|
- Output: :math:`(N, 4)`
|
|
|
|
Example:
|
|
>>> input = torch.rand(4, 3, 4) # Nx3x4
|
|
>>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4
|
|
"""
|
|
if not torch.is_tensor(rotation_matrix):
|
|
raise TypeError("Input type is not a torch.Tensor. Got {}".format(
|
|
type(rotation_matrix)))
|
|
|
|
input_shape = rotation_matrix.shape
|
|
if len(input_shape) == 2:
|
|
rotation_matrix = rotation_matrix.unsqueeze(0)
|
|
|
|
if len(rotation_matrix.shape) > 3:
|
|
raise ValueError(
|
|
"Input size must be a three dimensional tensor. Got {}".format(
|
|
rotation_matrix.shape))
|
|
if not rotation_matrix.shape[-2:] == (3, 4):
|
|
raise ValueError(
|
|
"Input size must be a N x 3 x 4 tensor. Got {}".format(
|
|
rotation_matrix.shape))
|
|
|
|
rmat_t = torch.transpose(rotation_matrix, 1, 2)
|
|
|
|
mask_d2 = rmat_t[:, 2, 2] < eps
|
|
|
|
mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
|
|
mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
|
|
|
|
t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
|
q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
|
t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
|
rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
|
|
t0_rep = t0.repeat(4, 1).t()
|
|
|
|
t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
|
q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
|
rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
|
t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
|
|
t1_rep = t1.repeat(4, 1).t()
|
|
|
|
t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
|
q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
|
|
rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
|
|
rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
|
|
t2_rep = t2.repeat(4, 1).t()
|
|
|
|
t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
|
q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
|
rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
|
rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
|
|
t3_rep = t3.repeat(4, 1).t()
|
|
|
|
mask_c0 = mask_d2.float() * mask_d0_d1.float()
|
|
mask_c1 = mask_d2.float() * (1 - mask_d0_d1.float())
|
|
mask_c2 = (1 - mask_d2.float()) * mask_d0_nd1.float()
|
|
mask_c3 = (1 - mask_d2.float()) * (1 - mask_d0_nd1.float())
|
|
mask_c0 = mask_c0.view(-1, 1).type_as(q0)
|
|
mask_c1 = mask_c1.view(-1, 1).type_as(q1)
|
|
mask_c2 = mask_c2.view(-1, 1).type_as(q2)
|
|
mask_c3 = mask_c3.view(-1, 1).type_as(q3)
|
|
|
|
q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
|
|
q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa
|
|
t2_rep * mask_c2 + t3_rep * mask_c3)
|
|
q *= 0.5
|
|
|
|
if len(input_shape) == 2:
|
|
q = q.squeeze(0)
|
|
return q
|
|
|
|
def rot6d_to_axis_angle(x):
|
|
batch_size = x.shape[0]
|
|
|
|
x = x.view(-1, 3, 2)
|
|
a1 = x[:, :, 0]
|
|
a2 = x[:, :, 1]
|
|
b1 = F.normalize(a1)
|
|
b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
|
|
b3 = torch.cross(b1, b2)
|
|
rot_mat = torch.stack((b1, b2, b3), dim=-1) # 3x3 rotation matrix
|
|
|
|
rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1)).cuda().float()], 2) # 3x4 rotation matrix
|
|
axis_angle = rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3) # axis-angle
|
|
axis_angle[torch.isnan(axis_angle)] = 0.0
|
|
return axis_angle
|
|
|
|
def rot6d_to_rotmat(x):
|
|
"""Convert 6D rotation representation to 3x3 rotation matrix.
|
|
Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
|
|
Input:
|
|
(B,6) Batch of 6-D rotation representations
|
|
Output:
|
|
(B,3,3) Batch of corresponding rotation matrices
|
|
"""
|
|
if x.shape[-1] == 6:
|
|
batch_size = x.shape[0]
|
|
if len(x.shape) == 3:
|
|
num = x.shape[1]
|
|
x = rearrange(x, 'b n d -> (b n) d', d=6)
|
|
else:
|
|
num = 1
|
|
x = rearrange(x, 'b (k l) -> b k l', k=3, l=2)
|
|
# x = x.view(-1,3,2)
|
|
a1 = x[:, :, 0]
|
|
a2 = x[:, :, 1]
|
|
b1 = F.normalize(a1)
|
|
b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
|
|
b3 = torch.cross(b1, b2, dim=-1)
|
|
|
|
mat = torch.stack((b1, b2, b3), dim=-1)
|
|
if num > 1:
|
|
mat = rearrange(mat, '(b n) h w-> b n h w', b=batch_size, n=num, h=3, w=3)
|
|
else:
|
|
x = x.view(-1,3,2)
|
|
a1 = x[:, :, 0]
|
|
a2 = x[:, :, 1]
|
|
b1 = F.normalize(a1)
|
|
b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
|
|
b3 = torch.cross(b1, b2, dim=-1)
|
|
mat = torch.stack((b1, b2, b3), dim=-1)
|
|
return mat
|
|
|
|
def batch_rodrigues(theta):
|
|
"""Convert axis-angle representation to rotation matrix.
|
|
Args:
|
|
theta: size = [B, 3]
|
|
Returns:
|
|
Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
|
|
"""
|
|
l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
|
|
angle = torch.unsqueeze(l1norm, -1)
|
|
normalized = torch.div(theta, angle)
|
|
angle = angle * 0.5
|
|
v_cos = torch.cos(angle)
|
|
v_sin = torch.sin(angle)
|
|
quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
|
|
return quat_to_rotmat(quat)
|
|
|
|
def quat_to_rotmat(quat):
|
|
"""Convert quaternion coefficients to rotation matrix.
|
|
Args:
|
|
quat: size = [B, 4] 4 <===>(w, x, y, z)
|
|
Returns:
|
|
Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
|
|
"""
|
|
norm_quat = quat
|
|
norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)
|
|
w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]
|
|
|
|
B = quat.size(0)
|
|
|
|
w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
|
|
wx, wy, wz = w*x, w*y, w*z
|
|
xy, xz, yz = x*y, x*z, y*z
|
|
|
|
rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
|
|
2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
|
|
2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
|
|
return rotMat
|
|
|
|
def sample_joint_features(img_feat, joint_xy):
|
|
height, width = img_feat.shape[2:]
|
|
x = joint_xy[:, :, 0] / (width - 1) * 2 - 1
|
|
y = joint_xy[:, :, 1] / (height - 1) * 2 - 1
|
|
grid = torch.stack((x, y), 2)[:, :, None, :]
|
|
img_feat = F.grid_sample(img_feat, grid, align_corners=True)[:, :, :, 0] # batch_size, channel_dim, joint_num
|
|
img_feat = img_feat.permute(0, 2, 1).contiguous() # batch_size, joint_num, channel_dim
|
|
return img_feat
|
|
|
|
|
|
def soft_argmax_2d(heatmap2d):
|
|
batch_size = heatmap2d.shape[0]
|
|
height, width = heatmap2d.shape[2:]
|
|
heatmap2d = heatmap2d.reshape((batch_size, -1, height * width))
|
|
heatmap2d = F.softmax(heatmap2d, 2)
|
|
heatmap2d = heatmap2d.reshape((batch_size, -1, height, width))
|
|
|
|
accu_x = heatmap2d.sum(dim=(2))
|
|
accu_y = heatmap2d.sum(dim=(3))
|
|
|
|
accu_x = accu_x * torch.arange(width).float().cuda()[None, None, :]
|
|
accu_y = accu_y * torch.arange(height).float().cuda()[None, None, :]
|
|
|
|
accu_x = accu_x.sum(dim=2, keepdim=True)
|
|
accu_y = accu_y.sum(dim=2, keepdim=True)
|
|
|
|
coord_out = torch.cat((accu_x, accu_y), dim=2)
|
|
return coord_out
|
|
|
|
|
|
def soft_argmax_3d(heatmap3d):
|
|
batch_size = heatmap3d.shape[0]
|
|
depth, height, width = heatmap3d.shape[2:]
|
|
heatmap3d = heatmap3d.reshape((batch_size, -1, depth * height * width))
|
|
heatmap3d = F.softmax(heatmap3d, 2)
|
|
heatmap3d = heatmap3d.reshape((batch_size, -1, depth, height, width))
|
|
|
|
accu_x = heatmap3d.sum(dim=(2, 3))
|
|
accu_y = heatmap3d.sum(dim=(2, 4))
|
|
accu_z = heatmap3d.sum(dim=(3, 4))
|
|
|
|
accu_x = accu_x * torch.arange(width).float().cuda()[None, None, :]
|
|
accu_y = accu_y * torch.arange(height).float().cuda()[None, None, :]
|
|
accu_z = accu_z * torch.arange(depth).float().cuda()[None, None, :]
|
|
|
|
accu_x = accu_x.sum(dim=2, keepdim=True)
|
|
accu_y = accu_y.sum(dim=2, keepdim=True)
|
|
accu_z = accu_z.sum(dim=2, keepdim=True)
|
|
|
|
coord_out = torch.cat((accu_x, accu_y, accu_z), dim=2)
|
|
return coord_out
|