ego vehicle coord system parallel to world z plane nuscenes

"""
This code is to sort out a reply to
https://github.com/nutonomy/nuscenes-devkit/issues/122
"""

import numpy as np
import matplotlib.pyplot as plt

from pyquaternion import Quaternion

from nuscenes import NuScenes

GLOBAL_X_AXIS = np.array([1, 0, 0])
GLOBAL_Z_AXIS = np.array([0, 0, 1])


def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)


def angle_between(v1, v2):
    """ Returns the angle in degrees between vectors 'v1' and 'v2' """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return 180/np.pi * np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))


def main():
    nusc = NuScenes('v1.0-mini')

    scene_number = 6

    # Grab first sample
    sample = nusc.get('sample', nusc.scene[scene_number]['first_sample_token'])

    # Cycle through all samples and store ego pose translation and rotation.
    poses = []
    while not sample['next'] == '':
        sample = nusc.get('sample', sample['next'])
        lidar = nusc.get('sample_data', sample['data']['LIDAR_TOP'])
        pose = nusc.get('ego_pose', lidar['ego_pose_token'])
        poses.append((pose['translation'], Quaternion(pose['rotation'])))

    # Print values for first pose.
    (x, y, _), q = poses[0]
    print('First location (x, y): {:.1f}, {:.1f}'.format(x, y))
    print('Orientation at first location: {:.2f} degrees'.format(q.degrees))

    # Plot all poses on global x, y plan along with orientation.
    # This doesn't really answer the question, it's just as a sanity check.

    for (x, y, z), q in poses:
        plt.plot(x, y, 'r.')

        # The x axis of the ego frame is pointing towards the front of the vehicle.
        ego_x_axis = q.rotate(GLOBAL_X_AXIS)
        plt.arrow(x, y, ego_x_axis[0], ego_x_axis[1])

    plt.axis('equal')
    plt.xlabel('x-global')
    plt.ylabel('y-global')

    # Calculate deviations from global vertical axis and plot.
    # This demonstrates that the z-axis of the ego poses are not
    # all pointing upwards.

    vertical_offsets = []
    for (_, _, _), q in poses:

        # The z axis of the ego frame indicates deviation from global vertical axis (direction of gravity)
        ego_z_axis = q.rotate(GLOBAL_Z_AXIS)
        vertical_offsets.append(angle_between(ego_z_axis, GLOBAL_Z_AXIS))

    plt.savefig('xyposes.png')
    plt.figure()
    plt.hist(vertical_offsets)
    plt.xlabel('Angle (degrees)')
    plt.ylabel('Count')
    plt.title('Offset from global vertical axis.')
    plt.savefig('vertical_offsets.png')

    # Finally show that all calibrated sensor tokens for a particular scene is indeed the same
    sample = nusc.get('sample', nusc.scene[scene_number]['first_sample_token'])

    cali_sensor_tokens = []
    while not sample['next'] == '':
        sample = nusc.get('sample', sample['next'])
        lidar = nusc.get('sample_data', sample['data']['LIDAR_TOP'])
        cali_sensor = nusc.get('calibrated_sensor', lidar['calibrated_sensor_token'])
        cali_sensor_tokens.append(cali_sensor['token'])

    # Assert that they all point to the same calibrated sensor record
    assert len(set(cali_sensor_tokens)) == 1


if __name__ == "__main__":
    main()