Python tensorflow.cross() Examples

def rectangle_homogeneous_lines(rv):
    """

    # Arguments

    rv: rectangle vectors [v0yx, v1yx, v2yx, v3yx]


    # Returns

    [r0abc, r1abc, r2abc, r3abc]

    """
    # ax + by + c = 0
    dv = rv[0] - rv[1]
    # TODO(ahundt) make sure cross product doesn't need to be in xy order
    r0abc = K.concatenate([dv[0], dv[1], tf.cross(rv[0], rv[1])])
    dv = rv[1] - rv[2]
    r1abc = K.concatenate([dv[1], dv[2], tf.cross(rv[1], rv[2])])
    dv = rv[2] - rv[3]
    r2abc = K.concatenate([dv[2], dv[3], tf.cross(rv[2], rv[3])])
    dv = rv[3] - rv[0]
    r3abc = K.concatenate([dv[3], dv[0], tf.cross(rv[3], rv[0])])
    return [r0abc, r1abc, r2abc, r3abc] 

def Compute_norm(self,face_shape,facemodel):
		shape = face_shape
		face_id = facemodel.face_buf
		point_id = facemodel.point_buf

		# face_id and point_id index starts from 1
		face_id = tf.cast(face_id - 1,tf.int32)
		point_id = tf.cast(point_id - 1,tf.int32)

		#compute normal for each face
		v1 = tf.gather(shape,face_id[:,0], axis = 1)
		v2 = tf.gather(shape,face_id[:,1], axis = 1)
		v3 = tf.gather(shape,face_id[:,2], axis = 1)
		e1 = v1 - v2
		e2 = v2 - v3
		face_norm = tf.cross(e1,e2)

		face_norm = tf.nn.l2_normalize(face_norm, dim = 2) # normalized face_norm first
		face_norm = tf.concat([face_norm,tf.zeros([tf.shape(face_shape)[0],1,3])], axis = 1)

		#compute normal for each vertex using one-ring neighborhood
		v_norm = tf.reduce_sum(tf.gather(face_norm, point_id, axis = 1), axis = 2)
		v_norm = tf.nn.l2_normalize(v_norm, dim = 2)
		
		return v_norm 

def compute_tri_normal(vertex,tri, vertex_tri):
    # Unit normals to the faces
    # vertex : 3xvertex_num
    # tri : 3xtri_num

    vertex = tf.transpose(vertex)

    vt1_indices, vt2_indices, vt3_indices = tf.split(tf.transpose(tri), num_or_size_splits = 3, axis = 1)

    vt1 = tf.gather_nd(vertex, vt1_indices)
    vt2 = tf.gather_nd(vertex, vt2_indices)
    vt3 = tf.gather_nd(vertex, vt3_indices)

    normalf = tf.cross(vt2 - vt1, vt3 - vt1)
    normalf = tf.nn.l2_normalize(normalf, dim = 1)

    return normalf 

def cameraMat(param):
    theta = param[0] * np.pi / 180.0
    camy = param[3] * tf.sin(param[1] * np.pi / 180.0)
    lens = param[3] * tf.cos(param[1] * np.pi / 180.0)
    camx = lens * tf.cos(theta)
    camz = lens * tf.sin(theta)
    Z = tf.stack([camx, camy, camz])

    x = camy * tf.cos(theta + np.pi)
    z = camy * tf.sin(theta + np.pi)
    Y = tf.stack([x, lens, z])
    X = tf.cross(Y, Z)

    cm_mat = tf.stack([normal(X), normal(Y), normal(Z)])
    return cm_mat, Z 

def test_Cross(self):
        t = tf.cross(*self.random((4, 3), (4, 3)))
        self.check(t)


    #
    # basic math ops
    # 

def compute_consistent_plane_frame(normal):
    # Input:  normal is Bx3
    # Returns: x_axis, y_axis, both of dimension Bx3
    batch_size = tf.shape(normal)[0]
    candidate_axes = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] # Actually, 2 should be enough. This may still cause singularity TODO!!!
    y_axes = []
    for tmp_axis in candidate_axes:
        tf_axis = tf.tile(tf.expand_dims(tf.constant(dtype=tf.float32, value=tmp_axis), axis=0), [batch_size, 1]) # Bx3
        y_axes.append(tf.cross(normal, tf_axis))
    y_axes = tf.stack(y_axes, axis=0) # QxBx3
    y_axes_norm = tf.norm(y_axes, axis=2) # QxB
    # choose the axis with largest norm
    y_axes_chosen_idx = tf.argmax(y_axes_norm, axis=0) # B
    # y_axes_chosen[b, :] = y_axes[y_axes_chosen_idx[b], b, :]
    indices_0 = tf.tile(tf.expand_dims(y_axes_chosen_idx, axis=1), [1, 3]) # Bx3
    indices_1 = tf.tile(tf.expand_dims(tf.range(batch_size), axis=1), [1, 3]) # Bx3
    indices_2 = tf.tile(tf.expand_dims(tf.range(3), axis=0), [batch_size, 1]) # Bx3
    indices = tf.stack([tf.cast(indices_0, tf.int32), indices_1, indices_2], axis=2) # Bx3x3
    y_axes = tf.gather_nd(y_axes, indices=indices) # Bx3
    if tf.VERSION == '1.4.1':
        y_axes = tf.nn.l2_normalize(y_axes, dim=1)
    else:
        y_axes = tf.nn.l2_normalize(y_axes, axis=1)
    x_axes = tf.cross(y_axes, normal) # Bx3

    return x_axes, y_axes 

def testGradientRandomValues(self):
    with self.test_session():
      us = [2, 3]
      u = tf.reshape([0.854, -0.616, 0.767, 0.725, -0.927, 0.159], shape=us)
      v = tf.reshape([-0.522, 0.755, 0.407, -0.652, 0.241, 0.247], shape=us)
      s = tf.cross(u, v)
      jacob_u, jacob_v = tf.test.compute_gradient([u, v], [us, us], s, us)

    self.assertAllClose(jacob_u[0], jacob_u[1], rtol=1e-3, atol=1e-3)
    self.assertAllClose(jacob_v[0], jacob_v[1], rtol=1e-3, atol=1e-3) 

def rotate_shape(m, mshape, output_size = 224): 

    n_size = get_shape(m)    
    n_size = n_size[0]

    m_single     = tf.split(axis = 0, num_or_size_splits = n_size, value = m)
    shape_single = tf.split(axis = 0, num_or_size_splits = n_size, value = mshape)
    
    vertex2ds = []

    for i in range(n_size):

        m_i = tf.transpose(tf.reshape(m_single[i], [4,2]))
        m_i_row1 = tf.nn.l2_normalize(m_i[0,0:3], dim = 0)
        m_i_row2 = tf.nn.l2_normalize(m_i[1,0:3], dim = 0)
        m_i_row3 = tf.concat([tf.reshape(tf.cross(m_i_row1, m_i_row2), shape = [1, 3]), tf.zeros([1, 1])], axis = 1)
                  
        m_i = tf.concat([m_i, m_i_row3], axis = 0)

        vertex3d_rs = tf.transpose(tf.reshape( shape_single[i], shape = [-1, 3] ))

        vertex4d = tf.concat(axis = 0, values = [vertex3d_rs, tf.ones([1, get_shape(vertex3d_rs)[1]], tf.float32)])
        
        vertex2d = tf.matmul(m_i, vertex4d, False, False)
        vertex2d = tf.transpose(vertex2d)
        
        [vertex2d_u, vertex2d_v, vertex2d_z]   = tf.split(axis=1, num_or_size_splits=3, value=vertex2d)
        vertex2d_u = vertex2d_u - 1
        vertex2d_v = output_size - vertex2d_v

        vertex2d = tf.concat(axis=1, values=[vertex2d_v, vertex2d_u, vertex2d_z])
        vertex2d = tf.transpose(vertex2d)

        vertex2ds.append(vertex2d)

    return tf.stack(vertex2ds) 

def get_edge_cross_feature(point_cloud, nn_idx, k=20):
  """Construct edge feature for each point
  Args:
    point_cloud: (batch_size, num_points, 1, num_dims)
    nn_idx: (batch_size, num_points, k)
    k: int

  Returns:
    edge features: (batch_size, num_points, k, num_dims)
  """
  og_batch_size = point_cloud.get_shape().as_list()[0]
  point_cloud = tf.squeeze(point_cloud)
  if og_batch_size == 1:
    point_cloud = tf.expand_dims(point_cloud, 0)

  point_cloud_central = point_cloud

  point_cloud_shape = point_cloud.get_shape()
  batch_size = point_cloud_shape[0].value
  num_points = point_cloud_shape[1].value
  num_dims = point_cloud_shape[2].value

  idx_ = tf.range(batch_size) * num_points
  idx_ = tf.reshape(idx_, [batch_size, 1, 1]) 

  point_cloud_flat = tf.reshape(point_cloud, [-1, num_dims])
  point_cloud_neighbors = tf.gather(point_cloud_flat, nn_idx+idx_)
  point_cloud_central = tf.expand_dims(point_cloud_central, axis=-2)

  point_cloud_central = tf.tile(point_cloud_central, [1, 1, k, 1])

  edge_feature = tf.concat([point_cloud_central, point_cloud_neighbors-point_cloud_central, 
                            tf.cross(point_cloud_central, 
                                            point_cloud_neighbors-point_cloud_central)], axis=-1)
  return edge_feature 

def rotate_vector_by_quaternion(q, v, q_ndims=None, v_ndims=None):
    """Rotate a vector (or tensor with last dimension of 3) by q.

    This function computes v' = q * v * conjugate(q) but faster.
    Fast version can be found here:
    https://blog.molecular-matters.com/2013/05/24/a-faster-quaternion-vector-multiplication/

    Args:
        q: A `Quaternion` or `tf.Tensor` with shape (..., 4)
        v: A `tf.Tensor` with shape (..., 3)
        q_ndims: The number of dimensions of q. Only necessary to specify if
            the shape of q is unknown.
        v_ndims: The number of dimensions of v. Only necessary to specify if
            the shape of v is unknown.

    Returns: A `tf.Tensor` with the broadcasted shape of v and q.
    """
    v = tf.convert_to_tensor(v)
    q = q.normalized()
    w = q.value()[..., 0]
    q_xyz = q.value()[..., 1:]
    # Broadcast shapes. Todo(phil): Prepare a pull request which adds
    # broadcasting support to tf.cross
    if q_xyz.shape.ndims is not None:
        q_ndims = q_xyz.shape.ndims
    if v.shape.ndims is not None:
        v_ndims = v.shape.ndims
    for _ in range(v_ndims - q_ndims):
        q_xyz = tf.expand_dims(q_xyz, axis=0)
    for _ in range(q_ndims - v_ndims):
        v = tf.expand_dims(v, axis=0) + tf.zeros_like(q_xyz)
    q_xyz += tf.zeros_like(v)
    v += tf.zeros_like(q_xyz)
    t = 2 * tf.cross(q_xyz, v)
    return v + tf.expand_dims(w, axis=-1) * t + tf.cross(q_xyz, t)


# ____________________________________________________________________________
#                      The quaternion class 

def look_at(eye, center, world_up):
  """Computes camera viewing matrices.

  Functionality mimes gluLookAt (third_party/GL/glu/include/GLU/glu.h).

  Args:
    eye: 2-D float32 tensor with shape [batch_size, 3] containing the XYZ world
        space position of the camera.
    center: 2-D float32 tensor with shape [batch_size, 3] containing a position
        along the center of the camera's gaze.
    world_up: 2-D float32 tensor with shape [batch_size, 3] specifying the
        world's up direction; the output camera will have no tilt with respect
        to this direction.

  Returns:
    A [batch_size, 4, 4] float tensor containing a right-handed camera
    extrinsics matrix that maps points from world space to points in eye space.
  """
  batch_size = center.shape[0].value
  vector_degeneracy_cutoff = 1e-6
  forward = center - eye
  forward_norm = tf.norm(forward, ord='euclidean', axis=1, keepdims=True)
  tf.assert_greater(
      forward_norm,
      vector_degeneracy_cutoff,
      message='Camera matrix is degenerate because eye and center are close.')
  forward = tf.divide(forward, forward_norm)

  to_side = tf.cross(forward, world_up)
  to_side_norm = tf.norm(to_side, ord='euclidean', axis=1, keepdims=True)
  tf.assert_greater(
      to_side_norm,
      vector_degeneracy_cutoff,
      message='Camera matrix is degenerate because up and gaze are close or'
      'because up is degenerate.')
  to_side = tf.divide(to_side, to_side_norm)
  cam_up = tf.cross(to_side, forward)

  w_column = tf.constant(
      batch_size * [[0., 0., 0., 1.]], dtype=tf.float32)  # [batch_size, 4]
  w_column = tf.reshape(w_column, [batch_size, 4, 1])
  view_rotation = tf.stack(
      [to_side, cam_up, -forward,
       tf.zeros_like(to_side, dtype=tf.float32)],
      axis=1)  # [batch_size, 4, 3] matrix
  view_rotation = tf.concat(
      [view_rotation, w_column], axis=2)  # [batch_size, 4, 4]

  identity_batch = tf.tile(tf.expand_dims(tf.eye(3), 0), [batch_size, 1, 1])
  view_translation = tf.concat([identity_batch, tf.expand_dims(-eye, 2)], 2)
  view_translation = tf.concat(
      [view_translation,
       tf.reshape(w_column, [batch_size, 1, 4])], 1)
  camera_matrices = tf.matmul(view_rotation, view_translation)
  return camera_matrices 

def _DEPRECATED_compute_normal(vertex, tri, vertex_tri):
    # Unit normals to the faces
    # vertex : 3xvertex_num
    # tri : 3xtri_num

    vertex = tf.transpose(vertex)

    vt1_indices, vt2_indices, vt3_indices = tf.split(tf.transpose(tri), num_or_size_splits = 3, axis = 1)
    

    vt1 = tf.gather_nd(vertex, vt1_indices)
    #print('get_shape(vt1)')
    #print(get_shape(vt1))
    vt2 = tf.gather_nd(vertex, vt2_indices)
    vt3 = tf.gather_nd(vertex, vt3_indices)


    normalf = tf.cross(vt2 - vt1, vt3 - vt1)
    normalf = tf.nn.l2_normalize(normalf, dim = 1)

    mask = tf.tile( tf.expand_dims(  tf.not_equal(vertex_tri, tri.shape[1] - 1), 2), multiples = [1, 1, 3])
    mask = tf.cast( mask, vertex.dtype  )
    vertex_tri = tf.reshape(vertex_tri, shape = [-1, 1])
    normal = tf.reshape(tf.gather_nd(normalf, vertex_tri), shape = [8, -1, 3])

    normal = tf.reduce_sum( tf.multiply( normal, mask ),  axis = 0)
    normal = tf.nn.l2_normalize(normal, dim = 1)


    #print('get_shape(normalf)')
    #print(get_shape(normalf))


    #print('get_shape(normal)')
    #print(get_shape(normal))


    # enforce that the normal are outward
    v = vertex - tf.reduce_mean(vertex,0)
    s = tf.reduce_sum( tf.multiply(v, normal), 0 )

    count_s_greater_0 = tf.count_nonzero( tf.greater(s, 0) )
    count_s_less_0 = tf.count_nonzero( tf.less(s, 0) )

    sign = 2 * tf.cast(tf.greater(count_s_greater_0, count_s_less_0), tf.float32) - 1
    normal = tf.multiply(normal, sign)
    normalf = tf.multiply(normalf, sign)

    return normal, normalf 

def depth2normal_layer_batch(depth_map, intrinsics, inverse, nei=3):

    ## depth_map is in rank 3 [batch, h, w], intrinsics are in rank 2 [batch,4]
    ## mask is used to filter the background with infinite depth
    mask = tf.greater(depth_map, tf.zeros(depth_map.get_shape().as_list()))

    if inverse:
        mask_clip = 1e-8 * (1.0-tf.cast(mask, tf.float32)) ## Add black pixels (depth = infinite) with delta
        depth_map += mask_clip
        depth_map = 1.0/depth_map ## inverse depth map
    kitti_shape = depth_map.get_shape().as_list()
    pts_3d_map = compute_3dpts_batch(depth_map, intrinsics)

    ## shift the 3d pts map by nei along 8 directions
    pts_3d_map_ctr = pts_3d_map[:,nei:-nei, nei:-nei, :]
    pts_3d_map_x0 = pts_3d_map[:,nei:-nei, 0:-(2*nei), :]
    pts_3d_map_y0 = pts_3d_map[:,0:-(2*nei), nei:-nei, :]
    pts_3d_map_x1 = pts_3d_map[:,nei:-nei, 2*nei:, :]
    pts_3d_map_y1 = pts_3d_map[:,2*nei:, nei:-nei, :]
    pts_3d_map_x0y0 = pts_3d_map[:,0:-(2*nei), 0:-(2*nei), :]
    pts_3d_map_x0y1 = pts_3d_map[:,2*nei:, 0:-(2*nei), :]
    pts_3d_map_x1y0 = pts_3d_map[:,0:-(2*nei), 2*nei:, :]
    pts_3d_map_x1y1 = pts_3d_map[:,2*nei:, 2*nei:, :]

    ## generate difference between the central pixel and one of 8 neighboring pixels
    diff_x0 = pts_3d_map_ctr - pts_3d_map_x0
    diff_x1 = pts_3d_map_ctr - pts_3d_map_x1
    diff_y0 = pts_3d_map_y0 - pts_3d_map_ctr
    diff_y1 = pts_3d_map_y1 - pts_3d_map_ctr
    diff_x0y0 = pts_3d_map_x0y0 - pts_3d_map_ctr
    diff_x0y1 = pts_3d_map_ctr - pts_3d_map_x0y1
    diff_x1y0 = pts_3d_map_x1y0 - pts_3d_map_ctr
    diff_x1y1 = pts_3d_map_ctr - pts_3d_map_x1y1

    ## flatten the diff to a #pixle by 3 matrix
    pix_num = kitti_shape[0] * (kitti_shape[1]-2*nei) * (kitti_shape[2]-2*nei)
    diff_x0 = tf.reshape(diff_x0, [pix_num, 3])
    diff_y0 = tf.reshape(diff_y0, [pix_num, 3])
    diff_x1 = tf.reshape(diff_x1, [pix_num, 3])
    diff_y1 = tf.reshape(diff_y1, [pix_num, 3])
    diff_x0y0 = tf.reshape(diff_x0y0, [pix_num, 3])
    diff_x0y1 = tf.reshape(diff_x0y1, [pix_num, 3])
    diff_x1y0 = tf.reshape(diff_x1y0, [pix_num, 3])
    diff_x1y1 = tf.reshape(diff_x1y1, [pix_num, 3])

    ## calculate normal by cross product of two vectors
    normals0 = normalize_l2(tf.cross(diff_x1, diff_y1)) #* tf.tile(normals0_mask[:, None], [1,3])
    normals1 = normalize_l2(tf.cross(diff_x0, diff_y0)) #* tf.tile(normals1_mask[:, None], [1,3])
    normals2 = normalize_l2(tf.cross(diff_x0y1, diff_x0y0)) #* tf.tile(normals2_mask[:, None], [1,3])
    normals3 = normalize_l2(tf.cross(diff_x1y0, diff_x1y1)) #* tf.tile(normals3_mask[:, None], [1,3])
    
    normal_vector = tf.reduce_sum(tf.concat([[normals0], [normals1], [normals2], [normals3]], 0),0)
    normal_vector = normalize_l2(normals0)
    normal_map = tf.reshape(tf.squeeze(normal_vector), [kitti_shape[0]]+[kitti_shape[1]-2*nei]+[kitti_shape[2]-2*nei]+[3])

    normal_map *= tf.tile(tf.expand_dims(tf.cast(mask[:, nei:-nei, nei:-nei], tf.float32), -1), [1,1,1,3])
    normal_map = tf.pad(normal_map, [[0,0], [nei, nei], [nei, nei], [0,0]] ,"CONSTANT")

    return normal_map 

def depth_to_normals_tf(depth, intrinsics, scope=None, eps=1e-4):
    """
    :param depth: real depth (B,1,H,W) 
    :param intrinsics: (B,4)
    :return: normals (B,3,H,W)
    """
    with tf.name_scope(scope, 'depth_to_normals_tf', [depth, intrinsics]):
        H, W = depth.shape.as_list()[-2:]
        B = tf.shape(depth)[0]  # config.batch_size
        depth = tf.reshape(depth, [B, H, W])

        # fx_rel = fx_abs / W, cx_real = cx_abs / W
        fx, fy, cx, cy = tf.split(tf.expand_dims(intrinsics, 2), 4, axis=1)  # (B,1,1)
        inv_fx = tf.div(1.0, fx * W)
        inv_fy = tf.div(1.0, fy * H)
        cx = cx * W
        cy = cy * H

        X, Y = tf.meshgrid(tf.range(W), tf.range(H))
        X = tf.cast(tf.tile(tf.expand_dims(X, axis=0), [B, 1, 1]), tf.float32)  # (B,H,W)
        Y = tf.cast(tf.tile(tf.expand_dims(Y, axis=0), [B, 1, 1]), tf.float32)

        x_cord = (X - cx) * inv_fx * depth
        y_cord = (Y - cy) * inv_fy * depth
        p = tf.stack([x_cord, y_cord, depth], axis=3, name='p_3d')  # (B,H,W,3)

        # vector of p_3d in west, south, east, north direction
        p_ctr = p[:, 1:-1, 1:-1, :]
        vw = p_ctr - p[:, 1:-1, 2:, :]
        vs = p[:, 2:, 1:-1, :] - p_ctr
        ve = p_ctr - p[:, 1:-1, :-2, :]
        vn = p[:, :-2, 1:-1, :] - p_ctr
        normal_1 = tf.cross(vs, vw, name='cross_1')  # (B,H-2,W-2,3)
        normal_2 = tf.cross(vn, ve, name='cross_2')
        normal_1 = tf.nn.l2_normalize(normal_1, 3, epsilon=eps)
        normal_2 = tf.nn.l2_normalize(normal_2, 3, epsilon=eps)
        normal = normal_1 + normal_2
        # unused = tf.less(tf.norm(normal, axis=3), np.sqrt(eps))
        # unused = tf.stack([unused] * 3, axis=3)
        normal = tf.nn.l2_normalize(normal, 3, epsilon=eps, name='normal')
        # normal = tf.where(unused, tf.zeros_like(normal), normal)

        paddings = [[0, 0], [1, 1], [1, 1], [0, 0]]
        normal = tf.pad(normal, paddings)  # (B,H,W,3)
        normal = convertNHWC2NCHW(normal, 'normal_NCHW')
        return normal 

def get_ee_pos(states, are_tensors=False):
        theta1, theta2, theta3, theta4, theta5, theta6, theta7 = \
            states[:, :1], states[:, 1:2], states[:, 2:3], states[:, 3:4], states[:, 4:5], states[:, 5:6], states[:, 6:]
        if are_tensors:
            rot_axis = tf.concat([tf.cos(theta2) * tf.cos(theta1), tf.cos(theta2) * tf.sin(theta1), -tf.sin(theta2)],
                                 axis=1)
            rot_perp_axis = tf.concat([-tf.sin(theta1), tf.cos(theta1), tf.zeros(tf.shape(theta1))], axis=1)
            cur_end = tf.concat([
                0.1 * tf.cos(theta1) + 0.4 * tf.cos(theta1) * tf.cos(theta2),
                0.1 * tf.sin(theta1) + 0.4 * tf.sin(theta1) * tf.cos(theta2) - 0.188,
                -0.4 * tf.sin(theta2)
            ], axis=1)

            for length, hinge, roll in [(0.321, theta4, theta3), (0.16828, theta6, theta5)]:
                perp_all_axis = tf.cross(rot_axis, rot_perp_axis)
                x = tf.cos(hinge) * rot_axis
                y = tf.sin(hinge) * tf.sin(roll) * rot_perp_axis
                z = -tf.sin(hinge) * tf.cos(roll) * perp_all_axis
                new_rot_axis = x + y + z
                new_rot_perp_axis = tf.cross(new_rot_axis, rot_axis)
                new_rot_perp_axis = tf.where(tf.less(tf.norm(new_rot_perp_axis, axis=1), 1e-30),
                                             rot_perp_axis, new_rot_perp_axis)
                new_rot_perp_axis /= tf.norm(new_rot_perp_axis, axis=1, keepdims=True)
                rot_axis, rot_perp_axis, cur_end = new_rot_axis, new_rot_perp_axis, cur_end + length * new_rot_axis
        else:
            rot_axis = np.concatenate([np.cos(theta2) * np.cos(theta1), np.cos(theta2) * np.sin(theta1), -np.sin(theta2)],
                                      axis=1)
            rot_perp_axis = np.concatenate([-np.sin(theta1), np.cos(theta1), np.zeros(theta1.shape)], axis=1)
            cur_end = np.concatenate([
                0.1 * np.cos(theta1) + 0.4 * np.cos(theta1) * np.cos(theta2),
                0.1 * np.sin(theta1) + 0.4 * np.sin(theta1) * np.cos(theta2) - 0.188,
                -0.4 * np.sin(theta2)
            ], axis=1)

            for length, hinge, roll in [(0.321, theta4, theta3), (0.16828, theta6, theta5)]:
                perp_all_axis = np.cross(rot_axis, rot_perp_axis)
                x = np.cos(hinge) * rot_axis
                y = np.sin(hinge) * np.sin(roll) * rot_perp_axis
                z = -np.sin(hinge) * np.cos(roll) * perp_all_axis
                new_rot_axis = x + y + z
                new_rot_perp_axis = np.cross(new_rot_axis, rot_axis)
                new_rot_perp_axis[np.linalg.norm(new_rot_perp_axis, axis=1) < 1e-30] = \
                    rot_perp_axis[np.linalg.norm(new_rot_perp_axis, axis=1) < 1e-30]
                new_rot_perp_axis /= np.linalg.norm(new_rot_perp_axis, axis=1, keepdims=True)
                rot_axis, rot_perp_axis, cur_end = new_rot_axis, new_rot_perp_axis, cur_end + length * new_rot_axis

        return cur_end