├── .gitignore ├── doc ├── eq1.png ├── eq2.png └── results.png ├── .vscode └── settings.json ├── README.md ├── example.py └── rot_6d.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.py[cod] 3 | *$py.class -------------------------------------------------------------------------------- /doc/eq1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Janus-Shiau/6d_rot_tensorflow/HEAD/doc/eq1.png -------------------------------------------------------------------------------- /doc/eq2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Janus-Shiau/6d_rot_tensorflow/HEAD/doc/eq2.png -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "python.pythonPath": "/home/jiayau/env/basic-py3/bin/python" 3 | } -------------------------------------------------------------------------------- /doc/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Janus-Shiau/6d_rot_tensorflow/HEAD/doc/results.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 6d_rot_tensorflow 2 | 6D rotation representation (["On the Continuity of Rotation Representations in Neural Networks"](https://arxiv.org/abs/1812.07035)) for tensorflow. 3 | 4 | ### Environment 5 | This code is implemmented and tested with [tensorflow](https://www.tensorflow.org/) 1.11.0. \ 6 | I didn't use any spetial operator, so it should also work for other version of tensorflow. 7 | 8 | ### Usage 9 | Just add the tf_rotation6d_to_matrix after your output, whose last dimension of tensor should be 6. 10 | ``` 11 | """ 12 | Any model output whose last dimension is 6. 13 | e.g. output = tf.layers.dense(hidden, 6) 14 | """ 15 | rot = tf_rotation6d_to_matrix(output) 16 | ``` 17 | 18 | I very simple example of transformation between 6D continuous representation and SO(3) can be found in `example.py` 19 | 20 | ### Details 21 | Here I crop some parts of the context from the paper, FYI. 22 | 23 | According to the Section 3 and 4 of the paper, the target transformation between continuous representation and SO(n) can be formulate as follows. It's derived based on a Gram-Schmidt process. 24 | 25 | Continuous representation of SO(n) 27 | 28 | If you found it looks a little bit complicated, you can directly go to Appendix B. There is a very simple formulation of the 6D and SO(3). 29 | 30 | Continuous representation of SO(3) 32 | 33 | Besides, according to the features of rotation matrix, the formulation can be quite concise. You can found the concise transformation between 6D and SO(3) in the source code. 34 | 35 | ### Contact & Copy Right 36 | Code work by Jia-Yau Shiau . 37 | 38 | -------------------------------------------------------------------------------- /example.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2019 [Jia-Yau Shiau] 3 | Code work by Jia-Yau (jiayau.shiau@gmail.com). 4 | -------------------------------------------------- 5 | A very simple example code for converting between the rotation6d and rotation matrix. 6 | ''' 7 | 8 | import math 9 | from functools import reduce 10 | 11 | import numpy as np 12 | import tensorflow as tf 13 | from termcolor import colored 14 | 15 | from rot_6d import tf_matrix_to_rotation6d, tf_rotation6d_to_matrix 16 | 17 | 18 | def euler2mat(z=0, y=0, x=0): 19 | ''' Return matrix for rotations around z, y and x axes 20 | 21 | Uses the z, then y, then x convention above 22 | 23 | Parameters 24 | ---------- 25 | z : scalar 26 | Rotation angle in radians around z-axis (performed first) 27 | y : scalar 28 | Rotation angle in radians around y-axis 29 | x : scalar 30 | Rotation angle in radians around x-axis (performed last) 31 | 32 | Returns 33 | ------- 34 | M : array shape (3,3) 35 | Rotation matrix giving same rotation as for given angles 36 | 37 | Code fork from "https://afni.nimh.nih.gov/pub/dist/src/pkundu/meica.libs/nibabel/eulerangles.py". 38 | ''' 39 | Ms = [] 40 | if z: 41 | cosz = math.cos(z) 42 | sinz = math.sin(z) 43 | Ms.append(np.array( 44 | [[cosz, -sinz, 0], 45 | [sinz, cosz, 0], 46 | [0, 0, 1]])) 47 | if y: 48 | cosy = math.cos(y) 49 | siny = math.sin(y) 50 | Ms.append(np.array( 51 | [[cosy, 0, siny], 52 | [0, 1, 0], 53 | [-siny, 0, cosy]])) 54 | if x: 55 | cosx = math.cos(x) 56 | sinx = math.sin(x) 57 | Ms.append(np.array( 58 | [[1, 0, 0], 59 | [0, cosx, -sinx], 60 | [0, sinx, cosx]])) 61 | if Ms: 62 | return reduce(np.dot, Ms[::-1]) 63 | return np.eye(3) 64 | 65 | 66 | if __name__ == "__main__": 67 | np_mat = [] 68 | for i in range(3): 69 | np_mat.append(euler2mat(i, i, i)) 70 | 71 | np_mat = np.array(np_mat) 72 | 73 | tf_mat = tf.convert_to_tensor(np_mat) 74 | tf_r6d = tf_matrix_to_rotation6d(tf_mat) 75 | tf_mat_from_r6d = tf_rotation6d_to_matrix(tf_r6d) 76 | with tf.Session() as sess: 77 | mat, r6d, mat_from_r6d = sess.run([tf_mat, tf_r6d, tf_mat_from_r6d]) 78 | 79 | print (colored("[Original Rotation Matrix]", 'yellow')) 80 | print (mat) 81 | print (colored("[Rotation 6D from Rotation Matrix]", 'yellow')) 82 | print (r6d) 83 | print (colored("[Rotation Matrix from Rotation 6d]", 'yellow')) 84 | print (mat_from_r6d.reshape(-1,3,3)) 85 | -------------------------------------------------------------------------------- /rot_6d.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2019 [Jia-Yau Shiau] 3 | Code work by Jia-Yau (jiayau.shiau@gmail.com). 4 | -------------------------------------------------- 5 | The implementation of 6D rotatiton representation, 6 | based on 7 | 8 | https://arxiv.org/abs/1812.07035 9 | 10 | "On the continuity of rotation representations in neural networks" 11 | Yi Zhou, Connelly Barnes, Jingwan Lu, Jimei Yang, Hao Li. 12 | Conference on Neural Information Processing Systems (NeurIPS) 2019. 13 | ''' 14 | import tensorflow as tf 15 | 16 | 17 | def tf_rotation6d_to_matrix(r6d): 18 | """ Compute rotation matrix from 6D rotation representation. 19 | Implementation base on 20 | https://arxiv.org/abs/1812.07035 21 | [Inputs] 22 | 6D rotation representation (last dimension is 6) 23 | [Returns] 24 | flattened rotation matrix (last dimension is 9) 25 | """ 26 | tensor_shape = r6d.get_shape().as_list() 27 | 28 | if not tensor_shape[-1] == 6: 29 | raise AttributeError("The last demension of the inputs in tf_rotation6d_to_matrix should be 6, \ 30 | but found tensor with shape {}".format(tensor_shape[-1])) 31 | 32 | with tf.variable_scope('rot6d_to_matrix'): 33 | r6d = tf.reshape(r6d, [-1,6]) 34 | x_raw = r6d[:,0:3] 35 | y_raw = r6d[:,3:6] 36 | 37 | x = tf.nn.l2_normalize(x_raw, axis=-1) 38 | z = tf.cross(x, y_raw) 39 | z = tf.nn.l2_normalize(z, axis=-1) 40 | y = tf.cross(z, x) 41 | 42 | x = tf.reshape(x, [-1,3,1]) 43 | y = tf.reshape(y, [-1,3,1]) 44 | z = tf.reshape(z, [-1,3,1]) 45 | matrix = tf.concat([x,y,z], axis=-1) 46 | 47 | if len(tensor_shape) == 1: 48 | matrix = tf.reshape(matrix, [9]) 49 | else: 50 | output_shape = tensor_shape[:-1] + [9] 51 | matrix = tf.reshape(matrix, output_shape) 52 | 53 | return matrix 54 | 55 | def tf_matrix_to_rotation6d(mat): 56 | """ Get 6D rotation representation for rotation matrix. 57 | Implementation base on 58 | https://arxiv.org/abs/1812.07035 59 | [Inputs] 60 | flattened rotation matrix (last dimension is 9) 61 | [Returns] 62 | 6D rotation representation (last dimension is 6) 63 | """ 64 | tensor_shape = mat.get_shape().as_list() 65 | 66 | if not ((tensor_shape[-1] == 3 and tensor_shape[-2] == 3) or (tensor_shape[-1] == 9)): 67 | raise AttributeError("The inputs in tf_matrix_to_rotation6d should be [...,9] or [...,3,3], \ 68 | but found tensor with shape {}".format(tensor_shape[-1])) 69 | 70 | with tf.variable_scope('matrix_to_ration_6d'): 71 | mat = tf.reshape(mat, [-1, 3, 3]) 72 | r6d = tf.concat([mat[...,0], mat[...,1]], axis=-1) 73 | 74 | if len(tensor_shape) == 1: 75 | r6d = tf.reshape(r6d, [6]) 76 | 77 | return r6d 78 | --------------------------------------------------------------------------------