def while_body(i, est_pose, est_inputs): with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE): features = mlp_conv(est_inputs, [128, 256], bn=self.bn) features_global = tf.reduce_max(features, axis=1, keep_dims=True, name='maxpool_0') features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2) with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE): features = mlp_conv(features, [512, 1024], bn=self.bn) features = tf.reduce_max(features, axis=1, name='maxpool_1') with tf.variable_scope('fc', reuse=tf.AUTO_REUSE): if self.rot_representation == 'quat': est_pose_rep_i = mlp(features, [1024, 1024, 512, 512, 256, 7], bn=self.bn) elif self.rot_representation == '6dof': est_pose_rep_i = mlp(features, [1024, 1024, 512, 512, 256, 9], bn=self.bn) with tf.variable_scope('est', reuse=tf.AUTO_REUSE): if self.rot_representation == 'quat': t = tf.expand_dims(est_pose_rep_i[:, 4:], axis=2) q = est_pose_rep_i[:, 0:4] R = quat2rotm_tf(q) elif self.rot_representation == '6dof': t = tf.expand_dims(est_pose_rep_i[:, 6:], axis=2) mat6d = est_pose_rep_i[:, 0:6] R = mat6d2rotm_tf(mat6d) est_pose_T_i = tf.concat([ tf.concat([R, t], axis=2), tf.concat([tf.zeros([B, 1, 3]), tf.ones([B, 1, 1])], axis=2)], axis=1) est_inputs = transform_tf(est_inputs, est_pose_T_i) est_pose = tf.linalg.matmul(est_pose_T_i, est_pose) return [tf.add(i, 1), est_pose, est_inputs]
def test_identity(self): rotm = euldeg2rotm(0, 0, 0) rotm = tf.convert_to_tensor( np.tile(np.expand_dims(rotm, axis=0), (24, 1, 1))) mat6d_gt = rotm2mat6d_tf(rotm) rotm_est = mat6d2rotm_tf(mat6d_gt) assert np.all(np.isclose(rotm, rotm_est))
def test_rand(self): rotm = euldeg2rotm(23, 82, 13) rotm = tf.convert_to_tensor(np.expand_dims(rotm, axis=0)) mat6d_gt = rotm2mat6d_tf(rotm) rotm_est = mat6d2rotm_tf(mat6d_gt) assert np.all(np.isclose(rotm, rotm_est))