path: root/src/argaze/utils/tobii_stream_arcube_display.py

#!/usr/bin/env python

import argparse
import os, json
import math
import threading

from argaze import DataStructures
from argaze import GazeFeatures
from argaze.TobiiGlassesPro2 import *
from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *
from argaze.utils import MiscFeatures

import cv2 as cv
import numpy

def make_rotation_matrix(x, y, z):

    # Create rotation matrix around x axis
    c = numpy.cos(numpy.deg2rad(x))
    s = numpy.sin(numpy.deg2rad(x))
    Rx = numpy.array([[1, 0, 0], [0, c, -s], [0, s, c]])

    # Create rotation matrix around y axis
    c = numpy.cos(numpy.deg2rad(y))
    s = numpy.sin(numpy.deg2rad(y))
    Ry = numpy.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])

    # Create rotation matrix around z axis
    c = numpy.cos(numpy.deg2rad(z))
    s = numpy.sin(numpy.deg2rad(z))
    Rz = numpy.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])

    # Return intrinsic rotation matrix
    return Rx.dot(Ry.dot(Rz))

def main():
    """
    Track ArUcoCube into Tobii Glasses Pro 2 camera video stream.
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('-t', '--tobii_ip', metavar='TOBII_IP', type=str, default=None, help='tobii glasses ip')
    parser.add_argument('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default=None, help='json camera calibration filepath')
    parser.add_argument('-p', '--aruco_tracker_configuration', metavar='TRACK_CONFIG', type=str, default=None, help='json aruco tracker configuration filepath')
    parser.add_argument('-i', '--imu_calibration', metavar='IMU_CALIB', type=str, default=None, help='json imu calibration filepath')
    parser.add_argument('-ac', '--aruco_cube', metavar='ARUCO_CUBE', type=str, help='json aruco cube description filepath')
    parser.add_argument('-s', '--aoi_scene', metavar='AOI_SCENE', type=str, help='obj aoi 3D scene description filepath')
    parser.add_argument('-w', '--window', metavar='DISPLAY', type=bool, default=True, help='enable window display', action=argparse.BooleanOptionalAction)
    args = parser.parse_args()

    # Create tobii controller (with auto discovery network process if no ip argument is provided)
    print('\nLooking for a Tobii Glasses Pro 2 device ...')

    try:

        tobii_controller = TobiiController.TobiiController(args.tobii_ip)
        print(f'Tobii Glasses Pro 2 device found at {tobii_controller.address} address.')

    except ConnectionError as e:

        print(e)
        exit()

    # Setup camera at 25 fps to work on Full HD video stream
    tobii_controller.set_scene_camera_freq_25()

    # Print current confirugration
    print(f'Tobii Glasses Pro 2 configuration:')
    for key, value in tobii_controller.get_configuration().items():
        print(f'\t{key}: {value}')

    # Enable tobii data stream 
    tobii_data_stream = tobii_controller.enable_data_stream()

    # Enable tobii video stream
    tobii_video_stream = tobii_controller.enable_video_stream()

    # Load aruco cube description
    aruco_cube = ArUcoCube.ArUcoCube(args.aruco_cube)
    aruco_cube.print_cache()

    # Load AOI 3D scene centered onto aruco cube
    aoi3D_scene = AOI3DScene.AOI3DScene()
    aoi3D_scene.load(args.aoi_scene)

    print(f'\nAOI in {os.path.basename(args.aoi_scene)} scene related to ArCube:')
    for aoi in aoi3D_scene.keys():
        print(f'\t{aoi}')

    # Create aruco camera
    aruco_camera = ArUcoCamera.ArUcoCamera()
    
     # Load calibration file
    if args.camera_calibration != None:

        aruco_camera.load_calibration_file(args.camera_calibration)

    else:

        raise UserWarning('.json camera calibration filepath required. Use -c option.')

    # Create aruco tracker
    aruco_tracker = ArUcoTracker.ArUcoTracker(aruco_cube.dictionary, aruco_cube.marker_size, aruco_camera)

    # Load specific configuration file
    if args.aruco_tracker_configuration != None:

        aruco_tracker.load_configuration_file(args.aruco_tracker_configuration)

        print(f'\nArUcoTracker configuration for markers detection:')
        aruco_tracker.print_configuration()

    # Create tobii imu handler to track head pose changes when arcuco cube pose can't be estimated
    # So, the resulting head pose is relative to last pose estimation
    tobii_imu = TobiiInertialMeasureUnit.TobiiInertialMeasureUnit()
    
    # Load optional imu calibration file
    if args.imu_calibration != None:

        tobii_imu.load_calibration_file(args.imu_calibration)

    # Init tobii imu lock
    tobii_imu_lock = threading.Lock()

    # TEST : DIFF ACC
    last_accl = numpy.zeros(3)

    # Init data timestamped in millisecond
    data_ts_ms = 0
    
    # Assess temporal performance
    loop_chrono = MiscFeatures.TimeProbe()
    loop_ps = 0

    def data_stream_callback(data_ts, data_object, data_object_type):

        nonlocal tobii_imu
        nonlocal tobii_imu_lock
        nonlocal data_ts_ms

        #TEST
        nonlocal last_accl

        data_ts_ms = data_ts / 1e3

        # Don't update imu when it is used
        if tobii_imu_lock.locked():
            return

        # Lock tobii imu updates
        tobii_imu_lock.acquire()

        match data_object_type:

            case 'Gyroscope':

                data_object = tobii_imu.apply_gyroscope_offset(data_object)

                tobii_imu.update_rotation(data_ts, data_object)

            
            case 'Accelerometer':
                pass
                '''
                print('raw accelerometer(m/s2)=', data_object.value)

                # TEST : 
                diff_accl = last_accl - numpy.array(data_object.value)
                last_accl = numpy.array(data_object.value)
                print('\tdiff(cm/s2)=', 100 * numpy.linalg.norm(diff_accl))

                # TEST : ignore acceleration double
                if numpy.linalg.norm(diff_accl) > 0.:
                
                    data_object = tobii_imu.apply_accelerometer_coefficients(data_object)

                    print('corrected accelerometer(m/s2)=', data_object.value)

                    print('current plumb=', tobii_imu.get_plumb())

                    data_object = tobii_imu.apply_plumb(data_object)

                    print('corrected accelerometer - gravity(m/s2)=', data_object.value)
                    print('\tnorm(cm/s2)=', 100 * numpy.linalg.norm(data_object.value))

                    tobii_imu.update_translation(data_ts, data_object)
                '''
        # Unlock tobii imu updates
        tobii_imu_lock.release()
        
    tobii_data_stream.reading_callback = data_stream_callback

    # Start streaming
    tobii_controller.start_streaming()

    # Live video stream capture loop
    try:

        # Assess loop performance
        loop_chrono = MiscFeatures.TimeProbe()
        fps = 0

        # Track aruco cube pose
        aruco_cube_tvec = numpy.zeros(3)
        aruco_cube_rvec = numpy.zeros(3)
        aruco_cube_success = False
        aruco_cube_validity = False
        aruco_cube_ts_ms = 0

        while tobii_video_stream.is_alive():

            # Read video stream
            video_ts, video_frame = tobii_video_stream.read()
            video_ts_ms = video_ts / 1e3

            # Copy video frame to edit visualisation on it without disrupting aruco tracking
            visu_frame = video_frame.copy()

            # Process video and data frame
            try:

                # Track markers with pose estimation
                aruco_tracker.track(video_frame.matrix)

                # Estimate cube pose from tracked markers
                tvec, rvec, success, validity = aruco_cube.estimate_pose(aruco_tracker.tracked_markers)

                # Cube pose estimation succeed and is validated by 2 faces at least
                if success and validity >= 2:

                    # Lock tobii imu updates
                    tobii_imu_lock.acquire()

                    # Reset head rotation, translation and translation speed (cm/s)
                    # Note : head translation speed is also estimated thanks to accelerometer sensor (see upward)
                    tobii_imu.reset_rotation()
                    #tobii_imu.reset_translation(translation_speed = (tvec - aruco_cube_tvec) / (video_ts_ms - aruco_cube_ts_ms))

                    # Create a rotation matrix to transform cube rotation from camera referential to imu referential
                    F = make_rotation_matrix(*TobiiInertialMeasureUnit.CAMERA_TO_IMU_ROTATION_VECTOR)
                    R, _ = cv.Rodrigues(rvec)
                    rvec_flipped, _ = cv.Rodrigues(F.dot(R))

                    # Update head plumb orientation with flipped cube orientation
                    tobii_imu.rotate_plumb(rvec_flipped)

                    # Unlock tobii imu updates
                    tobii_imu_lock.release()

                    # Store cube pose
                    aruco_cube_tvec = tvec
                    aruco_cube_rvec = rvec
                    aruco_cube_success = success
                    aruco_cube_validity = validity
                    aruco_cube_ts_ms = video_ts_ms

                # Cube pose estimation fails
                elif aruco_cube_success:

                    # Use tobii glasses inertial sensors to estimate cube pose from last estimated pose

                    # Translate cube into imu referential
                    imu_tvec = aruco_cube_tvec + numpy.array(TobiiInertialMeasureUnit.CAMERA_TO_IMU_TRANSLATION_VECTOR)

                    # Translate cube according head translation
                    imu_tvec = imu_tvec + tobii_imu.translation

                    # Rotate cube around imu origin according head rotation
                    imu_rvec = tobii_imu.rotation * numpy.array([-1., -1., 1.])
                    imu_R = make_rotation_matrix(*imu_rvec)
                    new_imu_tvec = imu_tvec.dot(imu_R)

                    # Translate back cube into camera referential
                    new_tvec = new_imu_tvec - numpy.array(TobiiInertialMeasureUnit.CAMERA_TO_IMU_TRANSLATION_VECTOR)

                    # Rotate cube orientation (supposing cube top is up in )
                    imu_rvec = tobii_imu.rotation * numpy.array([1., -1., 1.])
                    imu_R = make_rotation_matrix(*imu_rvec)

                    C, _ = cv.Rodrigues(aruco_cube_rvec)
                    C = C.dot(imu_R)
                    new_rvec, _ = cv.Rodrigues(C)
                    #new_rvec = aruco_cube_rvec

                    # Set cube pose estimation
                    aruco_cube.translation = new_tvec
                    aruco_cube.rotation = new_rvec

                else:

                    raise UserWarning('Cube pose estimation fails.')

                # Project AOI 3 scene onto camera frame

                # DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
                # This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
                aoi2D_scene = aoi3D_scene.project(aruco_cube.translation, aruco_cube.rotation, aruco_camera.K)

                # Draw projected scene
                #aoi2D_scene.draw(visu_frame.matrix)

                # Draw markers pose estimation
                #aruco_tracker.draw_tracked_markers(visu_frame.matrix)

                # Draw cube pose estimation (without camera distorsion)
                aruco_cube.draw(visu_frame.matrix, aruco_camera.K, aruco_camera.D, draw_faces=False)

                # Warn about cube pose validity
                if not aruco_cube_validity:

                    raise UserWarning('Cube pose estimation is not validated.')

            # Write warning
            except UserWarning as w:

                cv.rectangle(visu_frame.matrix, (0, 100), (600, 150), (127, 127, 127), -1)
                cv.putText(visu_frame.matrix, str(w), (20, 140), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

            # Assess loop performance
            lap_time, lap_counter, elapsed_time = loop_chrono.lap()

            # Update fps each 10 loops
            if lap_counter >= 10:

                loop_ps = 1e3 * lap_counter / elapsed_time
                loop_chrono.restart()

            # Draw center
            cv.line(visu_frame.matrix, (int(visu_frame.width/2) - 50, int(visu_frame.height/2)), (int(visu_frame.width/2) + 50, int(visu_frame.height/2)), (255, 150, 150), 1)
            cv.line(visu_frame.matrix, (int(visu_frame.width/2), int(visu_frame.height/2) - 50), (int(visu_frame.width/2), int(visu_frame.height/2) + 50), (255, 150, 150), 1)

            # Write stream timing
            cv.rectangle(visu_frame.matrix, (0, 0), (1100, 50), (63, 63, 63), -1)
            cv.putText(visu_frame.matrix, f'Data stream time: {int(data_ts_ms)} ms', (20, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
            cv.putText(visu_frame.matrix, f'Video delay: {int(data_ts_ms - video_ts_ms)} ms', (550, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
            cv.putText(visu_frame.matrix, f'Fps: {int(loop_ps)}', (950, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)

            cv.imshow(f'Stream ArUcoCube', visu_frame.matrix)

            # Close window using 'Esc' key
            if cv.waitKey(1) == 27:
                break

    # Exit on 'ctrl+C' interruption
    except KeyboardInterrupt:
        pass

    # Stop frame display
    cv.destroyAllWindows()

    # Stop streaming
    tobii_controller.stop_streaming()
    
if __name__ == '__main__':

    main()