105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
|
|
def euler_to_rotation_matrix(x, y, z):
|
|
"""
|
|
|
|
:param x: Tensor of shape (B, 1) - x axis rotation
|
|
:param y: Tensor of shape (B, 1) - y axis rotation
|
|
:param z: Tensor of shape (B, 1) - z axis rotation
|
|
:return: Rotation matrix for the given euler anglers, in the order rotation(x).rotation(y).rotation(z)
|
|
"""
|
|
B = tf.shape(z)[0]
|
|
|
|
# Euler angles should be between -pi and pi, clip so the pose network is coerced to this range
|
|
z = tf.clip_by_value(z, -np.pi, np.pi)
|
|
y = tf.clip_by_value(y, -np.pi, np.pi)
|
|
x = tf.clip_by_value(x, -np.pi, np.pi)
|
|
|
|
# Expand to B x 1 x 1
|
|
z = tf.expand_dims(tf.expand_dims(z, -1), -1)
|
|
y = tf.expand_dims(tf.expand_dims(y, -1), -1)
|
|
x = tf.expand_dims(tf.expand_dims(x, -1), -1)
|
|
|
|
zeros = tf.zeros([B, 1, 1])
|
|
ones = tf.ones([B, 1, 1])
|
|
|
|
cosx = tf.cos(x)
|
|
sinx = tf.sin(x)
|
|
rotx_1 = tf.concat([ones, zeros, zeros], axis=3)
|
|
rotx_2 = tf.concat([zeros, cosx, -sinx], axis=3)
|
|
rotx_3 = tf.concat([zeros, sinx, cosx], axis=3)
|
|
xmat = tf.concat([rotx_1, rotx_2, rotx_3], axis=2)
|
|
|
|
cosz = tf.cos(z)
|
|
sinz = tf.sin(z)
|
|
rotz_1 = tf.concat([cosz, -sinz, zeros], axis=3)
|
|
rotz_2 = tf.concat([sinz, cosz, zeros], axis=3)
|
|
rotz_3 = tf.concat([zeros, zeros, ones], axis=3)
|
|
zmat = tf.concat([rotz_1, rotz_2, rotz_3], axis=2)
|
|
|
|
cosy = tf.cos(y)
|
|
siny = tf.sin(y)
|
|
roty_1 = tf.concat([cosy, zeros, siny], axis=3)
|
|
roty_2 = tf.concat([zeros, ones, zeros], axis=3)
|
|
roty_3 = tf.concat([-siny, zeros, cosy], axis=3)
|
|
ymat = tf.concat([roty_1, roty_2, roty_3], axis=2)
|
|
|
|
rotMat = tf.matmul(tf.matmul(xmat, ymat), zmat)
|
|
return rotMat
|
|
|
|
|
|
def pose_vec2mat(vec):
|
|
"""Converts 6DoF parameters to transformation matrix
|
|
Args:
|
|
vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz -- [B, 6]
|
|
Returns:
|
|
A transformation matrix -- [B, 4, 4]
|
|
"""
|
|
batch_size, _ = vec.get_shape().as_list()
|
|
translation = tf.slice(vec, [0, 0], [-1, 3])
|
|
translation = tf.expand_dims(translation, -1)
|
|
rx = tf.slice(vec, [0, 3], [-1, 1])
|
|
ry = tf.slice(vec, [0, 4], [-1, 1])
|
|
rz = tf.slice(vec, [0, 5], [-1, 1])
|
|
rot_mat = euler_to_rotation_matrix(rx, ry, rz)
|
|
rot_mat = tf.squeeze(rot_mat, axis=[1])
|
|
filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
|
|
filler = tf.tile(filler, [batch_size, 1, 1])
|
|
transform_mat = tf.concat([rot_mat, translation], axis=2)
|
|
transform_mat = tf.concat([transform_mat, filler], axis=1)
|
|
return transform_mat
|
|
|
|
|
|
def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics):
|
|
"""
|
|
Calculate the reprojected image from the source to the target, based on the given depth, pose and intrinsics
|
|
|
|
SFM Learner inverse warp step
|
|
ps ~ K.T(t->s).Dt(pt).K^-1.pt
|
|
|
|
Idea is to map the pixel coordinates of the target image to 3d space (Dt(pt).K^-1.pt), then map these onto
|
|
the source image in pixel coordinates (K.T(t->s).{3d coord}), then using the projected coordinates we sample
|
|
the pixels in the source image (ps) to reconstruct the target image.
|
|
|
|
:param target_img: Tensor (batch, height, width, 3)
|
|
:param source_img: Tensor, same shape as target_img
|
|
:param depth: Tensor, (batch, height, width, 1)
|
|
:param pose: (batch, 6)
|
|
:param intrinsics: (batch, 3, 3)
|
|
:return: The source image reprojected to the target
|
|
"""
|
|
# Convert pose vector (output of pose net) to pose matrix (4x4)
|
|
|
|
# Convert intrinsics matrix (3x3) to (4x4) so it can be multiplied by the pose net
|
|
intrinsics_4x4 =
|
|
|
|
# Calculate inverse of the 4x4 intrinsics matrix
|
|
tf.linalg.inv()
|
|
|
|
# Create grid of homogenous coordinates
|
|
|
|
#
|
|
pass
|