def tf2d_compose(pose1, pose12):
"""
Composing pose1 with pose12, i.e. T2 = T1*T12
:param pose1: n x 3 tensor [x,y,yaw]
:param pose12: n x 3 tensor [x,y,yaw]
:return: pose2 n x 3 tensor [x,y,yaw]
"""
num_data = pose1.shape[0]
rot1 = torch.cat([torch.zeros(num_data, 1), torch.zeros(
num_data, 1), pose1[:, 2][:, None]], 1)
rot12 = torch.cat([torch.zeros(num_data, 1), torch.zeros(
num_data, 1), pose12[:, 2][:, None]], 1)
t1 = torch.cat([pose1[:, 0][:, None], pose1[:, 1]
[:, None], torch.zeros(num_data, 1)], 1)
t12 = torch.cat([pose12[:, 0][:, None], pose12[:, 1]
[:, None], torch.zeros(num_data, 1)], 1)
R1 = p3d_t.euler_angles_to_matrix(rot1, "XYZ")
R12 = p3d_t.euler_angles_to_matrix(rot12, "XYZ")
R2 = torch.matmul(R1, R12)
rot2 = p3d_t.matrix_to_euler_angles(R2, "XYZ")
t2 = torch.matmul(R1, t12[:, :, None]) + t1[:, :, None]
t2 = t2[:, :, 0]
tx = t2[:, 0][:, None]
ty = t2[:, 1][:, None]
yaw = rot2[:, 2][:, None]
pose2 = torch.cat([tx, ty, yaw], 1)
return pose2
def tf2d_between(pose1, pose2, device=None, requires_grad=None):
"""
Relative transform of pose2 in pose1 frame, i.e. T12 = T1^{1}*T2
:param pose1: n x 3 tensor [x,y,yaw]
:param pose2: n x 3 tensor [x,y,yaw]
:return: pose12 n x 3 tensor [x,y,yaw]
"""
num_data = pose1.shape[0]
rot1 = torch.cat([torch.zeros(num_data, 1, device=device, requires_grad=requires_grad), torch.zeros(
num_data, 1, device=device, requires_grad=requires_grad), pose1[:, 2][:, None]], 1)
rot2 = torch.cat([torch.zeros(num_data, 1, device=device, requires_grad=requires_grad), torch.zeros(
num_data, 1, device=device, requires_grad=requires_grad), pose2[:, 2][:, None]], 1)
t1 = torch.cat([pose1[:, 0][:, None], pose1[:, 1]
[:, None], torch.zeros(num_data, 1, device=device, requires_grad=requires_grad)], 1)
t2 = torch.cat([pose2[:, 0][:, None], pose2[:, 1]
[:, None], torch.zeros(num_data, 1, device=device, requires_grad=requires_grad)], 1)
R1 = p3d_t.euler_angles_to_matrix(rot1, "XYZ")
R2 = p3d_t.euler_angles_to_matrix(rot2, "XYZ")
R1t = torch.inverse(R1)
R12 = torch.matmul(R1t, R2)
rot12 = p3d_t.matrix_to_euler_angles(R12, "XYZ")
t12 = torch.matmul(R1t, (t2-t1)[:, :, None])
t12 = t12[:, :, 0]
tx = t12[:, 0][:, None]
ty = t12[:, 1][:, None]
yaw = rot12[:, 2][:, None]
pose12 = torch.cat([tx, ty, yaw], 1)
return pose12
def random_rotation(pc, theta=45., euler_order='XYZ', dev=default_dev):
'''
Apply random rotation to point cloud and return both
Output: Angles - radian, Rot - (3, 3) matrix, PC - Point cloud
'''
angles = torch.rand(size=(3, ), device=dev) * np.radians(theta)
rot = transforms.euler_angles_to_matrix(angles, euler_order)
return angles, rot, pc @ rot.T # x.T * R.T <=> R * x
def get_extrinsic_matrix(pose):
"""
Returns the rotation matrix representation of the
rotations and translations from pose.
"""
batch_size, _ = pose.shape
rot = pose[:,:3]
trans = pose[:,3:]
rot = transforms.euler_angles_to_matrix(rot,convention="XYZ")
pose = torch.cat((rot,trans.view(batch_size, 3, 1)), -1)
return pose
def __init__(self, num_hypotheses=8, device="cuda", **kwargs):
"""
:param num_hypotheses: number of camera poses which should be predicted.
:param kwargs: arguments which are passed through to the single camera predictors
"""
super(MultiCameraPredictor, self).__init__()
_encoder = get_encoder(trainable=False)
self.num_hypotheses = num_hypotheses
self.cam_preds = nn.ModuleList([
CameraPredictor(encoder=_encoder, **kwargs)
for _ in range(num_hypotheses)
])
# taken from the original repo
base_rotation = matrix_to_quaternion(
euler_angles_to_matrix(
torch.FloatTensor([0.5, 0, 0]) * np.pi,
"XYZ")).unsqueeze(0) # rotation by PI/2 around the x-axis
base_bias = matrix_to_quaternion(
euler_angles_to_matrix(
torch.FloatTensor([0, 0.25, 0]) * np.pi,
"XYZ")).unsqueeze(0) # rotation by PI/4 around the y-axis
# base_rotation = torch.FloatTensor(
# [0.9239, 0, 0.3827, 0]).unsqueeze(0).unsqueeze(0) # pi/4 (45° )
# # base_rotation = torch.FloatTensor([ 0.7071, 0 , 0.7071, 0]).unsqueeze(0).unsqueeze(0) ## pi/2
# base_bias = torch.FloatTensor(
# [0.7071, 0.7071, 0, 0]).unsqueeze(0).unsqueeze(0) # 90° by x-axis
# taken from the original repo
self.cam_biases = [base_bias]
for i in range(1, self.num_hypotheses):
self.cam_biases.append(
quaternion_multiply(base_rotation, self.cam_biases[i - 1]))
self.cam_biases = torch.stack(self.cam_biases).squeeze().to(device)
def infer_and_save_pose(input_file_refs, input_file, model_wrapper,
image_shape, half, save):
"""
Process a single input file to produce and save visualization
Parameters
----------
input_file_refs : list(str)
Reference image file paths
input_file : str
Image file for pose estimation
model_wrapper : nn.Module
Model wrapper used for inference
image_shape : Image shape
Input image shape
half: bool
use half precision (fp16)
save: str
Save format (npz or png)
"""
base_name = os.path.basename(input_file)
# change to half precision for evaluation if requested
dtype = torch.float16 if half else None
# Load image
def process_image(filename):
image = load_image(filename)
# Resize and to tensor
image = resize_image(image, image_shape)
image = to_tensor(image).unsqueeze(0)
# Send image to GPU if available
if torch.cuda.is_available():
image = image.to('cuda:{}'.format(rank()), dtype=dtype)
return image
image_ref = [
process_image(input_file_ref) for input_file_ref in input_file_refs
]
image = process_image(input_file)
# Depth inference (returns predicted inverse depth)
pose_tensor = model_wrapper.pose(
image, image_ref)[0][0] # take the pose from 1st to 2nd image
rot_matrix = transforms.euler_angles_to_matrix(pose_tensor[3:],
convention="ZYX")
translation = pose_tensor[:3]
poses[base_name] = (rot_matrix, translation)
def transpose_rotation(rotation: Rotation) -> Rotation:
""" As for ortho6d and euler, we will first convert to matrix, then transpose.
As others, we will use
"""
# 1. TRANSPOSE ORTHO6D
matrix = compute_rotation_matrix_from_ortho6d(rotation.ortho6d).transpose(
2, 1)
ortho6d = compute_ortho6d_from_rotation_matrix(matrix)
# 2. TRANSPOSE EULER
matrix = euler_angles_to_matrix(rotation.euler,
convention="XYZ").transpose(2, 1)
euler = matrix_to_euler_angles(matrix, convention="XYZ")
# 3. TRANSPOSE QUAT
quat = quaternion_invert(rotation.quat)
# 4. TRANSPOSE MATRIX
matrix = rotation.matrix.transpose(2, 1)
return Rotation(ortho6d=ortho6d, quat=quat, matrix=matrix, euler=euler)
def matrix_from_angles(rot):
"""Create a rotation matrix from a triplet of rotation angles.
Args:
rot: a torch.Tensor of shape [..., 3], where the last dimension is the rotation
angles, along x, y, and z.
Returns:
A torch.tensor of shape [..., 3, 3], where the last two dimensions are the
rotation matrix.
This function mimics _euler2mat from struct2depth/project.py, for backward
compatibility, but wraps tensorflow_graphics instead of reimplementing it.
The negation and transposition are needed to bridge the differences between
the two.
"""
rank = rot.dim()
# Swap the two last dimensions
perm = torch.cat([
torch.arange(0, rank - 1, dtype=torch.long),
torch.tensor([rank]),
torch.tensor([rank - 1])
],
dim=0)
return euler_angles_to_matrix(-rot, convention="XYZ").permute(*perm)
def build_rotation(rotation, format="matrix") -> Rotation:
""" Convert roation (with format) into Rotation.
format: matrix, ortho6d, quat, euler
"""
# 1. CONVERT SPECIFIED FORMAT TO MATRIX FIRST
if format == "matrix":
matrix = rotation
elif format == "ortho6d":
matrix = compute_rotation_matrix_from_ortho6d(rotation)
elif format == "euler":
matrix = euler_angles_to_matrix(rotation, convention="XYZ")
elif format == "quat":
matrix = quaternion_to_matrix(rotation)
else:
raise TypeError
# 2. BUILD ROTATION
return Rotation(
ortho6d=rotation if format == "ortho6d" else
compute_ortho6d_from_rotation_matrix(matrix),
quat=rotation if format == "quat" else matrix_to_quaternion(matrix),
matrix=rotation if format == "matrix" else matrix,
euler=rotation if format == "euler" else matrix_to_euler_angles(
matrix, convention="XYZ"))