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

#!/usr/bin/env python

import argparse
import os, json

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('-ac', '--aruco_cube', metavar='ARUCO_CUBE', type=str, help='json aruco cube 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()

    # 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()

    # 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()

    # Init head movment estimation
    last_accelerometer = numpy.array([])
    last_accelerometer_ts_ms = 0
    last_gyroscope = numpy.array([])
    last_gyroscope_ts_ms = 0
    gyroscope_drift = numpy.zeros(3)
    head_translation_speed = numpy.zeros(3)
    head_rotation_speed = numpy.zeros(3)
    gravity = -9.81
    
    # 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 last_accelerometer
        nonlocal last_accelerometer_ts_ms
        nonlocal last_gyroscope
        nonlocal last_gyroscope_ts_ms
        nonlocal gyroscope_drift
        nonlocal head_translation_speed
        nonlocal head_rotation_speed
        nonlocal data_ts_ms

        data_ts_ms = data_ts / 1e3

        match data_object_type:

            case 'Accelerometer':

                # Convert m/s2 into cm/ms2
                current_accelerometer = numpy.array(data_object.value) * 1e-4

                # Init accelerometer time integration
                if last_accelerometer_ts_ms == 0:

                    last_accelerometer = current_accelerometer
                    last_accelerometer_ts_ms = data_ts_ms

                # Integrate accelerometer values in time to get head translation speed (cm/ms)
                else:

                    head_translation_speed +=  (last_accelerometer + current_accelerometer) * (data_ts_ms - last_accelerometer_ts_ms) / 2

                   # print(head_translation_speed)

            case 'Gyroscope':

                # convert deg/s into deg/ms
                current_gyroscope = numpy.array(data_object.value) * 1e-3

                # Init gyroscope derivation
                if last_gyroscope_ts_ms == 0:

                    last_gyroscope = current_gyroscope
                    last_gyroscope_ts_ms = data_ts_ms
                    gyroscope_derivation = 0

                # Derivate gyroscope
                else:
                    gyroscope_derivation = (current_gyroscope - last_gyroscope) / (data_ts_ms - last_gyroscope_ts_ms)

                # Update gyroscope drift when is not moving
                if numpy.linalg.norm(gyroscope_derivation) < 1e-7:

                     gyroscope_drift = current_gyroscope
                     
                 # Drift compensation
                head_rotation_speed = current_gyroscope - gyroscope_drift

                #print(head_rotation_speed)
        
    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.get_tracked_markers())

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

                    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: use tobii glasses inertial sensors to estimate cube pose from last estimated pose
                elif aruco_cube_success:

                    # Translate cube according head translation speed
                    #aruco_cube_tvec += head_translation_speed * (video_ts_ms - aruco_cube_ts_ms)

                    #print('after:')
                    #print(aruco_cube_tvec)

                    # Rotate cube around origin according head rotation speed
                    R = make_rotation_matrix(* (head_rotation_speed * (video_ts_ms - aruco_cube_ts_ms)))
                    aruco_cube_tvec = aruco_cube_tvec.dot(R.T)
                    aruco_cube_ts_ms = video_ts_ms

                    #print('berore rotation: ', aruco_cube_tvec)
                    #print(R)
                    #print('after rotation: ', aruco_cube_tvec)

                    # Set cube pose estimation
                    aruco_cube.set_pose(tvec = aruco_cube_tvec, rvec = aruco_cube_rvec)

                else:

                    raise UserWarning('Cube pose estimation fails.')

                # 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.get_K(), (0, 0, 0, 0))

                # 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), (500, 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()